Datasetpublic

shashwatsaxena136/ogpo-flow-toolhang-umap-ckpts

OGPO (flow policy) tool_hang checkpoints for the UMAP analysis Checkpoints of one online-RL run of the public OGPO code (https://github.com/simchowitzlabpublic/OGPO_public), scripts/ogpo/toolhang.sh defaults: robomimic tool_hang-ph-low_dim, flow-matching policy (10 flow steps, horizon 8), conservative group advantages (OGPO-CA), 10-head Q ensemble, seed 1. Wandb: https://wandb.ai/pluralistic-goal-conditioning/OGPO/runs/kshv94if file phase note params_200000.pkl end of… See the full description on the dataset page: https://huggingface.co/datasets/shashwatsaxena136/ogpo-flow-toolhang-umap-ckpts.

sourceHugging Facemitupdated 10d agoView on Hugging Face

0likes132downloads

Dataset Card

OGPO (flow policy) tool_hang checkpoints for the UMAP analysis

Checkpoints of one online-RL run of the public OGPO code (https://github.com/simchowitzlabpublic/OGPOpublic), `scripts/ogpo/toolhang.sh` defaults: robomimic `toolhang-ph-low_dim`, flow-matching policy (10 flow steps, horizon 8), conservative group advantages (OGPO-CA), 10-head Q ensemble, seed 1. Wandb: https://wandb.ai/pluralistic-goal-conditioning/OGPO/runs/kshv94if

file	phase	note
params_200000.pkl	end of BC	BC clipped at 200k (SDE eval success 0.50)
params_450000.pkl	early RL (250k online steps)
params_700000.pkl	RL 500k
params_950000.pkl	RL 750k
params_1200000.pkl	mid RL (1M online steps)
params_1450000.pkl	RL 1.25M
params_1700000.pkl	RL 1.5M
params_1950000.pkl	RL 1.75M
params_2200000.pkl	end RL (2M online steps)	final ODE success 0.95 (64 eps), SDE 0.88

Each params_<step>.pkl is {'agent': flax.serialization.to_state_dict(agent)} (actor + critic + optimizer states); flags.json is the full run config needed to rebuild the agent (ogpo.agents.ogpo.OGPOAgent.create, then flax.serialization.from_state_dict). eval.csv / eval_sde.csv / online_agent.csv are the run's logged curves.

Used for the flow-vs-AR action-UMAP comparison (demo 133 of tool_hang-ph, keyframes t=9/22/30/59 and every 4th chunk).

Run an evaluation (with the public OGPO repo)

bash

git clone https://github.com/simchowitzlabpublic/OGPO_public && cd OGPO_public
uv sync --extra cuda12 --extra robomimic            # see the repo README; CPU also works (JAX_PLATFORMS=cpu)
export ROBOMIMIC_DATASET_ROOT=/path/to/robomimic    # needs tool_hang/ph/low_dim_v15.hdf5 (env metadata)
hf download shashwatsaxena136/ogpo-flow-toolhang-umap-ckpts --repo-type dataset --local-dir ckpts/toolhang
cp ckpts/toolhang/eval_checkpoint.py scripts/       # (until it is merged into the repo)
MUJOCO_GL=egl uv run python scripts/eval_checkpoint.py \
    --ckpts ckpts/toolhang/params_200000.pkl ckpts/toolhang/params_2200000.pkl --n_eval_envs 32 --eval_episodes 2

eval_checkpoint.py rebuilds the agent from flags.json, restores the pickle, and runs the repo's own parallel evaluator: ode = deterministic flow (the "eval/success" curve), sde = the stochastic flow used during RL (the "eval_sde/success" curve). Reference numbers on 64 episodes: end BC (200k) ODE 0.22 / SDE 0.50; end RL (2.2M) ODE 0.95 / SDE 0.88 (training-time evals; re-evaluations vary by a few points with the seed).