"""End-to-end RTC demo on openpi pi0.5 — NO ROBOT, NO CAMERAS. Loads a pi0.5 base checkpoint through openpi, seeds an initial action chunk, then runs the Real-Time Chunking controller (paper Algorithm 1) against a synthetic observation stream so you can watch the async inference/execution loop work and verify the real-time constraint ``d <= H - s`` holds. This exercises the real RTC guided-inpainting inference path — it just feeds fake observations instead of driving a robot. Usage ----- # from the repo root, in an env that has openpi + torch + this package: python examples/pi05_realtime.py \ --config pi05_libero \ --checkpoint /path/to/pi05_checkpoint_dir \ --device cuda --steps 60 --hz 30 # feed a real recorded observation instead of the synthetic one: python examples/pi05_realtime.py --config ... --checkpoint ... \ --obs-npz my_obs.npz # keys = your checkpoint's raw obs dict Notes ----- * The raw observation keys (before the policy's input transforms) are checkpoint-specific. The synthetic default below uses libero-style keys (``observation/image``, ``observation/wrist_image``, ``observation/state``, ``prompt``); adapt ``make_synthetic_obs`` or pass ``--obs-npz`` to match yours. * Chunks live in the model's *normalized* action space (width = action_dim, often padded to 32). Decode to robot space with ``rtc.pi0.decode_actions`` — shown at the end for the current chunk. * JAX is only used for openpi's input transforms; we pin it to CPU so it doesn't fight torch for the GPU. """ from __future__ import annotations import argparse import os import time # Keep JAX (used only for openpi's data transforms) off the GPU. os.environ.setdefault("JAX_PLATFORMS", "cpu") import numpy as np def make_synthetic_obs(rng: np.random.Generator, prompt: str) -> dict: """A single fake observation dict (pre-input-transform). ADAPT these keys to your checkpoint's input transform. Defaults are libero-style. Images are HWC uint8; state is a float32 vector. """ img = rng.integers(0, 255, (224, 224, 3), dtype=np.uint8) wrist = rng.integers(0, 255, (224, 224, 3), dtype=np.uint8) state = rng.standard_normal(8).astype(np.float32) return { "observation/image": img, "observation/wrist_image": wrist, "observation/state": state, "prompt": prompt, } def main() -> int: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--config", required=True, help="openpi TrainConfig name, e.g. pi05_libero") ap.add_argument("--checkpoint", required=True, help="checkpoint dir (model.safetensors + assets/)") ap.add_argument("--device", default="cuda", help="torch device (default cuda)") ap.add_argument("--prompt", default="do something", help="task instruction for the synthetic obs") ap.add_argument("--obs-npz", default=None, help="optional .npz of a real raw obs dict (overrides synthetic)") # control loop ap.add_argument("--steps", type=int, default=60, help="how many control steps to simulate") ap.add_argument("--hz", type=float, default=30.0, help="control rate (1/dt)") # RTC hyperparameters (paper symbols in help) ap.add_argument("--smin", type=int, default=6, help="s_min: min execution horizon") ap.add_argument("--initial-delay", type=int, default=4, help="d_init: delay-buffer seed") ap.add_argument("--delay-buffer", type=int, default=10, help="b: recent-delay buffer size") ap.add_argument("--num-steps", type=int, default=5, help="n: denoising steps") ap.add_argument("--beta", type=float, default=5.0, help="guidance-weight clip") ap.add_argument("--schedule", default="exp", choices=["ones", "zeros", "linear", "exp"], help="soft-mask schedule ('zeros' = naive hard-masking baseline)") ap.add_argument("--seed", type=int, default=0) args = ap.parse_args() import torch from openpi.policies import policy_config as _policy_config from openpi.training import config as _config from rtc.controller import RealTimeChunkingController, execution_horizon_for from rtc.pi0 import build_observation, decode_actions, rtc_sample_actions, sample_initial_chunk rng = np.random.default_rng(args.seed) torch.manual_seed(args.seed) # ---- load the pi0.5 base policy (user brings the checkpoint) -------------- print(f"[load] config={args.config} checkpoint={args.checkpoint}") train_config = _config.get_config(args.config) policy = _policy_config.create_trained_policy(train_config, args.checkpoint, pytorch_device=args.device) model = policy._model H = int(model.config.action_horizon) A = int(model.config.action_dim) print(f"[load] pi05={getattr(model.config, 'pi05', '?')} action_horizon H={H} action_dim A={A}") # ---- one observation (synthetic unless --obs-npz) ------------------------ if args.obs_npz: raw = dict(np.load(args.obs_npz, allow_pickle=True)) raw = {k: (v.item() if getattr(v, "ndim", None) == 0 else v) for k, v in raw.items()} obs_dict = raw else: obs_dict = make_synthetic_obs(rng, args.prompt) def build_obs(perturb: bool = False) -> "object": d = dict(obs_dict) if perturb and "observation/state" in d: d["observation/state"] = (np.asarray(d["observation/state"], np.float32) + 0.01 * rng.standard_normal(np.shape(d["observation/state"])).astype(np.float32)) return build_observation(policy, d) # ---- seed the first chunk A_init (plain flow sample, model space) --------- print("[init] sampling initial chunk (first inference is slow: compile/load)...") obs0 = build_obs() t0 = time.monotonic() init_chunk_b, _ = sample_initial_chunk(policy, obs0, num_steps=args.num_steps) print(f"[init] A_init ready in {time.monotonic() - t0:.2f}s shape={tuple(init_chunk_b.shape)}") init_chunk = init_chunk_b[0] # controller stores (H, A); inference_fn re-adds the batch dim # ---- the RTC inference closure the controller calls in its bg thread ------ def inference_fn(obs, prev_chunk, inference_delay, execution_horizon): # prev_chunk is (H, A) torch (model space); rtc_sample_actions wants (B, H, A). prev_b = prev_chunk.unsqueeze(0) new_b, _ = rtc_sample_actions( policy, obs, prev_b, inference_delay=inference_delay, execution_horizon=execution_horizon, num_steps=args.num_steps, max_guidance_weight=args.beta, prefix_attention_schedule=args.schedule, ) return new_b[0] controller = RealTimeChunkingController( inference_fn, horizon=H, initial_chunk=init_chunk, smin=args.smin, initial_delay=args.initial_delay, delay_buffer_size=args.delay_buffer, ) controller.start() print(f"[run] simulating {args.steps} steps @ {args.hz} Hz " f"(smin={args.smin} d_init={args.initial_delay} n={args.num_steps} beta={args.beta} sched={args.schedule})") print(f"[run] real-time constraint: observed delay d must stay <= H - s = {H} - max(d,smin)") dt = 1.0 / args.hz last_delay = None try: for step in range(args.steps): tick = time.monotonic() observation = build_obs(perturb=True) action = controller.get_action(observation) # (A,) torch, model space — returns immediately d = controller.last_observed_delay s = execution_horizon_for(d, args.smin) if d != last_delay: print(f" step {step:3d}: observed delay d={d} s=max(d,smin)={s} " f"headroom H-s={H - s} action[:3]={action[:3].detach().cpu().numpy().round(3)}") last_delay = d sleep = dt - (time.monotonic() - tick) if sleep > 0: time.sleep(sleep) print(f"[run] OK — completed {args.steps} steps, real-time constraint held throughout.") except RuntimeError as e: print(f"[run] REAL-TIME CONSTRAINT VIOLATED: {e}") print(" -> lower --hz, raise --smin, or lower --num-steps.") return 1 finally: controller.stop() # ---- decode the current chunk to robot space (what you'd send to a robot) - with controller._lock: # noqa: SLF001 (demo: read the live chunk safely) cur = controller._cur.clone() robot_actions = decode_actions(policy, cur, {"state": None}) print(f"[decode] current chunk -> robot space {np.asarray(robot_actions).shape}; " f"row0={np.asarray(robot_actions)[0][:4].round(4)}") return 0 if __name__ == "__main__": raise SystemExit(main())