ludde73865/expo-ft-square-peg-residual-actor-sft

EXPO-FT Square-Peg residual actor + critic (offline SFT)

Proof-of-concept / smoke checkpoint — not a production model.

EXPO-FT's small trainable networks — residual actor (TanhNormal over MLPResNetV2), Q-ensemble critic, batch encoder (CNN), temperature — trained by offline SFT on the robosuite NutAssemblySquare ("square peg") demos.

What this is

• Trained here: residual actor, critic, batch encoder, temperature (the small EXPO-FT nets).
• Frozen, NOT trained, NOT included: the pi0.5 VLA base (mkohegyi/openpi-pi05-simdata22-upright-3ff6e758). Load it separately.
• Base actions: taken straight from the dataset demo action chunks (EXPO_FT_DATASET_BASE_ACTIONS=1) — the VLA is not sampled during the update, so the jitted JAX update stays torch-free.

Training

• Dataset: fabiankerj/simdata22-upright-v2.1 (LeRobot v2.1), 50 episodes / 10,561 transitions.
• 4000 gradient updates, batch 32, utd 8, single A100 (CUDA_VISIBLE_DEVICES=0).
• critic_loss 0.47 → ~0.0003 (converged).
• Driver: rl/expo-ft/train_offline_sft.py. checkpoints/4000/ is the final step; checkpoints/200/ is an earlier 5-episode smoke run.

Format

Orbax checkpoint under checkpoints/4000/ (agent/ + params/ pytrees). Restore with the expo-ft restore_checkpoint after building the agent (see train_offline_sft.py).

Caveats

• Trained on a sim dataset (square-peg) at modest scale — a research artifact, not a production policy. Not evaluated in-sim here.
• The SAC entropy-temperature auto-tuning drifts upward over training (temperature ~4.2 by step 4000, residual_actor_loss inflating with it) because the target entropy isn't adjusted in expo_ft_pi_config. The critic converges fine; tune entropy_scale / adjust_target_entropy before a production run.