Ryukijano/marl-world-model

MARL + World Model for Warehouse Cube Sorting in Isaac Sim

Independent PPO (IPPO) multi-agent policies with a pluggable world-model wrapper for color-coded cube sorting in NVIDIA Isaac Sim, with an SO-101 robot bridge and LeRobot export.

This is a research demo that trains 2–4 differential-drive mobile manipulators to pick red, green and blue cubes and place them in matching bins. Each agent has its own on-board RGB camera and a continuous 4-D action space. A learned world model can optionally generate imagined rollouts for data augmentation; on this Windows workstation it falls back to a tiny mock dynamics network when the Cosmos3 / DreamDojo backends are not available.

What it is

• Scenario: warehouse floor with colored cubes and matching bins.
• Agents: 2–4 small robots with onboard 128×128 RGB cameras and a sticky gripper.
• Algorithm: Independent Proximal Policy Optimisation (IPPO) — one PPO actor-critic per agent.
• World model: pluggable mock / cosmos3 / dreamdojo backend that predicts the next observation + reward.
• Sim-to-real bridge: so101_bridge.py extracts one agent’s policy and deploys it on a real SO-101 6-DOF arm.
• Dataset export: convert_to_lerobot.py turns rollouts into LeRobot v2.1 format for warm-starting GR00T / SmolVLA policies.

Model architecture

The per-agent policy is a CNN+MLP actor-critic defined in marl_ippo.py. Full metadata is in config.json.

Component	Value
Number of agents	2 (configurable up to 4)
Observation image	128 × 128 RGB, stored as (C, H, W) = (3, 128, 128)
State dimension	12
State vector	[x, y, yaw, vx, vy, omega, nearest_cube_dx, nearest_cube_dy, nearest_cube_dist, target_bin_dx, target_bin_dy, target_bin_dist]
Action space	Continuous 4-D in [-1, 1]: [forward, lateral, angular, gripper]
Image encoder	3 stride-2 conv layers (32, 64, 64) + AdaptiveAvgPool2d(4, 4)
Shared MLP	conv_flat + state_dim -> 256 -> 256 with ReLU
Value head	256 -> 1
Actor head	256 -> 4 Gaussian mean + log-std (continuous)
Hidden size	256
Default PPO	lr=3e-4, gamma=0.99, gae_lambda=0.95, clip_eps=0.2, update_epochs=4, batch_size=64

Files in this repo

File	Purpose
env_multi_agent_isaac.py	Multi-agent Isaac Sim warehouse scene (robots, cubes, bins, cameras, rewards, headless D3D12).
marl_ippo.py	Independent PPO policy and trainer (SimpleActorCritic, PPOAgent, IPPOTrainer).
world_model.py	Pluggable world-model wrapper: mock tiny network, cosmos3 (NIM / direct / subprocess) and dreamdojo (direct / subprocess) with transparent fallback.
run_marl.py	End-to-end training/evaluation runner with CLI for agents, episodes, world model, imagination, and checkpoints.
so101_bridge.py	Deploy one agent from a checkpoint to the SO-101 arm, in --simulate or --deploy mode.
convert_to_lerobot.py	Convert saved rollouts or a checkpoint into LeRobot v2.1 parquet+video datasets.
DESIGN.md	Original architecture and scenario design document.
AGENTS.md	Canonical map of the broader H:\Robotics physical-AI workspace.
outputs/checkpoints/	Trained IPPO checkpoints, e.g. marl_ippo_ep100.pt (Git-LFS).
outputs/logs/marl_runs.jsonl	Per-episode reward and timing logs.
config.json	Model-card style metadata (shapes, hyperparameters, backend list).
README.md	This file.

Installation

1. Clone this repo

   git clone https://huggingface.co/Ryukijano/marl-world-model
   cd marl-world-model
   

2. NVIDIA Isaac Sim

   The code is tested with Isaac Sim 5.1.0 standalone and the Isaac Sim 6.0.1 pip environment.

   • 5.1.0 standalone: use the bundled python.bat (e.g. H:\Robotics\ISSAC_SIM_5.1.0\python.bat).
   • 6.0.1 pip: create/activate the isaac-sim-601 conda env and set OMNI_KIT_ACCEPT_EULA=YES.

3. General Python deps (for so101_bridge.py and convert_to_lerobot.py)

   pip install torch numpy opencv-python pillow imageio requests pandas
   

4. SO-101 / LeRobot deps (only for real-robot bridge)

   pip install lerobot
   

How to run

All training and simulation commands below use Isaac Sim 5.1.0. For the 6.0.1 pip env, replace H:\Robotics\ISSAC_SIM_5.1.0\python.bat with conda run -n isaac-sim-601 python (and set OMNI_KIT_ACCEPT_EULA=YES before running).

Smoke test

H:\Robotics\ISSAC_SIM_5.1.0\python.bat run_marl.py --smoke --device cpu

This launches the multi-agent scene, takes random actions for a few steps, and saves frames.

Short IPPO sanity run

H:\Robotics\ISSAC_SIM_5.1.0\python.bat run_marl.py --num-episodes 2 --max-steps-per-episode 30 --batch-size 16 --device cpu

Full 100-episode training

H:\Robotics\ISSAC_SIM_5.1.0\python.bat run_marl.py --num-episodes 100 --max-steps-per-episode 200 --batch-size 64 --device cuda --checkpoint-every 20

Checkpoints are written to outputs/checkpoints/marl_ippo_ep{N}.pt.

IPPO with world-model imagination

H:\Robotics\ISSAC_SIM_5.1.0\python.bat run_marl.py --num-episodes 1 --max-steps-per-episode 20 --batch-size 16 --use-imagination --imagination-horizon 5 --device cpu

Cosmos3 / DreamDojo backend

:: Will try Cosmos3; on Windows it falls back to mock if the backend is unavailable
H:\Robotics\ISSAC_SIM_5.1.0\python.bat run_marl.py --num-episodes 1 --max-steps-per-episode 5 --batch-size 4 --world-model cosmos3 --device cpu

:: Same for DreamDojo
H:\Robotics\ISSAC_SIM_5.1.0\python.bat run_marl.py --num-episodes 1 --max-steps-per-episode 5 --batch-size 4 --world-model dreamdojo --device cpu

4-agent training

H:\Robotics\ISSAC_SIM_5.1.0\python.bat run_marl.py --num-agents 4 --num-episodes 1 --max-steps-per-episode 10 --batch-size 8 --device cpu

Using the SO-101 bridge

The bridge is not run inside Isaac Sim. It loads a single agent from a marl_ippo_ep*.pt checkpoint and runs on the SO-101 (or a dummy robot in simulation).

Simulation mode (no hardware required)

conda run -n gr00t python so101_bridge.py --checkpoint outputs/checkpoints/marl_ippo_ep100.pt --simulate --max-steps 20

Real hardware mode

conda run -n gr00t python so101_bridge.py --checkpoint outputs/checkpoints/marl_ippo_ep100.pt --deploy --follower-port COM7 --cam-idx 0

The bridge preprocesses the camera frame, runs the IPPO policy, and maps the 4-D continuous action to SO-101 joint commands in --action-mode delta (default) or direct. The first three action components typically map to shoulder/elbow/wrist joints and the last component maps to the gripper.

Link to the LeRobot dataset

Generated rollouts can be exported to the companion LeRobot dataset repo:

Ryukijano/marl-world-model-lerobot

Create a local LeRobot v2.1 dataset from the trained policy:

conda run -n gr00t python convert_to_lerobot.py --checkpoint outputs/checkpoints/marl_ippo_ep100.pt --num-episodes 10 --episode-length 200 --output-dir so101_marl_dataset --task "pick up the cube"

You can then upload so101_marl_dataset to the LeRobot repo above for GR00T / SmolVLA fine-tuning.

Known limitations

• Cosmos3 / DreamDojo on Windows: these backends require Linux-native Cosmos / DreamDojo installs or a running NIM HTTP server. On the Windows development workstation they transparently fall back to the mock tiny world model.
• 100 episodes is not enough for full cube sorting: the runs are intended as a proof-of-concept. Reliable sorting will need far more episodes / curriculum / fine-tuning.
• Isaac Sim 5.1.0 headless: the renderer runs with D3D12 and vulkan=false because the local GPU driver is 595.79. GUI mode may crash.
• Shutdown workaround: run_marl.py calls os._exit(0) after writing outputs to avoid a known omni.syntheticdata.plugin.dll shutdown crash during Py_FinalizeEx.
• SO-101 mapping is heuristic: the trained agent has a 4-D base action space, while the SO-101 has 6 joints. The bridge pads/maps actions and uses simple delta or direct control.
• Sticky gripper: cube grasping is approximated by a distance threshold and gripper action > 0.5; there is no explicit articulated gripper in the base environment.
• Center-cropped camera images: 128×128 frames are cropped/padded from the raw camera output.

License and citation

This repository is released under the Apache-2.0 License. Please check individual third-party dependencies (Isaac Sim, LeRobot, Cosmos, DreamDojo) for their respective licenses.

If you use this code in your research, please cite it as:

@software{ryukijano_marl_world_model,
  title={MARL + World Model for Warehouse Cube Sorting in Isaac Sim},
  author={Ryukijano},
  year={2026},
  url={https://huggingface.co/Ryukijano/marl-world-model}
}