Past2Next β€” LIBERO-10 checkpoints

Two-stage discrete action modelling for robot manipulation, trained on LIBERO-10 with 500 demonstrations (10 tasks Γ— 50).

Code: https://github.com/seanliu7081/past2next_clean

File Stage Metric Size
tokenizer_oattok_so3aug_ep4960_mse0.001.ckpt Stage 1 β€” action tokenizer reconstruction MSE 0.001 @ epoch 4960 93 MB
policy_past2next_self_past_ep0200_sr0.772.ckpt Stage 2 β€” policy success rate 0.772 @ epoch 200 692 MB

Stage 1 β€” action tokenizer (OATTokSO3Aug)

Action-only autoencoder with an FSQ discrete bottleneck. Compresses a 16-step chunk of continuous 7-DoF actions into 8 ordered tokens and reconstructs it; it never sees observations. Trained with SO(3) action-chunk augmentation (p=0.6, max_angle_deg=30, mode=left_noise, rotation dims only).

FSQ levels: [8, 5, 5, 5, 5] β€” a 5-dim latent, i.e. 5000 codes. Checkpoints trained under the earlier [8, 5, 5, 5, 5, 5] (6-dim) setting are not interchangeable with this one; the policy below was trained against these levels.

Stage 2 β€” policy (Past2NextSelfPastPolicy)

Causal transformer conditioned on observations and the policy's own previously generated actions, autoregressively predicting the next action chunk's tokens, decoding them with the frozen Stage-1 tokenizer above and executing receding-horizon.

The "self-past" variant sources past_action from the policy's own generated actions (re-running itself on the previous rollout window) rather than ground-truth dataset actions: self_past_p=1.0 after self_past_warmup_steps=500.

Config: horizon=16, n_action_steps=8, n_obs_steps=2, past_n=7, embed_dim=256, n_layers=4, n_heads=4.

Loading

Each checkpoint embeds its own Hydra config, so no YAML is needed β€” from_checkpoint rebuilds the workspace from the payload's _target_:

from oat.tokenizer.oat.tokenizer import OATTok
from oat.policy.past2next_self_past import Past2NextSelfPastPolicy

tokenizer = OATTok.from_checkpoint("tokenizer_oattok_so3aug_ep4960_mse0.001.ckpt")
policy = Past2NextSelfPastPolicy.from_checkpoint("policy_past2next_self_past_ep0200_sr0.772.ckpt")

You still need the repo code on sys.path and the LIBERO submodule (github.com/Chaoqi-LIU/LIBERO, branch oat) for the environment.

Notes

  • These are training-state checkpoints (torch.save + dill), not safetensors. They carry optimizer and EMA state, which is why the policy file is 692 MB β€” that supports resuming training, but means loading them executes pickled code. Only load checkpoints you trust.
  • The policy checkpoint is from an in-progress run at epoch 200 of 5001; success rate may improve in later revisions of this repo.
Downloads last month

-

Downloads are not tracked for this model. How to track
Video Preview
loading