This repository unifies two papers that share a common philosophy — planning is more robust when multiple perspectives on the world model are explored in parallel — into a single, easy-to-run codebase.

Both methods are model-based, target visual continuous control, and are compared against the Dreamer / PlaNet family.

Classical model-based RL is brittle in two ways:

1. World-model bias — a single learned dynamics model is treated as ground truth.
2. Singular-policy bias — one policy is rolled out inside that model.

EPN and MSPP attack the second source of bias. Instead of trusting one policy, they roll out a population of policies with different scales / pruning views, then aggregate. The result is more reliable planning inside the same imperfect world model.

📄 Full paper PDF is included in the repo: Understanding world models through multi-step pruning policy via reinforcement learning.pdf

Verified on Python 3.10 – 3.12 with modern PyTorch (≥ 2.x). The exact dependency set lives in requirements.txt.

git clone https://github.com/tinyzqh/MSPP.git
cd MSPP
# 1. fresh environment (conda or venv — pick one)
conda create -n mspp python=3.10 -y && conda activate mspp
# or: python -m venv .venv && source .venv/bin/activate
# 2. core training stack
pip install -r requirements.txt
# 3. (optional) extras only needed for utils_plot.py post-hoc plotting
pip install -r requirements-plot.txt

• GPU: any CUDA-capable PyTorch build works. Multi-GPU is auto-enabled via torch.nn.DataParallel when more than one CUDA device is visible.
• CPU only: pass --disable-cuda. Training works but is orders of magnitude slower; useful for smoke-tests.

export MUJOCO_GL=osmesa # or 'egl' on machines with EGL
python asynchronous_main.py --symbolic-env ...

DM Control will otherwise try to open a GLFW window and emit GLFWError: DISPLAY environment variable is missing.

After install, this should finish in well under a minute on CPU and confirm everything is wired up:

python asynchronous_main.py \
 --algo dreamer --env cartpole-balance --symbolic-env \
 --episodes 3 --seed-episodes 1 --collect-interval 2 \
 --batch-size 2 --chunk-size 10 --max-episode-length 50 \
 --disable-cuda --test-interval 999

You should see two training loop + Data collection blocks and EXIT 0, and results/.../models_*.pth written to disk.

The entry point is asynchronous_main.py. Algorithms are selected with --algo:

python asynchronous_main.py \
 --algo aap \
 --env acrobot-swingup \
 --pool_len 3 \
 --top_planning-horizon 8 \
 --seed 5

python asynchronous_main.py --algo dreamer --env cartpole-balance --seed 5

python asynchronous_main.py --algo aap --env acrobot-swingup \
 --models results/.../models_500.pth

python asynchronous_main.py --test --models path/to/checkpoint.pth --env hopper-hop

Each run writes to results/{env}_seed_{seed}_{algo}_..._{top_planning_horizon}/:

The results/ tree is .gitignored.

Multiple parallel sub-agents reach high returns where a single-view planner stalls.

Head-to-head comparison against Dreamer, the state-of-the-art model-based baseline at the time.

Detailed curves, ablations, and the tabular convergence analysis are in the Information Sciences paper and the bundled PDF.

MSPP/
├── algorithms/
│ ├── dreamer.py # Dreamer baseline
│ ├── planet.py # PlaNet baseline
│ ├── plan_to_plan.py # p2p variant
│ ├── actor_pool_1.py # Actor-pool ablation
│ └── asynchronous_actor_planet.py # EPN / MSPP — multi-scale parallel sub-agents
├── asynchronous_main.py # Entry point — orchestrates training
├── asynchronous_actor.py # Async actor worker
├── asynchronous_init_sample.py # Seed-episode collection
├── planner.py # CEM planner
├── models.py # Transition / Observation / Reward / Encoder / Actor / Value
├── memory.py # Experience replay
├── env.py # DM Control + Gym wrappers
├── parameter.py # All CLI hyper-parameters
├── utils.py / utils_plot.py # Helpers & plotting
└── figures/ # Result plots used above

The most relevant flags (see parameter.py for the full list):

This fork has been audited and patched against the original release. The most user-visible changes:

• Checkpoints are now actually saved. The original repo had torch.save(model_state, ...) commented out in save_model_data, so multi-day runs left no resumable artefact. Resolved — each checkpoint_interval writes models_{episode}.pth containing the world model, per-algorithm actor/value pools, optimizer state, and metrics.
• Inter-process tensor transfer no longer hits disk. Previously every collect step wrote actor_states.pt + actor_beliefs.pt and the actor workers torch.loaded them back; replaced with direct pipe-based transfer (auto-IPC via PyTorch multiprocessing).
• Seed-episode workers now use distinct seeds — the original passed the same args.seed to every worker, so the "S random seed episodes" were S identical trajectories.
• data_collection now uses exploration noise (explore=True) — the original silently ran the deterministic policy, hurting buffer diversity.
• Actor workers shut down cleanly via a sentinel on the pipe — the original orphaned them at process exit (visible as EOFError tracebacks on shutdown).
• CEM std-dev is clamped to 1e-6 so collapsed elites don't produce Normal(scale=0) and get rejected by modern PyTorch.
• --action-scale defaults to 1 (was -1, which made imagine_ahead pop from an empty list and crash on the dreamer path).
• CPU-only path works — several .cuda() calls inside MPCPlanner are now .to(args.device).
• Misc: TanhBijector no longer instantiates the abstract constraints.Constraint; is not '' → != ''; tqdm total/initial are correct; per-step accounting uses t * action_repeat consistently; cuDNN determinism is enabled by default.

See git history for the precise diffs.

• The repo was originally engineered for multi-GPU; the MultiGPU chunk/concat trick inside Plan.upper_transition_model has not been re-tested under the modern stack — keep an eye on it if you train on multiple cards.
• gym env names are still pinned to a Pendulum-v0-era list — if you use Gym envs, update env.py to whatever your installed gym/gymnasium exposes.

If this code or these ideas helped your research, please cite both papers.

@article{he2024understanding,
 title = {Understanding World Models through Multi-Step Pruning Policy via Reinforcement Learning},
 author = {He, Zhiqiang and Qiu, Wen and Zhao, Wei and Shao, Xun and Liu, Zhi},
 journal = {Information Sciences},
 pages = {121361},
 year = {2024},
 publisher = {Elsevier}
}
@article{wang2022erlang,
 title = {Erlang planning network: An iterative model-based reinforcement learning with multi-perspective},
 author = {Wang, Jiao and Zhang, Lemin and He, Zhiqiang and Zhu, Can and Zhao, Zihui},
 journal = {Pattern Recognition},
 volume = {128},
 pages = {108668},
 year = {2022},
 publisher = {Elsevier}
}

Built on top of the open-source PlaNet and Dreamer line of work, and the DeepMind Control Suite.

Questions, issues, and PRs are welcome. For paper-related discussion, please open a GitHub issue or reach the corresponding author through the publication pages linked above.