You can deploy this model and start earning money today!

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

Name: arx5-mujoco-wm-delta-sweep-1M
Availability: InStock

Views

No views yet

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).

Config	Action Input	State Target	best_val_loss	best_step	params_hash
cfgA	absolute	absolute	0.6345	940k	`b7f10671...`
cfgB	absolute	delta	0.3303	995k	`0f111485...`
cfgC	delta	absolute	0.0585	900k	`d81f0a3e...`
cfgD	delta	delta	0.4970	980k	`a50ef808...`

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:

Config	SHA-256
cfgA	`b7f1067167440f189152294375febaedcfcc92fa6f26f20097843d9129a1fef3`
cfgB	`0f111485d9ba7fa9babbadaf99df7af4315414a78d6f73075c45b1a1c9bb6b12`
cfgC	`d81f0a3e7612b27a3e29045b6180b245c737944d5a7a63320dfa54911eb7190b`
cfgD	`a50ef80e6823a18063ac2d78b638febc8e0f721dfcfc6b674755ef270f1d41e0`

W&B

Training

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

Usage

python

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]}")

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