Summary

Implements two TODO items from the README:

• Mixture of Experts in SwiGLU FFN — New MoESwiGLUFFN class wraps N smaller SwiGLU experts with top-k routing and load-balancing auxiliary loss. Enabled via use_moe: true in config (dynamics model only). Default: 4 experts, top-2 routing, 0.01 aux loss coefficient.
• Muon optimizer — Newton-Schulz orthogonalized updates for 2D weight matrices, with AdamW handling biases/norms/embeddings. Standalone implementation with no external dependencies. Enabled via optimizer: "muon" in config.

Both features are opt-in with zero regression to default behavior. All existing configs work without modification.

Inference Visualization

MoE + Muon dynamics on PONG (2000 steps, 4 experts top-2, autoregressive generation without actions):

View inference result (GT vs Predicted frames) on W&B — file: inference_results_gt_vs_pred_no_actions_20260413_091147.png

• 12 frames generated autoregressively (2 context + 10 predicted)
• MSE (GT vs Pred): 0.004465
• All predicted frames maintain coherent paddle/ball positions

Also fixed two bugs found during inference testing:

• run_inference.py: os.path.isfile → os.path.exists to support directory-format checkpoints
• utils.py: model_sd.get('model', {}) → model_sd for correct conditioning_dim inference from state dict keys

W&B Training Runs

Short runs (500 steps, RTX 4090, PicoDoom)

W&B Project Dashboard

Longer runs — AdamW vs Muon convergence comparison (5000 steps, RTX 4090, PicoDoom)

Full dynamics-stage comparison with shared VT/LA prerequisites (2000 steps each):

*Run 4 completed ~4990/5000 steps before spot preemption; W&B has near-complete data.

Key findings:

• Both Muon and MoE independently improve over the AdamW baseline
• MoE gives slightly larger gains than Muon alone at 5000 steps
• Per-expert token utilization logged to W&B (moe/block{i}_expert{j}_frac) — router distributes tokens across all 4 experts with no expert collapse

All runs visible at: W&B Project Dashboard

Changes

New files

• models/muon.py — Standalone Muon optimizer (Newton-Schulz orthogonalization + momentum). DDP-compatible; FSDP limitation documented.
• utils/optimizer_utils.py — Shared create_optimizer() for all training stages. Handles AdamW (single optimizer) and Muon (split: Muon for 2D weights, AdamW for biases/norms/embeddings).

Modified files

• models/st_transformer.py — SwiGLUExpert, MoESwiGLUFFN classes with top-k routing + load-balancing aux loss. Per-expert utilization tracking. STTransformerBlock and STTransformer accept MoE kwargs.
• models/dynamics.py — Passes MoE config through to STTransformer.
• utils/config.py — MoE fields (use_moe, num_experts, top_k_experts, moe_aux_loss_coeff) and optimizer fields (optimizer, muon_momentum, muon_backend_steps) added to all config dataclasses.
• configs/training.yaml — Defaults for new config fields.
• scripts/train_dynamics.py — Multi-optimizer loop, MoE aux loss integration, per-expert utilization W&B logging, multi-optimizer checkpoint save.
• scripts/train_video_tokenizer.py — Multi-optimizer support.
• scripts/train_latent_actions.py — Multi-optimizer support.
• scripts/run_inference.py — Fix directory-format checkpoint detection.
• utils/utils.py — Fix conditioning_dim inference from state dict, add MoE config to checkpoint loader.

Usage

# Enable MoE (dynamics only)
use_moe: true
num_experts: 4
top_k_experts: 2
moe_aux_loss_coeff: 0.01

# Enable Muon (all stages)
optimizer: "muon"
muon_momentum: 0.95
muon_backend_steps: 5

Known limitations

• Muon only supports DDP, not FSDP (Newton-Schulz requires full weight matrices, not shards)
• Muon LR/momentum may need per-architecture tuning vs AdamW defaults
• MoE checkpoint only saves primary optimizer via save_training_state; secondary optimizer state saved in separate all_optimizers.pt file

Test plan

• Unit tests: MoE gradient flow, device safety, param coverage, eval mode
• Code review: bf16 grad corruption fix, device-safe aux loss, checkpoint completeness
• E2E training: 4 runs on RTX 4090 with W&B logging (see table above)
• Inference visualization: MoE+Muon on PONG (2000 steps, autoregressive)
• Longer training runs to compare AdamW vs Muon convergence curves (5000 steps, 4 configs)
• MoE expert utilization analysis (per-expert token fractions logged to W&B)

🤖 Generated with Claude Code