rtc-pi05: standalone Real-Time Chunking for openpi pi0.5

bdca518 verified about 1 month ago

8.85 kB

"""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())

	"""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())