comrad-bench 1.0.0__tar.gz

comrad_bench-1.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

comrad_bench-1.0.0/MANIFEST.in

@@ -0,0 +1,5 @@
+include README.md
+recursive-include comrad/scenarios *.cfg *.wad
+recursive-include comrad/docs *.md
+recursive-include comrad/experiments *.md
+recursive-include comrad/templates *.sh

comrad_bench-1.0.0/PKG-INFO

@@ -0,0 +1,229 @@
+Metadata-Version: 2.4
+Name: comrad-bench
+Version: 1.0.0
+Summary: MARL benchmark suite in ViZDoom
+Author-email: "Khoi H.B. Nguyen" <baokhoi136@gmail.com>, Tristan Tomilin <tristan.tomilin@hotmail.com>, Dimitar Zhekov <mitko.zh@gmail.com>
+License: MIT
+Project-URL: Homepage, https://github.com/TTomilin/COMRAD
+Project-URL: Repository, https://github.com/TTomilin/COMRAD
+Keywords: marl,vizdoom,reinforcement-learning,multi-agent,benchmark
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Games/Entertainment :: First Person Shooters
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy<2.0,>=1.26.4
+Requires-Dist: gymnasium[classic-control]<1.0,>=0.27
+Requires-Dist: pettingzoo<1.25.0
+Requires-Dist: opencv-python>=4.11.0
+Requires-Dist: pygame>=2.6.1
+Requires-Dist: pre-commit>=4.5.1
+Requires-Dist: pynput>=1.8.1
+Requires-Dist: wandb[media]==0.22.1
+Requires-Dist: torch>=2.10.0
+Requires-Dist: torchrl>=0.11.0
+Requires-Dist: torchvision>=0.25.0
+Requires-Dist: torchaudio>=2.10.0
+Requires-Dist: psutil>=7.2.2
+Requires-Dist: pyglet>=2.1.12
+Requires-Dist: tensorboard>=2.20.0
+Requires-Dist: tensorboardx>=2.6.4
+Requires-Dist: threadpoolctl>=3.6.0
+Requires-Dist: colorlog>=6.10.1
+Requires-Dist: signal-slot-mp<2.0,>=1.0.3
+Requires-Dist: filelock>=3.20.3
+Requires-Dist: huggingface-hub>=0.10.0
+Requires-Dist: pandas>=2.3.3
+Requires-Dist: onnx>=1.19.0
+Requires-Dist: onnxruntime>=1.23.2
+Requires-Dist: matplotlib>=3.10.8
+Requires-Dist: seaborn>=0.13.2
+Requires-Dist: scipy>=1.17.0
+Requires-Dist: vizdoom>=1.3.0
+Provides-Extra: sample-factory
+Requires-Dist: torch>=2.9.0; extra == "sample-factory"
+Requires-Dist: wandb>=0.22.1; extra == "sample-factory"
+Requires-Dist: tensorboard>=2.20.0; extra == "sample-factory"
+Requires-Dist: colorlog>=6.10.1; extra == "sample-factory"
+Requires-Dist: signal-slot-mp>=1.0.5; extra == "sample-factory"
+Requires-Dist: faster-fifo>=1.5.2; extra == "sample-factory"
+Requires-Dist: filelock>=3.19.1; extra == "sample-factory"
+Requires-Dist: psutil>=7.1.0; extra == "sample-factory"
+Requires-Dist: huggingface_hub>=0.36.0; extra == "sample-factory"
+Requires-Dist: scipy>=1.16.3; extra == "sample-factory"
+Provides-Extra: visual
+Requires-Dist: omgifol>=0.5.1; extra == "visual"
+Requires-Dist: pillow>=11.3.0; extra == "visual"
+Requires-Dist: pypng>=0.20220715.0; extra == "visual"
+Requires-Dist: shapely>=2.1.2; extra == "visual"
+Provides-Extra: doomgen
+Requires-Dist: doomgen; extra == "doomgen"
+Dynamic: license-file
+
+# COMRAD
+
+COMRAD is a cooperative multi-agent reinforcement learning benchmark built on ViZDoom for first-person visual coordination. This repository contains the benchmark runtime assets, COMRAD-specific environment and model code, a modified Sample Factory training stack, paper artifacts, and the scripts used to produce figures, tables, and diagnostic media.
+
+<p align="center">
+  <img src="results/demo_all_checkpoints/comrad_vid_gif_540.gif" alt="COMRAD demo" style="width: 100%; max-width: 960px;">
+</p>
+
+## Repository map
+
+| Path | Purpose |
+| --- | --- |
+| `comrad/` | Benchmark package: training entrypoints, environment registration, wrappers, models, shipped scenario assets, launchers, tests. |
+| `sample_factory/` | Modified training engine used by COMRAD. |
+| `scripts/` | Figure-generation, result aggregation, WAD rendering, and experiment helper scripts. |
+| `results/` | Committed figures, tables, result summaries, and selected media artifacts. |
+| `DoomGen/` | Procedural generator used to build and vary scenario WADs. |
+
+Directory-level guides are provided in:
+
+- [`comrad/README.md`](comrad/README.md)
+- [`sample_factory/README.md`](sample_factory/README.md)
+- [`scripts/README.md`](scripts/README.md)
+- [`results/README.md`](results/README.md)
+
+## Setup
+
+COMRAD is packaged through `pyproject.toml` and is intended to be run with `uv`.
+
+```bash
+uv sync --all-extras
+uv run pre-commit install
+```
+
+The project requires Python 3.10 or newer. The package metadata and CI workflow are the authoritative sources for dependency resolution.
+
+## Single-run training
+
+The main training entrypoint is `comrad.train`.
+
+```bash
+uv run python -m comrad.train \
+  --env=stag_hunt_arena \
+  --algo=IPPO \
+  --train_for_env_steps=10000 \
+  --num_workers=4 \
+  --num_envs_per_worker=4 \
+  --device=cpu
+```
+
+Training outputs are written under `train_dir/<experiment>/`. The main runtime log is:
+
+```bash
+tail -n 120 train_dir/<experiment>/sf_log.txt
+```
+
+For launcher-owned benchmark sweeps, use `comrad/train_all.py` instead of manually reproducing those long commands.
+
+## Recording and evaluation
+
+High-resolution checkpoint recording:
+
+```bash
+uv run python -m comrad.record_video \
+  --env=stag_hunt_arena \
+  --train_dir=train_dir \
+  --experiment=<experiment_name> \
+  --load_checkpoint_kind=best \
+  --resolution=1280x720
+```
+
+Automap/top-down heatmap rendering:
+
+```bash
+uv run python -m comrad.record_topdown_heatmap \
+  --env=ammo_carrier \
+  --train_dir=<train_root> \
+  --experiment=<experiment_name> \
+  --output_dir=results/videos/ammo_carrier_navigation_heatmap \
+  --device=cpu \
+  --overwrite
+```
+
+Both commands expect a full experiment directory with at least `config.json` and the corresponding checkpoint directory.
+
+## Benchmark launchers
+
+COMRAD ships several paper-facing launcher profiles through `comrad/train_all.py`.
+
+- `benchmark`: full benchmark suite
+- `agent_scaling`: Armory Siege agent-count sweep
+- `platform_chain_curriculum`: curriculum comparison on Platform Chain
+- `qmix_lr_sensitivity_pilot`: Short sensitivity study
+- `qmix_lr_sensitivity_full`: Full sensitivity study
+
+Run a launcher profile locally with:
+
+```bash
+COMRAD_TRAIN_ALL_PROFILE=benchmark \
+uv run python -m sample_factory.launcher.run \
+  --run=comrad.train_all \
+  --backend=processes \
+  --max_parallel=1 \
+  --pause_between=1
+```
+
+The exact benchmark definitions, algorithms, root directory names, and profile constants live at the top of [`comrad/train_all.py`](comrad/train_all.py).
+
+## Scenario assets and DoomGen
+
+Runtime scenario assets are stored in [`comrad/scenarios/`](comrad/scenarios/README.md) as `.cfg` and `.wad` pairs, plus pre-generated batch directories such as `batch_platform_chain_curriculum/`.
+
+Generation workflows are in `DoomGen/`. The usual split is:
+
+- `DoomGen/examples/benchmark`: procedural generators and geometry logic
+- `comrad/scenarios/`: runtime WADs and configs consumed by training and evaluation
+
+If a scenario generator changes, rebuild the corresponding WAD in `DoomGen/` and copy the regenerated artifact back into `comrad/scenarios/` before training or evaluation.
+
+### Using DoomGen as a Python library
+
+DoomGen is declared as an optional dependency in the uv workspace. It is **not** installed by a plain `uv sync`:
+
+```bash
+uv sync --extra doomgen      # install doomgen + its deps
+# or
+uv run --extra doomgen python -m comrad.train ...   # one-off with doomgen available
+# or
+uv pip install -e DoomGen/   # editable install directly from local checkout
+```
+
+After any of these, `import doomgen` works from the root Python environment. Without `--extra doomgen`, the core COMRAD package remains self-contained and DoomGen is not pulled in.
+
+## Tests
+
+Run COMRAD and Sample Factory tests with `uv`:
+
+```bash
+uv run pytest comrad/tests -v
+uv run pytest sample_factory/tests -v
+```
+
+Tests in `comrad/tests/` focus on benchmark logic, launchers, and recorder behavior. Tests in `sample_factory/tests/` cover the modified training stack.
+
+## Outputs and artifacts
+
+- `train_dir/`: experiment directories, configs, checkpoints, logs
+- `results/`: committed figures, tables, result summaries, selected media
+- `wandb/`: local W&B run state and uploads
+
+See [`results/README.md`](results/README.md) for the structure of the committed artifacts.
+
+## License and upstream assets
+
+The repository is released under the MIT license in [`LICENSE`](LICENSE). Upstream and third-party assets are documented in the repository itself, including the bundled `scripts/wad2image/` tool and the paper's citations to ViZDoom, Sample Factory, HARL, PyMARL, and QPLEX.
+
+## Acknowledgements
+
+We would like to thank all researchers, engineers, advisors involved in this work; and the TU/e Supercomputing Center for providing the infrastructure to run the experiments. The environments are built on top of [ViZDoom](https://github.com/Farama-Foundation/ViZDoom). Our infrastructure is built on top of [Sample Factory](https://github.com/alex-petrenko/sample-factory). Our experiments were managed using [Weights & Biases](https://wandb.ai/). For on-policy actor-critic methods, we implement IPPO and MAPPO from the original implementation in [Sample Factory](https://github.com/alex-petrenko/sample-factory), and HAPPO following the [HARL codebase](https://github.com/PKU-MARL/HARL). For off-policy value decomposition methods, we implement IDQN, VDN, QMIX following the [PyMARL framework](https://github.com/oxwhirl/pymarl), and QPLEX following [official QPLEX reference implementation](https://github.com/wjh720/QPLEX).
+
+<!-- ## Citation -->

comrad_bench-1.0.0/README.md

@@ -0,0 +1,161 @@
+# COMRAD
+
+COMRAD is a cooperative multi-agent reinforcement learning benchmark built on ViZDoom for first-person visual coordination. This repository contains the benchmark runtime assets, COMRAD-specific environment and model code, a modified Sample Factory training stack, paper artifacts, and the scripts used to produce figures, tables, and diagnostic media.
+
+<p align="center">
+  <img src="results/demo_all_checkpoints/comrad_vid_gif_540.gif" alt="COMRAD demo" style="width: 100%; max-width: 960px;">
+</p>
+
+## Repository map
+
+| Path | Purpose |
+| --- | --- |
+| `comrad/` | Benchmark package: training entrypoints, environment registration, wrappers, models, shipped scenario assets, launchers, tests. |
+| `sample_factory/` | Modified training engine used by COMRAD. |
+| `scripts/` | Figure-generation, result aggregation, WAD rendering, and experiment helper scripts. |
+| `results/` | Committed figures, tables, result summaries, and selected media artifacts. |
+| `DoomGen/` | Procedural generator used to build and vary scenario WADs. |
+
+Directory-level guides are provided in:
+
+- [`comrad/README.md`](comrad/README.md)
+- [`sample_factory/README.md`](sample_factory/README.md)
+- [`scripts/README.md`](scripts/README.md)
+- [`results/README.md`](results/README.md)
+
+## Setup
+
+COMRAD is packaged through `pyproject.toml` and is intended to be run with `uv`.
+
+```bash
+uv sync --all-extras
+uv run pre-commit install
+```
+
+The project requires Python 3.10 or newer. The package metadata and CI workflow are the authoritative sources for dependency resolution.
+
+## Single-run training
+
+The main training entrypoint is `comrad.train`.
+
+```bash
+uv run python -m comrad.train \
+  --env=stag_hunt_arena \
+  --algo=IPPO \
+  --train_for_env_steps=10000 \
+  --num_workers=4 \
+  --num_envs_per_worker=4 \
+  --device=cpu
+```
+
+Training outputs are written under `train_dir/<experiment>/`. The main runtime log is:
+
+```bash
+tail -n 120 train_dir/<experiment>/sf_log.txt
+```
+
+For launcher-owned benchmark sweeps, use `comrad/train_all.py` instead of manually reproducing those long commands.
+
+## Recording and evaluation
+
+High-resolution checkpoint recording:
+
+```bash
+uv run python -m comrad.record_video \
+  --env=stag_hunt_arena \
+  --train_dir=train_dir \
+  --experiment=<experiment_name> \
+  --load_checkpoint_kind=best \
+  --resolution=1280x720
+```
+
+Automap/top-down heatmap rendering:
+
+```bash
+uv run python -m comrad.record_topdown_heatmap \
+  --env=ammo_carrier \
+  --train_dir=<train_root> \
+  --experiment=<experiment_name> \
+  --output_dir=results/videos/ammo_carrier_navigation_heatmap \
+  --device=cpu \
+  --overwrite
+```
+
+Both commands expect a full experiment directory with at least `config.json` and the corresponding checkpoint directory.
+
+## Benchmark launchers
+
+COMRAD ships several paper-facing launcher profiles through `comrad/train_all.py`.
+
+- `benchmark`: full benchmark suite
+- `agent_scaling`: Armory Siege agent-count sweep
+- `platform_chain_curriculum`: curriculum comparison on Platform Chain
+- `qmix_lr_sensitivity_pilot`: Short sensitivity study
+- `qmix_lr_sensitivity_full`: Full sensitivity study
+
+Run a launcher profile locally with:
+
+```bash
+COMRAD_TRAIN_ALL_PROFILE=benchmark \
+uv run python -m sample_factory.launcher.run \
+  --run=comrad.train_all \
+  --backend=processes \
+  --max_parallel=1 \
+  --pause_between=1
+```
+
+The exact benchmark definitions, algorithms, root directory names, and profile constants live at the top of [`comrad/train_all.py`](comrad/train_all.py).
+
+## Scenario assets and DoomGen
+
+Runtime scenario assets are stored in [`comrad/scenarios/`](comrad/scenarios/README.md) as `.cfg` and `.wad` pairs, plus pre-generated batch directories such as `batch_platform_chain_curriculum/`.
+
+Generation workflows are in `DoomGen/`. The usual split is:
+
+- `DoomGen/examples/benchmark`: procedural generators and geometry logic
+- `comrad/scenarios/`: runtime WADs and configs consumed by training and evaluation
+
+If a scenario generator changes, rebuild the corresponding WAD in `DoomGen/` and copy the regenerated artifact back into `comrad/scenarios/` before training or evaluation.
+
+### Using DoomGen as a Python library
+
+DoomGen is declared as an optional dependency in the uv workspace. It is **not** installed by a plain `uv sync`:
+
+```bash
+uv sync --extra doomgen      # install doomgen + its deps
+# or
+uv run --extra doomgen python -m comrad.train ...   # one-off with doomgen available
+# or
+uv pip install -e DoomGen/   # editable install directly from local checkout
+```
+
+After any of these, `import doomgen` works from the root Python environment. Without `--extra doomgen`, the core COMRAD package remains self-contained and DoomGen is not pulled in.
+
+## Tests
+
+Run COMRAD and Sample Factory tests with `uv`:
+
+```bash
+uv run pytest comrad/tests -v
+uv run pytest sample_factory/tests -v
+```
+
+Tests in `comrad/tests/` focus on benchmark logic, launchers, and recorder behavior. Tests in `sample_factory/tests/` cover the modified training stack.
+
+## Outputs and artifacts
+
+- `train_dir/`: experiment directories, configs, checkpoints, logs
+- `results/`: committed figures, tables, result summaries, selected media
+- `wandb/`: local W&B run state and uploads
+
+See [`results/README.md`](results/README.md) for the structure of the committed artifacts.
+
+## License and upstream assets
+
+The repository is released under the MIT license in [`LICENSE`](LICENSE). Upstream and third-party assets are documented in the repository itself, including the bundled `scripts/wad2image/` tool and the paper's citations to ViZDoom, Sample Factory, HARL, PyMARL, and QPLEX.
+
+## Acknowledgements
+
+We would like to thank all researchers, engineers, advisors involved in this work; and the TU/e Supercomputing Center for providing the infrastructure to run the experiments. The environments are built on top of [ViZDoom](https://github.com/Farama-Foundation/ViZDoom). Our infrastructure is built on top of [Sample Factory](https://github.com/alex-petrenko/sample-factory). Our experiments were managed using [Weights & Biases](https://wandb.ai/). For on-policy actor-critic methods, we implement IPPO and MAPPO from the original implementation in [Sample Factory](https://github.com/alex-petrenko/sample-factory), and HAPPO following the [HARL codebase](https://github.com/PKU-MARL/HARL). For off-policy value decomposition methods, we implement IDQN, VDN, QMIX following the [PyMARL framework](https://github.com/oxwhirl/pymarl), and QPLEX following [official QPLEX reference implementation](https://github.com/wjh720/QPLEX).
+
+<!-- ## Citation -->

comrad_bench-1.0.0/comrad/curriculum/__init__.py

@@ -0,0 +1,30 @@
+from comrad.curriculum.base import Curriculum
+from comrad.curriculum.uniform import UniformCurriculum
+from comrad.curriculum.sequential import SequentialCurriculum
+from comrad.curriculum.learning_progress import LearningProgress
+from comrad.curriculum.omni import OMNICurriculum
+from comrad.curriculum.plr import PrioritizedLevelReplay
+
+def make_curriculum(n_tasks: int, strategy: str = "uniform", **kwargs) -> Curriculum:
+    """Factory to spin up proper curriculum based on configuration."""
+    strategies = {
+        "uniform": UniformCurriculum,
+        "sequential": SequentialCurriculum,
+        "learning_progress": LearningProgress,
+        "omni": OMNICurriculum,
+        "plr": PrioritizedLevelReplay,
+    }
+    
+    assert strategy in strategies, f"Unknown strategy '{strategy}'. Available: {list(strategies.keys())}"
+    
+    return strategies[strategy](n_tasks=n_tasks, **kwargs)
+
+__all__ = [
+    "Curriculum",
+    "UniformCurriculum",
+    "SequentialCurriculum",
+    "LearningProgress",
+    "OMNICurriculum",
+    "PrioritizedLevelReplay",
+    "make_curriculum"
+]

comrad_bench-1.0.0/comrad/curriculum/base.py

@@ -0,0 +1,133 @@
+import multiprocessing
+import os
+from typing import Optional
+import numpy as np
+
+
+class Curriculum:
+    """
+    Abstract base curriculum class for a fixed discrete batch of tasks (WADs).
+    Requires multiprocessing arrays and locks for disjoint worker/learner access.
+    """
+
+    EMA_ALPHA = 0.1
+    SCORE_MEAN_VALUE_L1   = "mean_value_l1"
+    SCORE_MAX_VALUE_L1    = "max_value_l1"
+    SCORE_MEAN_ADVANTAGE  = "mean_advantage"
+    SCORE_MAX_ADVANTAGE   = "max_advantage"
+    SCORE_MEAN_ENTROPY    = "mean_entropy"
+    SCORE_MAX_ENTROPY     = "max_entropy"
+
+    def __init__(self, n_tasks: int, seed: Optional[int] = None, tasks: Optional[list] = None, **kwargs):
+        self.ctx = multiprocessing.get_context("spawn")
+        self._n = n_tasks
+        self._base_seed = int(seed) if seed is not None else int.from_bytes(os.urandom(8), "little")
+        self._rng_counter = self.ctx.Value("Q", 0)
+        self._task_ids = tuple(tasks) if tasks is not None else None
+
+        self._weights = self.ctx.Array('d', [1.0 / n_tasks] * n_tasks)
+        self._counts = self.ctx.Array('l', [0] * n_tasks)
+        self._lock = self.ctx.Lock()
+
+    def update(self, task_idx: int, episode_return: float) -> None:
+        """Called at end of each episode."""
+        with self._lock:
+            self._counts[task_idx] += 1
+            self._update_logic(task_idx, float(episode_return))
+            self._recompute_weights()
+
+    def update_task_score(self, task_idx: int, score: dict, num_steps: int = 1) -> None:
+        """Used by sample factory learners (e.g. PLR)"""
+        pass
+
+    def sample(self) -> int:
+        with self._lock:
+            task_idx = self._sample_logic()
+        return task_idx
+
+    def _sample_logic(self) -> int:
+        """Isolated locked core for actually rolling the random sample"""
+        weights = np.array(self._weights[:])
+        return self._choice(self._n, p=weights)
+
+    def _next_rng(self) -> np.random.Generator:
+        """
+        Generate randomness from a single shared draw counter so spawned workers
+        do not each replay an identical local RNG stream.
+        """
+        counter = self._rng_counter.value
+        self._rng_counter.value = counter + 1
+        seed = (self._base_seed + 0x9E3779B97F4A7C15 * counter) & ((1 << 64) - 1)
+        return np.random.default_rng(seed)
+
+    def _choice(self, n: int, p=None) -> int:
+        return int(self._next_rng().choice(n, p=p))
+
+    def _random(self) -> float:
+        return float(self._next_rng().random())
+
+    def _update_logic(self, task_idx: int, episode_return: float) -> None:
+        """Sub-classes override this to implement update behavior."""
+        pass
+
+    def _recompute_weights(self) -> None:
+        """Sub-classes override this to update probabilities inside self._weights."""
+        pass
+
+    def task_weights(self) -> np.ndarray:
+        with self._lock:
+            return np.array(self._weights[:])
+
+    def task_returns(self) -> np.ndarray:
+        """Return fast moving average or appropriate returns proxy."""
+        # By default just returns zeroes, overridden in LP/OMNI classes
+        return np.zeros(self._n)
+
+    def task_scores(self) -> np.ndarray:
+        """PLR scores or similar metric."""
+        return np.zeros(self._n)
+
+    def metrics(self) -> dict:
+        """General metrics reporting."""
+        return {}
+
+    def state_dict(self) -> dict:
+        with self._lock:
+            state = {
+                "version": 1,
+                "curriculum_type": type(self).__name__,
+                "n_tasks": self._n,
+                "base_seed": int(self._base_seed),
+                "rng_counter": int(self._rng_counter.value),
+                "weights": list(self._weights[:]),
+                "counts": list(self._counts[:]),
+            }
+            if self._task_ids is not None:
+                state["task_ids"] = list(self._task_ids)
+            state.update(self._state_dict_locked())
+            return state
+
+    def load_state_dict(self, state: dict) -> None:
+        expected_type = type(self).__name__
+        actual_type = state.get("curriculum_type", expected_type)
+        if actual_type != expected_type:
+            raise ValueError(f"Expected curriculum state for {expected_type}, got {actual_type}")
+        if int(state.get("n_tasks", self._n)) != self._n:
+            raise ValueError(f"Expected curriculum state for {self._n} tasks, got {state.get('n_tasks')}")
+        expected_task_ids = list(self._task_ids) if self._task_ids is not None else None
+        actual_task_ids = state.get("task_ids", None)
+        if expected_task_ids is not None and actual_task_ids != expected_task_ids:
+            raise ValueError(f"Expected curriculum task ids {expected_task_ids}, got {actual_task_ids}")
+
+        with self._lock:
+            self._base_seed = int(state.get("base_seed", self._base_seed))
+            self._rng_counter.value = int(state.get("rng_counter", 0))
+            self._weights[:] = [float(weight) for weight in state.get("weights", self._weights[:])]
+            self._counts[:] = [int(count) for count in state.get("counts", self._counts[:])]
+            self._load_state_dict_locked(state)
+
+    def _state_dict_locked(self) -> dict:
+        return {}
+
+    def _load_state_dict_locked(self, state: dict) -> None:
+        pass