#!/usr/bin/env python3 """ rx1_vla_demo.py — a REAL VLA in the RX1 MuJoCo loop, honestly framed. Runs lerobot SmolVLA (or any lerobot-native checkpoint from `zelpi hub install`) closed-loop in the RX1 cube scene: MuJoCo cameras (ego + overview, real pixels) ─┐ right-arm joint state (6-dim) ─┼─▶ SmolVLA forward pass ─▶ 6-dof action task string, e.g. "stack the cubes" ─┘ │ naive retarget → RX1 right arm ▼ RX1Env.step() → MuJoCo Every part of the inference chain is genuine: real checkpoint weights, real camera renders from the live scene, real language conditioning, one real forward pass per action chunk. What it is NOT: competent. smolvla_base is trained on SO-100/SO-101 tabletop data; the RX1 is an unseen embodiment with different kinematics, cameras and action semantics, so zero-shot it will gesture, not stack — the final report measures and says so. Getting an actual VLA-driven stack needs fine-tuning on RX1 demonstrations (see README_v2.md / Finetune_SmolVLA_notebook.ipynb in the checkpoint). MUJOCO_GL=egl python scripts/rx1_vla_demo.py \ --checkpoint ~/.zelpi/models/smolvla --task "stack the cubes" """ from __future__ import annotations import argparse import sys import time from pathlib import Path import numpy as np import mujoco import torch sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from env.rx1_env import RX1Env # noqa: E402 from models.rx1_ik import (RX1IKSolver, CUBE_SIZE, # noqa: E402 RIGHT_ARM_ACT) ROBOT_CFG = { "simulation": {"physics_timestep": 0.002, "n_substeps": 4}, "vision": {"enabled": False}, # we render frames ourselves at the VLA's size "actuator": {"max_joint_speed": 3.0, "control_clip": 2.5}, } IMG = 256 # smolvla_base declares 3×256×256 camera features STATE_JOINTS = RIGHT_ARM_ACT[:6] # 6-dim state slot ← first 6 right-arm joints def load_policy(checkpoint: str): from lerobot.configs.policies import PreTrainedConfig from lerobot.policies.factory import get_policy_class, make_pre_post_processors config = PreTrainedConfig.from_pretrained(checkpoint) policy_cls = get_policy_class(config.type) print(f"[rx1-vla] checkpoint type: {config.type} -> {policy_cls.__name__}") policy = policy_cls.from_pretrained(checkpoint) policy.eval() device = "cuda" if torch.cuda.is_available() else "cpu" policy.to(device) device_override = {"device_processor": {"device": device}} pre, post = make_pre_post_processors( policy.config, pretrained_path=checkpoint, preprocessor_overrides=device_override, postprocessor_overrides=device_override) print(f"[rx1-vla] loaded on {device}; action dim " f"{policy.config.output_features['action'].shape}") return policy, pre, post class SceneCams: """Real RGB renders of the live scene at the VLA's declared resolution.""" def __init__(self, env): self._r = mujoco.Renderer(env.model, IMG, IMG) self._env = env def frame(self, camera: str) -> torch.Tensor: self._r.update_scene(self._env.data, camera=camera) rgb = self._r.render() # (H, W, 3) uint8 return torch.from_numpy(rgb.copy()).permute(2, 0, 1).float() / 255.0 def save(self, camera: str, path: str): self._r.update_scene(self._env.data, camera=camera) try: import PIL.Image as Image Image.fromarray(self._r.render()).save(path) except ImportError: pass def build_obs(env, cams, task: str, config, weld_eq: int | None): """Observation dict shaped by the CHECKPOINT's own declared features, so the same script runs zero-shot smolvla_base (3 generic cams, 6-dim state) and RX1 fine-tunes from collect_rx1_lerobot.py (ego+overview, 8-dim).""" obs = {} for key, feat in config.input_features.items(): if key.startswith("observation.images"): cam = next((c for c in ("top", "overview") if c in key), "ego") obs[key] = cams.frame(cam) elif key == "observation.state": dim = feat.shape[0] if dim >= 8 and weld_eq is not None: state = np.concatenate([ env.data.qpos[RIGHT_ARM_ACT], [float(env.data.eq_active[weld_eq])], ]) else: state = env.data.qpos[STATE_JOINTS][:dim] obs[key] = torch.tensor(state, dtype=torch.float32) obs["task"] = task return obs def target_cube_from_task(task: str, env) -> str | None: """The grasp bit convention from collect_rx1_lerobot.py: it drives the weld of the cube being stacked — the FIRST colour named in the task.""" for word in task.lower().replace(",", " ").split(): if f"{word}_cube" in env._weld_bit: return f"{word}_cube" return None def main(): ap = argparse.ArgumentParser(description="Zero-shot SmolVLA in the RX1 MuJoCo loop") ap.add_argument("--checkpoint", required=True) ap.add_argument("--task", default="stack the cubes") ap.add_argument("--steps", type=int, default=400, help="control steps to run") ap.add_argument("--frames-out", default="", help="dir to save first/last camera frames") ap.add_argument("--spawn-seed", type=int, default=-1, help="if >=0, teleport the task cubes to random spots inside " "the training distribution (see collect_rx1_lerobot.py)") ap.add_argument("--n-action-steps", type=int, default=0, help="if >0, re-plan after this many actions instead of the " "checkpoint's default (smaller = more reactive)") ap.add_argument("--render", action="store_true", help="open the MuJoCo viewer and pace to real time") ap.add_argument("--torch-seed", type=int, default=-1, help="if >=0, seed torch so the flow-matching action " "sampling (and thus the whole rollout) is reproducible") ap.add_argument("--hold-back", action="store_true", help="hold torso yaw/roll at start pose (use ONLY for " "checkpoints trained on data collected with the same " "hold — older checkpoints learned to compensate for " "the drifting torso and get worse with it)") args = ap.parse_args() if args.torch_seed >= 0: torch.manual_seed(args.torch_seed) policy, pre, post = load_policy(args.checkpoint) if args.n_action_steps > 0: policy.config.n_action_steps = args.n_action_steps print(f"[rx1-vla] re-planning every {args.n_action_steps} actions") env = RX1Env(ROBOT_CFG) ik = RX1IKSolver(env.model) mujoco.mj_forward(env.model, env.data) env.data.ctrl[:env._nu] = env.data.qpos[:env._nu] cams = SceneCams(env) nu = env._nu action_dim = policy.config.output_features["action"].shape[0] finetuned = action_dim >= 8 # collect_rx1_lerobot.py convention tgt_cube = target_cube_from_task(args.task, env) if finetuned else None base_cube = next((f"{w}_cube" for w in reversed(args.task.lower().split()) if f"{w}_cube" in env._weld_bit and f"{w}_cube" != tgt_cube), None) if finetuned else None weld_bit = env._weld_bit[tgt_cube] if tgt_cube else None weld_eq = dict(env._cube_welds)[tgt_cube] if tgt_cube else None act_joints = RIGHT_ARM_ACT if finetuned else STATE_JOINTS lo = env.model.actuator_ctrlrange[act_joints, 0] hi = env.model.actuator_ctrlrange[act_joints, 1] if finetuned: print(f"[rx1-vla] fine-tuned RX1 checkpoint: 7 arm joints + grasp bit " f"→ {tgt_cube or 'no cube named in task!'}") if args.spawn_seed >= 0 and tgt_cube is not None: # place base+top cubes at random spots inside the training workspace # (matches collect_rx1_lerobot.py's SPAWN ranges), distractors parked from scripts.collect_rx1_lerobot import (_set_cube, _sample_spots, TABLE_Z, COLORS) rng = np.random.default_rng(args.spawn_seed) spots = _sample_spots(rng, 0) _set_cube(env, base_cube, [*spots[0], TABLE_Z]) _set_cube(env, tgt_cube, [*spots[1], TABLE_Z]) for i, c in enumerate(COLORS): if f"{c}_cube" not in (base_cube, tgt_cube): _set_cube(env, f"{c}_cube", [1.0 + 0.1 * i, 1.0, 0.1]) mujoco.mj_forward(env.model, env.data) env.data.ctrl[:nu] = env.data.qpos[:nu] print(f"[rx1-vla] spawned {tgt_cube} at {np.round(spots[1], 3).tolist()}, " f"{base_cube} at {np.round(spots[0], 3).tolist()}") start_pos = {k: v["pos"].copy() for k, v in env.perceive_objects().items()} if args.frames_out: Path(args.frames_out).mkdir(parents=True, exist_ok=True) cams.save("ego", f"{args.frames_out}/ego_first.png") cams.save("overview", f"{args.frames_out}/overview_first.png") # The RX1 dataset is recorded at 31 Hz — every 4th control step (see # collect_rx1_lerobot.py RECORD_EVERY). Querying the policy at the raw # 125 Hz control rate would replay its trajectories 4× too fast, so hold # each action for the same number of steps it represented in training. action_repeat = 4 if finetuned else 1 print(f"[rx1-vla] task: '{args.task}' — running {args.steps} closed-loop steps " + ("(fine-tuned on RX1 demos)" if finetuned else "(zero-shot: this checkpoint has never seen the RX1)")) viewer = None if args.render: from mujoco import viewer as mj_viewer viewer = mj_viewer.launch_passive(env.model, env.data) # Torso regime must MATCH the checkpoint's training data. Fully-compliant # (the original demo regime) lets the arm's weight drag the torso sideways; # --hold-back pins torso yaw/roll at the start pose (pitch stays compliant — # the arm cannot reach the table from a bolt-upright back, models/expert.py) # and is only correct for checkpoints trained on hold-back data. TORSO_YAW, TORSO_ROLL = 0, 2 torso_hold = env.data.qpos[[TORSO_YAW, TORSO_ROLL]].copy() def compliant_base(): base = env.data.qpos[:nu].copy() if args.hold_back: base[TORSO_YAW], base[TORSO_ROLL] = torso_hold return base forwards, t_inf = 0, 0.0 full = np.zeros(env.N_ACTUATORS, dtype=np.float32) held = None # last policy targets, held between queries weld_events = [] # (step, on/off, cube pos) transitions weld_prev = False for step in range(args.steps): if step % action_repeat: full[:nu] = compliant_base() if finetuned else env.data.ctrl[:nu] if held is not None: full[act_joints] = held # weld bit persists in `full` env.step(full) if viewer is not None: viewer.sync() time.sleep(0.008) continue obs = build_obs(env, cams, args.task, policy.config, weld_eq) t0 = time.monotonic() batch = pre(obs) with torch.no_grad(): was_empty = len(getattr(policy, "_queues", {}).get("action", [])) == 0 \ if hasattr(policy, "_queues") else True action = policy.select_action(batch) action = post(action) dt = time.monotonic() - t0 if dt > 0.05 or was_empty: forwards += 1 t_inf += dt a = np.asarray(action.squeeze(0).cpu().numpy(), dtype=np.float64) full = np.zeros(env.N_ACTUATORS, dtype=np.float32) if finetuned: # Demo regime (compliant non-arm joints) with the yaw/roll hold # above; 7 right-arm targets from the policy, grasp bit → the # task cube's weld. full[:nu] = compliant_base() held = np.clip(a[:7], lo, hi) full[act_joints] = held if weld_bit is not None: weld_now = a[7] > 0.5 full[weld_bit] = 1.0 if weld_now else 0.0 if weld_now != weld_prev: weld_events.append( (step, "GRASP" if weld_now else "RELEASE", env.perceive_objects()[tgt_cube]["pos"].copy())) weld_prev = weld_now else: # Naive zero-shot retarget: 6 outputs as position targets for the # first 6 right-arm joints. An SO-100 policy on a 40-DoF humanoid # can't be mapped faithfully — the fine-tuned path above replaces it. full[:nu] = env.data.ctrl[:nu] held = np.clip(a[:6], lo, hi) full[act_joints] = held env.step(full) if viewer is not None: viewer.sync() if step % 100 == 0: print(f" step {step:4d} action={np.round(a, 3).tolist()}") if args.frames_out: cams.save("ego", f"{args.frames_out}/ego_last.png") cams.save("overview", f"{args.frames_out}/overview_last.png") # ── honest report ──────────────────────────────────────────────────────── end_pos = {k: v["pos"] for k, v in env.perceive_objects().items()} print(f"\n{'=' * 60}") print(f"[rx1-vla] {forwards} real forward passes, " f"avg {t_inf / max(forwards, 1):.2f}s each") moved = False for name in sorted(start_pos): d = float(np.linalg.norm(end_pos[name] - start_pos[name])) moved |= d > 0.005 print(f" {name:<12} moved {d * 100:5.1f} cm") if finetuned and tgt_cube and base_cube: goal = end_pos[base_cube] + np.array([0.0, 0.0, CUBE_SIZE]) d = float(np.linalg.norm(end_pos[tgt_cube] - goal)) print(f" {tgt_cube} final distance to stack goal on {base_cube}: {d * 100:.1f} cm") for s, kind, pos in weld_events: print(f" step {s:4d} {kind:<7} {tgt_cube} at {np.round(pos, 3).tolist()}") cubes = [n for n in end_pos if n != "box"] stacked = any( np.linalg.norm(end_pos[a][:2] - end_pos[b][:2]) < CUBE_SIZE / 2 and abs((end_pos[a][2] - end_pos[b][2]) - CUBE_SIZE) < CUBE_SIZE / 2 for a in cubes for b in cubes if a != b) if finetuned: tail = ("the fine-tuned VLA stacked the cube." if stacked else "fine-tuned but placement missed this episode.") else: tail = ("unexpected for zero-shot!" if stacked else "as expected zero-shot; fine-tune on RX1 demos for competence.") print(f"\nVERDICT: inference chain real (pixels → language → action → MuJoCo); " f"objects {'moved' if moved else 'did not move'}, " f"stack {'FORMED' if stacked else 'NOT formed'} — {tail}") if viewer is not None: print("(viewer stays open 20 s — Ctrl-C to exit sooner)") t0 = time.time() while viewer.is_running() and time.time() - t0 < 20: viewer.sync() time.sleep(0.05) viewer.close() return 0 if __name__ == "__main__": sys.exit(main())