arx5-mujoco-wm-delta-sweep-1M

ARX5 MuJoCo World Model — Delta Representation Sweep (1M steps)

Four-way ablation of action representation × state prediction mode for an ARX5 MuJoCo manipulation world model trained with RSL-RWM.

Sweep Configs

All configs share: GPT (6L/12H/384D, ~10.7M params), ensemble_size=2, history_horizon=31, batch_size=64, max_steps=1M, lr=1e-3 (cosine → 3e-5), max_grad_norm=1.0, state_context_noise=linear(0→0.05 over 50k steps).

Winner: cfgC — delta action inputs + absolute state targets.

Key Findings

• Delta action inputs are the single biggest factor (~10x gap: C vs A).
• Delta state targets hurt even with grad clipping (C vs D).
• Grad clipping rescued B/D from gradient explosion but did not make them competitive.

Dataset

• pravsels/arx5-mujoco-trajectories
• 1024 episodes, ~3.5M frames of ARX5 7-DOF manipulation in MuJoCo
• State: 14-dim (7 joint pos + 7 joint vel), Action: 7-dim (PD targets)

Repo Structure

README.md
TRAINING_LOG.md
cfgA/
  assets/normalization_stats.json
  checkpoints/params/params.pt
cfgB/
  assets/normalization_stats.json
  checkpoints/params/params.pt
cfgC/
  assets/normalization_stats.json
  checkpoints/params/params.pt
cfgD/
  assets/normalization_stats.json
  checkpoints/params/params.pt

Checkpoint Contents (params.pt)

Each params.pt contains:

• model_state_dict — model weights (inference + fine-tuning)
• best_step, best_val_loss, train_loss, val_loss — training metrics
• action_representation, state_prediction_mode — config labels
• noise_schedule_cfg — state context noise schedule used during training

Hashes

Computed via cd checkpoints && find params -type f | sort | xargs sha256sum | sha256sum:

W&B

• cfgA
• cfgB
• cfgC
• cfgD

Training

Trained on Isambard AI (GH200, ARM64) with Apptainer. See TRAINING_LOG.md for full details.

Usage

import torch

ckpt = torch.load("cfgC/checkpoints/params/params.pt", map_location="cpu")
state_dict = ckpt["model_state_dict"]
print(f"best_val_loss={ckpt[best_val_loss]:.6f} at step {ckpt[best_step]}")