TemporalRatio: bimanual YAM checkpoints

Video-action policies for the bimanual YAM robot, to be run with the TemporalRatio code. From three camera images, the 14-dimensional robot state and a language instruction, a policy predicts the next 50 actions, and it can also generate the 17-frame video of the future it expects.

Checkpoints

directory policy inference config (configs/infer/abc130k/)
base multi-task policy for 24 bimanual tasks (22 of them in ABC-130k) yam_pp24.yaml
hybrid the multi-task action head paired with a separately fine-tuned video model yam_pp24.yaml
task_finetunes/fold_tshirts base fine-tuned on folding and stacking T-shirts fold_pen_bottle.yaml with policy_config=videoVLA_abc130k_folding_short
task_finetunes/pen_caps base fine-tuned on opening pen caps fold_pen_bottle.yaml with policy_config=videoVLA_abc130k_pen_caps_short
task_finetunes/bottle_caps base fine-tuned on unscrewing bottle caps fold_pen_bottle.yaml with policy_config=videoVLA_abc130k_assembly_short
longhorizon_video/goal_conditioned base with its video model fine-tuned for goal-conditioned, long-horizon video of the three tasks above; the action head is base's fold_pen_bottle.yaml

Each directory (about 8.5 GB) holds model.safetensors (action head), video_model_full.pt (video diffusion transformer), metadata.pt and assets/**/norm_stats.json (normalization statistics).

Setup

Install the code as described in the repository README. On first use it also downloads nvidia/Cosmos-Predict2.5-2B and nvidia/Cosmos-Reason1-7B from Hugging Face.

Usage

Configs accept hf://imishutk0410/TemporalRatio/<directory> as the checkpoint and download that directory on first use. To fetch one ahead of time:

trvam download imishutk0410/TemporalRatio --subfolder task_finetunes/pen_caps --out checkpoints/

Serve a policy to a robot or any other client over websockets (the openpi client protocol):

trvam serve configs/infer/abc130k/yam_pp24.yaml --set checkpoint=hf://imishutk0410/TemporalRatio/base --port 8000

Run it open-loop on recorded episodes (ABC-130k converted as in docs/data.md), comparing the predicted actions with the recorded ones and saving the predicted future videos:

trvam infer configs/infer/abc130k/fold_pen_bottle.yaml \
    --set checkpoint=hf://imishutk0410/TemporalRatio/task_finetunes/pen_caps \
    --set policy_config=videoVLA_abc130k_pen_caps_short --set generate_video=true \
    --dataset abc130k/open_the_pen_caps_val --episodes 0 --max-steps 20 --out results/pen_caps

From Python:

from temporal_ratio.config import loader
from temporal_ratio.inference import ConfiguredPolicy

exp = loader.load_inference("configs/infer/abc130k/yam_pp24.yaml",
                            ["checkpoint=hf://imishutk0410/TemporalRatio/base",
                             "generate_video=true"])   # also loads the video decoder
policy = ConfiguredPolicy(exp)
out = policy.infer({
    "observation/left_image": base_image,             # H x W x 3 uint8, base camera
    "observation/left_wrist_left": left_wrist_image,  # H x W x 3 uint8
    "observation/right_wrist_left": right_wrist_image,
    "observation/state": state,                       # (14,)
    "prompt": instruction,
})
actions = out["actions"]        # (50, 14)
video = out["future_video"]     # 17 predicted frames, the three views side by side
policy.reset()                  # between episodes

All six checkpoints were checked on one NVIDIA H200 with the repository at commit a034e5a: each loads, returns finite actions and generates future videos of the task scene. Video generation used about 25 GB of GPU memory, and loading took up to 94 GB of host memory.

License

The TemporalRatio code is Apache 2.0. The model weights are subject to their own licenses (NVIDIA Open Model License, Gemma Terms of Use).

Downloads last month
384
Video Preview
loading

Model tree for imishutk0410/TemporalRatio

Finetuned
(25)
this model