Add README.md and config.json with architecture/metadata

README.md

@@ -0,0 +1,200 @@
+---
+tags:
+- marl
+- multi-agent-reinforcement-learning
+- ppo
+- ipppo
+- isaac-sim
+- world-model
+- cosmos3
+- dreamdojo
+- warehouse-robotics
+- cube-sorting
+- so-101
+- lerobot
+license: apache-2.0
+library_name: custom
+---
+
+# 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`](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**
+
+   ```bash
+   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`)
+
+   ```bash
+   pip install torch numpy opencv-python pillow imageio requests pandas
+   ```
+
+4. **SO-101 / LeRobot deps** (only for real-robot bridge)
+
+   ```bash
+   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
+
+```batch
+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
+
+```batch
+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
+
+```batch
+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
+
+```batch
+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
+
+```batch
+:: 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
+
+```batch
+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)
+
+```batch
+conda run -n gr00t python so101_bridge.py --checkpoint outputs/checkpoints/marl_ippo_ep100.pt --simulate --max-steps 20
+```
+
+### Real hardware mode
+
+```batch
+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](https://huggingface.co/datasets/Ryukijano/marl-world-model-lerobot)**
+
+Create a local LeRobot v2.1 dataset from the trained policy:
+
+```batch
+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:
+
+```bibtex
+@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}
+}
+```

config.json

@@ -0,0 +1,99 @@
+{
+  "model_type": "marl-ippo-actor-critic",
+  "name": "MARL IPPO Warehouse Cube Sorting Policy",
+  "description": "Independent PPO policies for 2-4 mobile warehouse robots in NVIDIA Isaac Sim, with optional world-model imagination.",
+  "task": "multi-agent-warehouse-cube-sorting",
+  "simulator": {
+    "name": "NVIDIA Isaac Sim",
+    "versions_tested": ["5.1.0", "6.0.1 pip"],
+    "headless": true,
+    "renderer": "D3D12"
+  },
+  "num_agents": 2,
+  "observation": {
+    "image": {
+      "shape": [3, 128, 128],
+      "dtype": "uint8",
+      "description": "RGB camera image per agent, stored/transposed to CHW for the policy network."
+    },
+    "state": {
+      "shape": [12],
+      "dtype": "float32",
+      "description": "[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": {
+    "type": "continuous",
+    "shape": [4],
+    "low": -1.0,
+    "high": 1.0,
+    "description": "[forward, lateral, angular, gripper]",
+    "gripper_threshold": 0.5
+  },
+  "policy": {
+    "algorithm": "IPPO",
+    "architecture": {
+      "image_encoder": {
+        "type": "CNN",
+        "input_shape": [3, 128, 128],
+        "conv_layers": [
+          {"in_channels": 3, "out_channels": 32, "kernel_size": 3, "stride": 2, "padding": 1},
+          {"in_channels": 32, "out_channels": 64, "kernel_size": 3, "stride": 2, "padding": 1},
+          {"in_channels": 64, "out_channels": 64, "kernel_size": 3, "stride": 2, "padding": 1}
+        ],
+        "pooling": "AdaptiveAvgPool2d(4, 4)"
+      },
+      "shared_mlp": {
+        "input_dim": "conv_flat + state_dim",
+        "hidden_dim": 256,
+        "num_layers": 2,
+        "activation": "ReLU"
+      },
+      "value_head": {"input_dim": 256, "output_dim": 1},
+      "actor_head": {
+        "type": "gaussian",
+        "input_dim": 256,
+        "output_dim": 4
+      },
+      "state_dim": 12,
+      "hidden_dim": 256,
+      "cnn_channels": [32, 64, 64],
+      "reduce_spatial": 4
+    },
+    "hyperparameters": {
+      "lr": 3e-4,
+      "gamma": 0.99,
+      "gae_lambda": 0.95,
+      "clip_eps": 0.2,
+      "value_loss_coef": 0.5,
+      "entropy_coef": 0.01,
+      "update_epochs": 4,
+      "batch_size": 64,
+      "max_grad_norm": 0.5,
+      "normalize_advantages": true
+    }
+  },
+  "world_model": {
+    "backends": ["mock", "cosmos3", "dreamdojo"],
+    "default": "mock",
+    "fallback_to_mock": true,
+    "imagination_horizon_default": 20
+  },
+  "so101_bridge": {
+    "input_image_size": [128, 128],
+    "state_dim_policy": 12,
+    "action_mode": "delta",
+    "max_delta": 0.05,
+    "joint_keys": [
+      "shoulder_pan.pos",
+      "shoulder_lift.pos",
+      "elbow_flex.pos",
+      "wrist_flex.pos",
+      "wrist_roll.pos",
+      "gripper.pos"
+    ]
+  },
+  "lerobot_dataset_repo": "Ryukijano/marl-world-model-lerobot",
+  "source_repo": "Ryukijano/marl-world-model",
+  "license": "apache-2.0"
+}