arx5-mujoco-rwm-sweep-500k

ARX5 MuJoCo World Model — Delta Representation Sweep (500k steps)

Ablation of action/state representation for an offline world model trained on ARX5 MuJoCo simulation data. Four configs sweeping action_representation (absolute|delta) × state_prediction_mode (absolute|delta).

Architecture

• Backbone: nanoGPT causal transformer (type: gpt)
• Size: 6 layers, 12 heads, 384 embed dim
• History horizon: 32 (h32)
• Ensemble: 2 members
• State dim: 14 (7 joint pos + 7 joint vel)
• Action dim: 7

Training

• Dataset: train.pt — 1024 episodes, 3.5M frames, ARX5 dual-arm MuJoCo
• Steps: 500,000
• Batch size: 64
• LR: 3e-4 → 3e-5 (cosine)
• Noise: linear context noise (0 → 0.05 over 50k steps)
• Hardware: NVIDIA GH200 (Isambard AI Phase 2)

Results

Config C (delta action input, absolute state targets) achieves 22% lower val loss than the baseline. Delta state targets (B, D) hurt significantly — the state is 14D (7 pos + 7 vel), so predicting s_{t+1} - s_t forces the velocity dims to regress acceleration (a second-order, high-variance quantity). This explains both the poor loss and the ~4x higher grad norms.

Checkpoint Hashes

Computed via cd checkpoints/cfg<X> && find params -type f | sort | xargs sha256sum | sha256sum (verified with double run):

Contents

README.md
TRAINING_LOG.md
checkpoints/
  cfgA/params/params.pt          # absolute/absolute baseline
  cfgA/params/normalization_stats.json
  cfgB/params/params.pt          # absolute/delta
  cfgB/params/normalization_stats.json
  cfgC/params/params.pt          # delta/absolute (winner)
  cfgC/params/normalization_stats.json
  cfgD/params/params.pt          # delta/delta
  cfgD/params/normalization_stats.json

Each params.pt contains: model_state_dict, best_val_loss, best_step, global_step, action_representation, state_prediction_mode, config_label.

Usage

import torch

ckpt = torch.load("checkpoints/cfgC/params/params.pt", map_location="cpu")
model.load_state_dict(ckpt["model_state_dict"])

W&B

• cfgA: https://wandb.ai/pravsels/arx5-mujoco-world-model/runs/4ijawk8v
• cfgB: https://wandb.ai/pravsels/arx5-mujoco-world-model/runs/5u1wo249
• cfgC: https://wandb.ai/pravsels/arx5-mujoco-world-model/runs/492kwx7j
• cfgD: https://wandb.ai/pravsels/arx5-mujoco-world-model/runs/iay4spcz