Self-Grounded Prediction (policy_training/) — released checkpoint

GitHub

This module implements self-grounded prediction, the mechanism behind HOST (Human-to-robot One-shot Skill Transfer) that resolves execution from a coupled visual demonstration: it localizes the robot's current progress within a visual demonstration, predicts the robot's own future observations conditioned on that localized segment, and derives motor commands from the predicted future. This is implemented as a single autoregressive diffusion model with dual video/action experts (a Mixture-of-Transformers architecture) built on top of the Fast-WAM codebase — see Acknowledgements.

This is a companion checkpoint for the policy_training/ module of the HOST code release; see that repo's policy_training/README.md / README_zh.md for the full pipeline (data_preprocessing/ → alignment/ → coupling/ → policy_training/).

Training regime: robot→robot, not human→robot. This checkpoint's task-video conditioning comes from peer robot demonstration episodes of the same task (via task_paths.json), not from human demonstration video. It exercises the self-grounded prediction architecture on robot-only data, independent of the human-video half of the HOST pipeline (alignment/) that gives HOST its human→robot transfer capability.

Index

Checkpoint Contents

Trained with configs/task/real_joint_2cam_224_1e-4_pac_headwise_ncp_ve.yaml (self_grounded_predictor_joint_cross_attn_ve model config + custom_cross_all data config).

.
├── config.yaml    # minimal inference config (model: + trimmed data.train:)
└── model.pt       # ~15.9GB, bf16 + a small fp32 group

model.pt is a plain torch.save of a nested state_dict dict: {"mot", "proprio_encoder", "progress_encoder", "progress_decoder", "visual_encoder", "step", "torch_dtype"}. No optimizer/scheduler/RNG state — inference-only. sha256: 89382e2a48c1d4a4c5b1791baff49c8cb5ded731a4b25209f0118e6e5016e86f.

This flat layout is simpler than, and not directly discoverable by, policy_training/'s own training-output / eval-harness convention ({checkpoint_dir}/config.yaml + {checkpoint_dir}/checkpoints/weights/step_*.pt, see SelfGroundedPredictorEval._find_checkpoint and Wan22Trainer.save_checkpoint in the source repo) — if you want to drive this checkpoint through scripts/eval_openloop.sh rather than loading it directly (see Evaluation), recreate that structure first:

mkdir -p <checkpoint_dir>/checkpoints/weights
cp config.yaml <checkpoint_dir>/config.yaml
cp model.pt <checkpoint_dir>/checkpoints/weights/step_000000.pt   # any step_<N>.pt name works

config.yaml is not the verbatim training config — it's trimmed to exactly what SelfGroundedPredictorEval / evaluate_openloop.py actually read from data.train (verified by grepping every dcfg.* access in the source repo's eval code), plus the full model: section (every field there is a real constructor kwarg of create_self_grounded_predictor, so none of it is droppable). The full field-for-field comparison against hydra.compose(config_name="train", overrides=["task=real_joint_2cam_224_1e-4_pac_headwise_ncp_ve"]) on the current policy_training/configs/ is documented inline in config.yaml's header comment, along with the two content deviations (both marked NOTE(released checkpoint)):

  1. _target_ paths (model._target_, data.train._target_) — this training run predates the fastwam → self_grounded_prediction rename, so its own saved config still pointed at the old internal module path. Updated to self_grounded_prediction.* here.
  2. model.visual_encoder.{backbone_local_repo,backbone_weights_path,siglip_local_weights_path} — set to null instead of this team's internal cluster paths. visual_encoder.py already falls back to downloading the public DINOv2 (torch.hub, facebookresearch/dinov2) and SigLIP (TIMM) base architectures when these are unset; safe here specifically because this checkpoint's own fine-tuned visual_encoder weights are loaded on top immediately after and overwrite whatever the fallback initialized. (This is the exact gap tracked in the source repo's OPEN_SOURCE_PATH_TODOS.md, row 1 of the load-bearing table — nulling it in this standalone checkpoint's config doesn't require that doc's "re-verify full training run" process, since this is a separate artifact from the repo's own shipped default model config.)

Dropped entirely (not read by the eval harness or by instantiate(cfg.model) + load_checkpoint()): data.val / data.extra_val, all top-level training-loop fields (wandb, output_dir, batch_size, learning_rate, staged_unfreeze, ...), data_path (this team's internal video-paths list — not shippable and not read at eval time anyway), and the data.train fields that only matter for constructing a live training dataset (augmentation flags, drop probabilities, indicator/prompt flags, etc.).

cam_mapping_dir and joint_action_mapping_dir are set to null and are required — the model construction/eval path reads them directly, and joint_action_mapping_dir in particular holds the per-dataset normalization stats (norm_min/norm_delta for actions and joints) needed to convert this model's [-1, 1]-normalized action outputs to physical joint angles / end-effector poses. Without it, eval/real_openloop/ cannot run, and loading the model directly via instantiate() + load_checkpoint() only gets you normalized-space outputs. See Data Preparation for the expected format.

Model Preparation

This step is required before loading this checkpoint (mirrors policy_training/README.md's own "Model Preparation" section — any user of that repo runs this regardless of whose weights they load).

Step 1: set the Wan model directory first (optional, default ./checkpoints):

cd policy_training
mkdir -p checkpoints
export DIFFSYNTH_MODEL_BASE_PATH="$(pwd)/checkpoints"

Step 2: pre-generate the ActionDiT backbone (interpolated from Wan2.2 DiT) — the model construction code reads this file to build the action expert's architecture before this checkpoint's weights are loaded on top and overwrite it:

python scripts/preprocess_action_dit_backbone.py \
  --model-config <this-checkpoint-dir>/config.yaml \
  --output checkpoints/ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt \
  --device cuda \
  --dtype bfloat16

Data Preparation

policy_training/ and alignment/ consume the same on-disk data convention — see data_preprocessing/README.md at the repo root for the full, single-source-of-truth schema (video-paths list, episode directory layout, camera mapping, joint/action normalization). This checkpoint's config.yaml ships cam_mapping_dir / joint_action_mapping_dir as null — fill in your own data prepared in that format before running evaluation. If you want to fine-tune rather than just load this checkpoint, start from the repo's own configs/data/custom_cross_all.yaml (the full, non-trimmed data config with data_path, val, extra_val, and augmentation settings) rather than expanding this minimal inference config back out by hand.

Evaluation

eval/real_openloop/ runs the model open-loop against your own episode data (the same format as Data Preparation), predicting an action chunk from the current observation and comparing it against ground truth — this is not a live-robot control loop, it replays a recorded dataset.

Option A — this repo's own eval harness (after recreating the checkpoints/weights/step_*.pt layout per Checkpoint Contents, filling in cam_mapping_dir / joint_action_mapping_dir, and pointing EVAL_DATA_PATH/DATASET_NAME at your own episodes):

cd policy_training
CHECKPOINT_DIR=<path-to-restructured-checkpoint-dir> bash scripts/eval_openloop.sh

Option B — load directly (works with this checkpoint's flat layout as-is):

import torch
from omegaconf import OmegaConf
from hydra.utils import instantiate

cfg = OmegaConf.load("config.yaml")
model = instantiate(cfg.model, model_dtype=torch.bfloat16, device="cuda")
model.load_checkpoint("model.pt")
model.eval()

load_checkpoint loads with strict=False per sub-module and logs any missing/unexpected keys — on a matching code version, all keys should match exactly. On first load, config.yaml triggers a public download of the Wan2.2-TI2V-5B VAE, the Wan2.1-T2V-1.3B text encoder/tokenizer, and the DINOv2/SigLIP base architectures (all immediately overwritten by this checkpoint's weights where applicable, except the VAE/text encoder, which aren't part of this checkpoint at all).

License

Released under the MIT License (same as policy_training/'s own LICENSE). This is a fine-tune built on top of publicly released base models, each under their own license: Wan2.2-TI2V-5B / Wan2.1-T2V-1.3B (Apache 2.0), DINOv2 (Apache 2.0), SigLIP SO400M (Apache 2.0).

Acknowledgements

This module's codebase is built on top of Fast-WAM: Do World Action Models Need Test-time Future Imagination? (Yuan et al.). We thank the Fast-WAM authors for releasing their codebase.

BibTeX

If you find our work helpful, please consider citing:

@misc{chen2026robotsacquiremanipulationskills,
  title={Robots Acquire Manipulation Skills in Seconds from a Single Human Video},
  author={Guangyan Chen and Meiling Wang and Te Cui and Zichen Zhou and Qi Shao and Shalfun Li and Hang Su and Roy Gan and Hao Wang and Mengyin Fu and Yi Yang and Yufeng Yue},
  year={2026},
  eprint={2607.20033},
  archivePrefix={arXiv},
  primaryClass={cs.RO},
  url={https://arxiv.org/abs/2607.20033},
}
Downloads last month
44
Video Preview
loading

Papers for Guangyan/HOST