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.
