TamasTran · Sep 25, 2026 · Sep 25, 2026 · Sep 25, 2026 · Sep 25, 2026 · Sep 25, 2026
diff --git a/GUIDE.md b/GUIDE.md
diff --git a/README.md b/README.md
@@ -1,47 +1,98 @@
 ![Project banner](https://raw.githubusercontent.com/mujocolab/mjlab/main/docs/source/_static/mjlab-banner.jpg)
 
-# mjlab
+# mjlab_MPOPI
 
-[![GitHub Actions](https://img.shields.io/github/actions/workflow/status/mujocolab/mjlab/ci.yml?branch=main)](https://github.com/mujocolab/mjlab/actions/workflows/ci.yml?query=branch%3Amain)
-[![Documentation](https://github.com/mujocolab/mjlab/actions/workflows/docs.yml/badge.svg)](https://mujocolab.github.io/mjlab/)
-[![License](https://img.shields.io/github/license/mujocolab/mjlab)](https://github.com/mujocolab/mjlab/blob/main/LICENSE)
+[![Based on mjlab](https://img.shields.io/badge/based_on-mjlab-blue)](https://github.com/mujocolab/mjlab)
 [![MuJoCo Warp](https://img.shields.io/badge/MuJoCo_Warp-3.11.0-blue)](https://github.com/google-deepmind/mujoco_warp/releases/tag/v3.11.0)
-[![Nightly Benchmarks](https://img.shields.io/badge/Nightly-Benchmarks-blue)](https://mujocolab.github.io/mjlab/nightly/)
-[![PyPI](https://img.shields.io/pypi/v/mjlab)](https://pypi.org/project/mjlab/)
-[![PyPI downloads](https://img.shields.io/pypi/dm/mjlab?color=blue)](https://pypistats.org/packages/mjlab)
+[![License](https://img.shields.io/github/license/mujocolab/mjlab)](LICENSE)
 
-mjlab combines [Isaac Lab](https://github.com/isaac-sim/IsaacLab)'s manager-based API with [MuJoCo Warp](https://github.com/google-deepmind/mujoco_warp), a GPU-accelerated version of [MuJoCo](https://github.com/google-deepmind/mujoco).
-The framework provides composable building blocks for environment design,
-with minimal dependencies and direct access to native MuJoCo data structures.
+This repository is [mjlab](https://github.com/mujocolab/mjlab) plus a research package,
+**`mpopi_train`**, that tries to **speed up policy training with control algorithms**: PPO reuses
+its own past rollouts with importance correction (**Replay-IS**), and a sampling MPC planner
+(MPPI / MPOPI) acts as a teacher that labels the policy's states (**DAgger**).
+
+mjlab combines [Isaac Lab](https://github.com/isaac-sim/IsaacLab)'s manager-based API with
+[MuJoCo Warp](https://github.com/google-deepmind/mujoco_warp), a GPU-accelerated version of
+[MuJoCo](https://github.com/google-deepmind/mujoco). `mpopi_train` uses it as a library:
+**nothing under `src/mjlab` is modified**, so upstream mjlab updates merge cleanly.
+
+> **New to the project?** Follow the step-by-step guide in Vietnamese: **[GUIDE.md](GUIDE.md)**.
 
 ## Getting Started
 
-mjlab requires an NVIDIA GPU for training. macOS is supported for evaluation only.
+Training requires an NVIDIA GPU (about 12 GB of memory for 4096 robots). Without one, use the
+Kaggle notebook (see [Notebooks](#notebooks)).
 
-**Try it now:**
+The research code is on the **`mpc-stage1`** branch:
+
+```bash
+git clone --branch mpc-stage1 https://github.com/TamasTran/mjlab_MPOPI.git && cd mjlab_MPOPI
+uv sync --extra cu128
+```
 
-Run the demo (no installation needed):
+mjlab's own demo still works:
 
 ```bash
-uvx --from mjlab --refresh demo
+uv run --extra cu128 demo
 ```
 
-Or try in [Google Colab](https://colab.research.google.com/github/mujocolab/mjlab/blob/main/notebooks/demo.ipynb) (no local setup required).
+## Training Examples
+
+### 1. Velocity Tracking with the Compared Methods
+
+The task is mjlab's flat Unitree G1 velocity task for **2000 iterations**, with a command
+curriculum that moves from (-1, 1) m/s to (-1, 1.5) m/s at iteration 500. Each method is a task:
 
-**Install from source:**
+| Task | Method |
+|---|---|
+| `Mpopi-G1-2k-PPO` | RSL-RL PPO (baseline) |
+| `Mpopi-G1-2k-Replay-IS` | PPO plus its last 4 rollouts, importance-corrected (clipped weights, V-trace) |
+| `Mpopi-G1-2k-DAgger` | PPO plus behavior cloning toward an MPOPI planner that labels the policy's own states |
+| `Mpopi-G1-2k-Replay-IS-DAgger` | Both, in separate buffers |
 
 ```bash
-git clone https://github.com/mujocolab/mjlab.git && cd mjlab
-uv run demo
+uv run --extra cu128 mpopi-train Mpopi-G1-2k-Replay-IS-DAgger --agent.logger tensorboard \
+  --agent.seed 1 --agent.run-name Replay-IS-DAgger_s1
 ```
 
-For alternative installation methods (PyPI, Docker), see the [Installation Guide](https://mujocolab.github.io/mjlab/main/source/installation.html).
+`mpopi-train` is mjlab's `train` with these tasks registered, so every mjlab option works. Method
+settings are in [`src/mpopi_train/presets.py`](src/mpopi_train/presets.py) and can be overridden,
+for example `--agent.algorithm.mpopi.mpc.num-envs 32`.
 
-## Training Examples
+**Several seeds** are several runs (results vary noticeably between runs, so compare at least two):
+
+```bash
+for s in 1 2 3; do
+  uv run --extra cu128 mpopi-train Mpopi-G1-2k-PPO --agent.logger tensorboard \
+    --agent.seed $s --agent.run-name PPO_s$s
+done
+```
 
-### 1. Velocity Tracking
+The tasks are registered through the `mjlab.tasks` entry point, so mjlab's own `train` and
+`play` commands (and the worker processes of multi-GPU training) see them too. Multi-GPU training
+of one run (`--gpu-ids` with more than one GPU) is untested; the experiments ran one training per
+GPU with `CUDA_VISIBLE_DEVICES`.
 
-Train a Unitree G1 humanoid to follow velocity commands on flat terrain:
+**Evaluate** a checkpoint at fixed forward speeds (64 robots, measured over 10 s after 1 s of
+acceleration). The PPO task is used for every method: same robot and network, and no MPC teacher
+is built:
+
+```bash
+uv run --extra cu128 mpopi-eval --task Mpopi-G1-2k-PPO --controllers policy \
+  --num-envs 64 --steps 550 --settle-steps 50 \
+  --checkpoint logs/rsl_rl/g1_velocity_2k/<run>/model_1999.pt
+```
+
+**Watch** a trained policy:
+
+```bash
+uv run --extra cu128 mpopi-play Mpopi-G1-2k-PPO \
+  --checkpoint-file logs/rsl_rl/g1_velocity_2k/<run>/model_1999.pt
+```
+
+### 2. mjlab's Original Velocity Task
+
+Train a Unitree G1 humanoid to follow velocity commands on flat terrain (PPO, 30,000 iterations):
 
 ```bash
 uv run train Mjlab-Velocity-Flat-Unitree-G1 --env.scene.num-envs 4096
@@ -63,7 +114,7 @@ Evaluate a policy while training (fetches latest checkpoint from Weights & Biase
 uv run play Mjlab-Velocity-Flat-Unitree-G1 --wandb-run-path your-org/mjlab/run-id
 ```
 
-### 2. Motion Imitation
+### 3. Motion Imitation
 
 Train a humanoid to mimic reference motions. See the [motion imitation guide](https://mujocolab.github.io/mjlab/main/source/training/motion_imitation.html) for preprocessing setup.
 
@@ -72,7 +123,7 @@ uv run train Mjlab-Tracking-Flat-Unitree-G1 --registry-name your-org/motions/mot
 uv run play Mjlab-Tracking-Flat-Unitree-G1 --wandb-run-path your-org/mjlab/run-id
 ```
 
-### 3. Sanity-check with Dummy Agents
+### 4. Sanity-check with Dummy Agents
 
 Use built-in agents to sanity check your MDP before training:
 
@@ -83,10 +134,31 @@ uv run play Mjlab-Your-Task-Id --agent random  # Sends uniform random actions
 
 When running motion-tracking tasks, add `--registry-name your-org/motions/motion-name` to the command.
 
+## Project Layout
+
+| Path | Content |
+|---|---|
+| [`src/mpopi_train/`](src/mpopi_train/) | The research package: algorithms, MPC, runner, presets, tasks, scripts ([README](src/mpopi_train/README.md)) |
+| [`docs/mpopi/`](docs/mpopi/) | Design notes |
+| [`notebooks/`](notebooks/) | Colab and Kaggle notebooks |
+| [`tests/`](tests/) (`test_mpopi_*.py`, `test_mpc_*.py`) | Tests of the package |
+| `src/mjlab/` | Unmodified mjlab |
+
+### Notebooks
+
+| Notebook | Purpose |
+|---|---|
+| `notebooks/mpc_g1_train_kaggle.ipynb` | Train and evaluate the four methods on Kaggle (2 T4 GPUs, up to 12 h) |
+| `notebooks/mpc_g1_eval_kaggle.ipynb` | Re-evaluate checkpoints of earlier Kaggle runs |
+| `notebooks/mpc_g1_bc_kaggle.ipynb` | How long DAgger should clone the MPC teacher (3 variants, 2 seeds) |
+| `notebooks/mpc_g1_gate_kaggle.ipynb` | Clone only while the teacher beats the policy (3 variants, 2 seeds) |
+| Other `mpc_*` and `mpopi_*` notebooks | Earlier experiments; they clone the tag `mpopi-before-module` |
 
 ## Documentation
 
-Full documentation is available at **[mujocolab.github.io/mjlab](https://mujocolab.github.io/mjlab/)**.
+- Step-by-step guide (Vietnamese): **[GUIDE.md](GUIDE.md)**
+- Package reference: [`src/mpopi_train/README.md`](src/mpopi_train/README.md)
+- mjlab documentation: **[mujocolab.github.io/mjlab](https://mujocolab.github.io/mjlab/)**
 
 ## Development
 
@@ -97,6 +169,12 @@ make format        # Format and lint
 make docs          # Build docs locally
 ```
 
+Tests of the research package only:
+
+```bash
+uv run --extra cu128 pytest tests/test_mpopi_*.py tests/test_mpc_*.py
+```
+
 For development setup: `uvx pre-commit install`
 
 ## Citation
@@ -117,6 +195,11 @@ If you use mjlab in your research, please consider citing:
 }
 ```
 
+The methods in `mpopi_train` build on PPO (Schulman et al., 2017,
+[arXiv:1707.06347](https://arxiv.org/abs/1707.06347)), V-trace (Espeholt et al., 2018,
+[arXiv:1802.01561](https://arxiv.org/abs/1802.01561)), DAgger (Ross et al., 2011,
+[arXiv:1011.0686](https://arxiv.org/abs/1011.0686)) and MPPI (Williams et al., 2017).
+
 ## License
 
 mjlab is licensed under the [Apache License, Version 2.0](LICENSE).

diff --git a/docs/mpopi/mpc_g1_velocity.md b/docs/mpopi/mpc_g1_velocity.md
@@ -0,0 +1,44 @@
+# Can sampling MPC walk G1 at 1.5 m/s? (gate before MPC → PPO on G1)
+
+Date: 2026-09-29. Branch `mpc-stage1`.
+
+MPC-DAgger and MPC-Inject need an MPC teacher that can do the task. Before
+comparing them with PPO on `Mjlab-Velocity-Flat-Unitree-G1` (random velocity
+commands that include 1.5 m/s), this checks the MPC alone.
+
+## Prerequisites done
+
+- **State sync** (`mpopi_train.mpc.state_sync`): the planner copies commands and
+  timers, previous actions, stateful reward terms, sensor histories, entity
+  data and domain-randomized model fields. On G1 the planner's per-term
+  rewards match the real env to about 1e-5; copying only the simulator state
+  gave errors of order 1 (`tests/test_mpc_state_sync.py`).
+- **Smooth planner noise** (`num_knots`): with independent per-step noise every
+  perturbed sequence scored far below the unperturbed plan (0.24 vs 0.72 over
+  10 steps at noise 0.3), mainly from the action-rate penalty over 29 joints.
+  Knot noise and a smaller scale narrow the gap (best perturbed 0.70 vs 0.72 at
+  noise 0.1, 3 knots, 16 samples, CPU). No perturbed sequence beat standing
+  still in that CPU probe: starting to walk first costs posture and upright
+  reward before velocity tracking pays off, a local optimum for a short
+  horizon.
+
+## Pre-registration (written before the GPU run)
+
+`src/mpopi_train/scripts/eval_velocity.py` (`mpopi-eval`): `play` env, commanded forward speeds 0.5,
+1.0 and 1.5 m/s (no lateral or yaw command), 8 envs × 250 steps (5 s), eval
+seed 10000; speed averaged after the first 50 steps.
+
+| Planner | Samples × iterations | Horizon | Noise | Knots |
+|---|---|---|---|---|
+| P1 MPPI | 64 × 1 | 20 | 0.3 | none (per step) |
+| P2 MPPI | 64 × 1 | 20 | 0.1 | 4 |
+| P3 MPPI | 128 × 1 | 40 | 0.2 | 4 |
+| P4 MPOPI | 32 × 4 | 40 | 0.2 | 4 |
+
+Baseline: zero actions (default standing pose).
+
+**Pass criterion (MPC is usable as a teacher):** some planner reaches a mean
+forward speed of at least 1.2 m/s at the 1.5 m/s command with at most 0.25
+falls per env in 5 s. If none passes, MPC → PPO on G1 is not pursued with this
+planner; the options are then a better planner (longer horizon, more samples,
+MPC-specific cost) or a different teacher.
diff --git a/docs/mpopi/mpc_injection_ppo.md b/docs/mpopi/mpc_injection_ppo.md
@@ -0,0 +1,114 @@
+# MPC-Injection for PPO on Cartpole swing-up
+
+Date: 2026-09-28. Branch `mpc-stage1`. CPU only.
+
+MPC-Injection (arXiv 2606.26392) mixes MPC transitions into the replay buffer
+of an off-policy learner (SAC / TD3) at a fixed fraction p, with no
+importance correction and no behavior-cloning loss. This study transfers the
+idea to PPO: every PPO update draws MPC samples uniformly from a fixed MPC
+dataset so that they make up a fraction p of the batch, for the whole run.
+
+For PPO this is not an unbiased estimator. Advantages are computed with GAE
+along the MPC's trajectories and the samples are treated as on-policy, so
+the update behaves like advantage-weighted imitation of the MPC. The
+off-policy learners in the paper avoid this because their critic is
+Q(s, a).
+
+## Setup
+
+`mpc_ppo` with `use_in_ppo=True`, `correction=False`, `bc_coef=0`,
+`execution_std=0` (the MPC action itself, as in the paper),
+`inject_fraction=p`, MPOPI teacher K = 8, L = 4, H = 20. MPC data: 8 envs × 16
+steps per iteration for the first 20 iterations, all kept (`max_age=None`,
+20 buffer segments, 2560 samples), then reused until the end of training like
+the paper's offline MPC dataset.
+
+Everything else matches the swing-up recipe study
+([`mpc_ppo_swingup_recipe.md`](mpc_ppo_swingup_recipe.md)): 64 PPO envs, 200
+iterations, eval every 10 iterations over 400 steps, seeds 400–404, one CPU
+thread per process. `PPO` and `MPC-DAgger` are taken from that study (same
+seeds and settings), not rerun.
+
+## Pre-registration (written before running)
+
+Arms: `MPC-Inject` with p = 0.25 (primary; the fraction the paper reports on
+the Go2), p = 0.10 and p = 0.50 (secondary).
+
+- **H1:** `MPC-Inject` (p = 0.25) has a higher AUC than `PPO`, paired by seed.
+  With 5 seeds the smallest two-sided Wilcoxon p is 0.0625, so this can only
+  show a direction (at least 4 of 5 seeds).
+- **H2 (descriptive):** `MPC-Inject` (p = 0.25) compared with `MPC-DAgger`.
+- Also reported: final score, whether the score still drops after the MPC
+  phase (the MPC data stays in every batch here), and wall-clock time.
+
+## Results (CPU)
+
+| Arm | Runs finished | AUC (finished runs) | Paired with `PPO` |
+|---|---|---|---|
+| `PPO` (recipe study) | 5/5 | 0.0273 | – |
+| `MPC-DAgger` (recipe study) | 5/5 | 0.0366 | +0.0093, 4/5 seeds |
+| `MPC-Inject` p = 0.10 | 2/5 | 0.0172 | −0.0101, 0/2 seeds |
+| `MPC-Inject` p = 0.25 | 3/5 | 0.0172 | −0.0147, 0/3 seeds |
+| `MPC-Inject` p = 0.50 | **0/5** | – | – |
+
+Mean normalized score of the finished runs: p = 0.25 reached 0.24 at
+iteration 20, 0.32 at 50 and 0.40 at 199 (`PPO`: 0.21, 0.51, 0.68;
+`MPC-DAgger`: 0.34, 0.84, 0.78). Injection did not speed up the start either.
+
+**10 of 15 runs crashed** with `RuntimeError: normal expects all elements of
+std >= 0.0` between iterations 120 and 180, the same failure as the BC-floor
+runs. The crash rate grows with the injected fraction (3/5, 2/5, 5/5).
+Correction (2026-09-29): an earlier version of this note said the std went
+negative. RSL-RL clamps the std to at least 1e-6 before sampling, so it cannot
+be negative; the error means the std was NaN, i.e. the actor parameters had
+become NaN. A plausible mechanism, not verified: the injected actions are the
+noise-free MPC action, and PPO's surrogate raises `log π(u0|o)` for samples
+with positive advantage; for a Gaussian this shrinks the std (the gradient of
+the log density with respect to σ grows like 1/σ as the mean approaches
+`u0`) toward the 1e-6 floor, where log-probabilities and PPO ratios overflow
+and NaN reaches the parameters. `min_action_std` (a higher std floor, applied
+to every method) was added to test this.
+
+**H1 not supported.** Every finished `MPC-Inject` run had a lower AUC than
+`PPO` on the same seed, and most runs did not finish. Injecting noise-free MPC
+data into PPO as if it were on-policy both hurt learning and destabilized the
+policy's exploration noise on this task.
+
+## Results (GPU, seeds 500–509)
+
+Colab GPU (`cuda:0`), notebook `notebooks/mpc_dagger_vs_inject_colab.ipynb`
+run with the defaults: `PPO` and `MPC-DAgger` were reused from the DAgger GPU
+confirmation (same seeds, byte-identical `curves.csv`), and `MPC-Inject` with
+p = 0.25 was run for the 10 seeds.
+
+| Arm | Runs finished | AUC | Final | Seeds reaching 0.040 |
+|---|---|---|---|---|
+| `PPO` | 10/10 | 0.0230 | 0.0261 | 4/10 |
+| `MPC-DAgger` | 10/10 | **0.0289** | **0.0344** | 9/10 |
+| `MPC-Inject` (p = 0.25) | **8/10** | 0.0206 | 0.0272 | 6/8 |
+
+The two failed runs (seeds 506 and 507) crashed with the NaN-std error,
+at iterations 181 and 141. Mean normalized score by iteration:
+
+| Arm | 0 | 20 | 30 | 40 | 50 | 80 | 120 | 160 | 199 |
+|---|---|---|---|---|---|---|---|---|---|
+| `PPO` | 0.06 | 0.21 | 0.33 | 0.43 | 0.47 | 0.52 | 0.52 | 0.56 | 0.54 |
+| `MPC-DAgger` | 0.10 | 0.38 | 0.53 | 0.60 | 0.75 | 0.75 | 0.57 | 0.51 | 0.72 |
+| `MPC-Inject` | 0.10 | 0.28 | 0.31 | 0.31 | 0.32 | 0.41 | 0.45 | 0.58 | 0.55 |
+
+Paired AUC (Wilcoxon, two-sided, over the 8 seeds where `MPC-Inject`
+finished):
+
+- **H1** `MPC-Inject − PPO`: −0.0024, 3/8 seeds positive, p = 0.84. **Not
+  supported.**
+- **H2** `MPC-Inject − MPC-DAgger`: −0.0094, 1/8 seeds positive, **p = 0.023**.
+  `MPC-Inject` is worse than `MPC-DAgger`. Counting the two crashed runs as
+  losses would only strengthen this.
+- Reference: `MPC-DAgger − PPO` +0.0058, 7/10, p = 0.13 (unchanged).
+
+The GPU run reproduces the CPU picture with fewer crashes: injection gives a
+slightly faster first 20 iterations than `PPO` (0.28 vs 0.21) and then stalls
+around 0.31–0.45 until iteration 120, while `MPC-DAgger` climbs to 0.75. It is
+the first result in this project that reaches p < 0.05. On this task,
+injecting MPC data into PPO as if it were on-policy is a worse use of the same
+MPC teacher than DAgger-style imitation.