mmahdavian/isaac_models

UR5e pi0.5: real demonstrations and obstacle-augmented simulation

Two independently fine-tuned pi0.5 models, each initialized from gs://openpi-assets/checkpoints/pi05_base/params and trained for one epoch over the same 171 episodes / 55,641 anchor frames. Training used GPUs 0, 1, 6, and 7. These are OpenPI JAX/Orbax inference bundles, with EMA parameters and normalization assets. They are not Hugging Face Transformers or Safetensors models.

Bundle	Global batch	Samples per GPU	Updates	Warmup updates	Final training loss
bs64	64	16	870	13	0.00427246
bs384	384	96	145	3	0.00621835

Each anchor frame was sampled once in a seeded permutation. Final-batch padding had zero loss and gradient weight. Warmup targets approximately 800 frames, rounded up to whole updates. Both runs used two-device FSDP groups, 64 CPU loading workers, and an action horizon of 15 steps at 15 Hz. Changing batch size changes the optimizer update count even though both runs cover one epoch.

These losses are training metrics, not held-out task success rates. No held-out robot or simulation evaluation was performed for these checkpoints.

Training data

Source	Episodes	Frames
Successful real UR5e demonstrations	43	10,398
hamidreza3252/real2sim_ep_02	43	8,723
hamidreza3252/real2sim_ep_03	42	8,508
Obstacle-augmented captures	43	28,012
Total	171	55,641

The obstacle selection uses exactly the lowest-numbered variation from each source episode, with all 43 selected episodes included. It contains 38 clearing motions and 5 detours, preserving each capture's instruction. Original archive split labels are retained in training.json; this mixture uses all selected sources for training, including the five originally labeled validation. sim_ep_01 is not part of this mixture.

Real/replay rows retain their commanded-state and next-command labels. Augmented rows use measured state at time t and the absolute command held over [t, t + 1/15), with no action shift. State/action normalization is computed over the full selected mixture and episode-clamped future action windows.

Bundle contents

Each bs64/ or bs384/ directory contains:

• params/: the complete portable Orbax/OCDBT inference parameter tree.
• assets/ur5e_real_sim_ep_02_03_augmented_success/norm_stats.json: normalization.
• data_setup.json: the single-arm adapter settings.
• training.json: run configuration, source revisions, selected variants and metrics.
• validation.json: strict model-restoration and normalization checks.
• checksums.json: byte sizes and SHA-256 hashes of the bundle files.

Optimizer/training state is not included in these inference bundles. Preserve all files under params/, including its nested OCDBT manifests and data objects.

Download and restore

Use RoboResearch's isolated deployment training environment with OpenPI pinned at 01177e0242a1c7e8fad2547caa0e987def614cda. The adapter settings correspond to ur5e-vla/training/ur5e_single_arm.py at commit f78302b29f9cb0c0aa263aaf356bfd4ff3877c13.

from pathlib import Path
from huggingface_hub import snapshot_download
import jax.numpy as jnp
from openpi.models.model import restore_params
from openpi.models.pi0_config import Pi0Config
from openpi.shared import normalize

bundle = "bs64"  # or "bs384"
snapshot = snapshot_download(
    repo_id="mmahdavian/isaac_models",
    allow_patterns=[f"{bundle}/**", "README.md", "manifest.json"],
)
checkpoint = Path(snapshot) / bundle
params = restore_params(checkpoint / "params", dtype=jnp.bfloat16)
model = Pi0Config(pi05=True, action_horizon=15).load(
    params, remove_extra_params=False,
)
norm_stats = normalize.load(
    checkpoint / "assets" / "ur5e_real_sim_ep_02_03_augmented_success"
)

This restores the model and its statistics. Robot inference additionally requires the single-arm UR5e observation/action transforms and prompt tokenization:

• Take the UR5e from the first seven slots of the 14-dimensional dataset vectors; the model representation is padded to 32 dimensions.
• Six arm joints are radians; actions use joint deltas relative to the input state.
• Gripper commands are absolute normalized values, 1 - native/255, with 1=open.
• Wall/countertop maps to base_0_rgb, wrist/left to left_wrist_0_rgb, and body-forward/right to right_wrist_0_rgb.
• Apply the supplied normalization and single-arm output transform when serving. Do not substitute Aloha sign or gripper transforms.

Run records

• Batch 64: W&B run er1uqbb8
• Batch 384: W&B run 6ti76ow1

Training extensions live in EAI-RSM/RoboResearch under deploy/; the frozen OpenPI source was not modified.