token-and-things

Discrete action tokenization for robot manipulation: a MinimalTok action tokenizer, an autoregressive token policy on top of it, and closed-loop evaluation in robomimic (low-dim tasks, absolute rot6d actions). Three subpackages, each run from its own directory:

dir	what	entry point
tokenization/	MinimalTok: register encoder -> FSQ -> decoder	src/training/train.py
policy/	AR token policy over a frozen tokenizer, sim-in-the-loop	arp/train.py
simulation/	receding-horizon rollouts, success rate, rollout video	eval_policy.py

Setup (one command)

Needs uv (curl -LsSf https://astral.sh/uv/install.sh | sh). Python 3.12 is downloaded automatically if missing.

uv sync --extra sim      # training + simulator (robomimic 0.4, robosuite 1.5.1, mujoco), from uv.lock
uv sync                  # training only, no MuJoCo

uv.lock pins every package (torch with CUDA on Linux); uv sync reproduces the exact environment. Datasets are fetched on first use from the ChaoyiPan/mip-dataset HF repo into tokenization/data/robomimic/<task>/ph/ (symlinks into the HF cache). Headless video rendering needs MUJOCO_GL=egl (set automatically by the eval script). Log in to wandb once with uv run wandb login, or pass wandb.mode=offline.

Run

# 1. tokenizer (square, ph, abs rot6d; FSQ noise dropout on by default)
cd tokenization && uv run python src/training/train.py --config configs/minimal_tok_robomimic_square.yaml

# 2. policy; default tokenizer_ckpt is the OAT multi-task no-token-dropout tokenizer converted with
#    tokenization/evals/convert_oat_ckpt.py (OAT .ckpt paths work directly too). Sim eval of latest.pt
#    every sim.eval_freq steps and of best.pt at the end (50 episodes, video of the first 2).
cd policy && uv run python arp/train.py --config configs/arpolicy_robomimic_square.yaml \
    tokenizer_ckpt=../tokenization/experiments/<exp>/checkpoints/best.pt

# 3. standalone sim eval (+ mp4)
cd simulation && uv run python eval_policy.py --ckpt ../policy/experiments/<exp>/checkpoints/best.pt \
    --episodes 50 --video out.mp4

Any key=value on the command line overrides the YAML. Experiments land in <subpackage>/experiments/exp_<stamp>_<name>/ with config.yaml, checkpoints/, metrics.jsonl, metrics.csv, results.json, and sim/ (policy runs), mirrored to wandb project token-and-things.

Pretrained checkpoints

Tokenizers (git-lfs, 23 MB each): MinimalTok trained on robomimic lift + can + square + tool_hang (ph, absolute rot6d 10-dim actions, horizon 16), 8 tokens per chunk, no token dropout, EMA weights, 310 epochs.

file	FSQ levels	vocab	used by
checkpoints/tokenizer_mt4_abs_rot6d_8tok_2kvocab_nodrop.pt	[10, 8, 5, 5]	2000	square policy; default tokenizer_ckpt of the square config
checkpoints/tokenizer_mt4_abs_rot6d_8tok_1kvocab_nodrop.pt	[8, 5, 5, 5]	1000	tool_hang policy

Square reconstruction MSE of the 2000-vocab tokenizer is 1.7e-4 in raw action units. To train a policy on one:

cd policy && uv run python arp/train.py --config configs/arpolicy_robomimic_square.yaml \
    tokenizer_ckpt=../checkpoints/tokenizer_mt4_abs_rot6d_8tok_2kvocab_nodrop.pt

or load it directly with arp.policy.load_tokenizer(path) (tokenize / detokenize on (B, 16, 10) chunks).

Policies (git-lfs, ~45 MB each) are behavior-cloned much-ado-about-noising ar_decoder models (MLP obs encoder + 4-layer causal token decoder over the tokenizer above), EMA weights only, with the obs normalization stats and the tokenizer config embedded, so eval_policy.py needs only the checkpoint and the dataset (for the env metadata):

file	task	steps	success in this simulator
checkpoints/square/policy_bc_ar_decoder_8tok_2kvocab_seed130.pt	square ph	45k	67% over 150 episodes (400-step limit)
checkpoints/tool_hang/policy_bc_ar_decoder_8tok_1kvocab_seed140.pt	tool_hang ph	250k	38% over 50 episodes (700-step limit)

cd simulation && uv run python eval_policy.py \
    --ckpt ../checkpoints/square/policy_bc_ar_decoder_8tok_2kvocab_seed130.pt \
    --hdf5 ../tokenization/data/robomimic/square/ph/low_dim_abs.hdf5 --episodes 50 --video square.mp4
cd simulation && uv run python eval_policy.py \
    --ckpt ../checkpoints/tool_hang/policy_bc_ar_decoder_8tok_1kvocab_seed140.pt \
    --hdf5 ../tokenization/data/robomimic/tool_hang/ph/low_dim_abs.hdf5 --episodes 50 --max-steps 700 \
    --video tool_hang.mp4

Passing --tokenizer-ckpt as well makes the loader check that the file matches the embedded weights.

Checkpoint compatibility

model.arch: oat (default) is weight-compatible with tokenizers trained in the OAT codebase, so those .ckpt files load directly (MinimalTok.from_oat_checkpoint, or point tokenizer_ckpt at the .ckpt). arch: torch is the earlier stock-transformer variant; checkpoints without an arch field resolve to it. simulation/sim/mip_policy.py evaluates much-ado-about-noising ar_decoder policy checkpoints (eval_policy.py --ckpt <mip.pt> --tokenizer-ckpt <oat.ckpt> --hdf5 <dataset>).