"""RTC adapter for openpi's PyTorch pi0.5 (``openpi.models_pytorch.pi0_pytorch.PI0Pytorch``). Bring your own pi0.5 *base* checkpoint (loaded through openpi's ``create_trained_policy``); this module supplies the differentiable velocity field RTC needs and wires it to :func:`rtc.sampler.rtc_guided_sample`. Why this works without inlining anything ---------------------------------------- openpi's sampler ``PI0Pytorch.sample_actions`` is ``@torch.no_grad`` *and* ``torch.compile``d, so RTC cannot route its autograd through it. But the single step it calls, ``PI0Pytorch.denoise_step(state, prefix_pad_masks, past_key_values, x_t, timestep)``, is **not** decorated ``@torch.no_grad`` -- so we can call it directly and gradients flow back to ``x_t``. (Contrast the vlash adapter, whose ``denoise_step`` *is* no-grad and had to be inlined.) We therefore: 1. run the prefix (VLM) once under ``torch.no_grad`` to build the KV cache (``embed_prefix`` -> ``paligemma_with_expert.forward(use_cache=True)``), exactly like ``scripts/sanity_check_ut_mse.py`` in openpi; 2. return ``velocity_fn(x, tau)`` that runs only ``denoise_step`` (the action expert / suffix) against that cache, with grad enabled. Time / velocity convention -------------------------- openpi integrates ``t=1`` noise -> ``t=0`` data and its expert regresses ``v_pi = noise - data`` (see ``PI0Pytorch.sample_actions``: ``dt<0``, loop from ``time=1``). RTC's core (:mod:`rtc.sampler`) integrates the *opposite* way (``tau=0`` noise -> ``tau=1`` data). The two are reconciled by the identity used throughout this package:: t = 1 - tau v_rtc(x, tau) = - v_pi(x, t = 1 - tau) which is what :func:`build_pi0_velocity_field` returns. """ from __future__ import annotations import torch from .masking import PrefixAttentionSchedule from .sampler import VelocityField, rtc_guided_sample def build_observation(policy, obs_dict: dict): """Turn a raw obs dict into a batched, on-device openpi ``Observation``. Replicates the input preamble of ``openpi.policies.policy.Policy.infer`` for the PyTorch path: apply the policy's input transforms (normalization, tokenization, image resize keys), convert to torch tensors, add a batch dim, and build the ``Observation`` dataclass. Args: policy: an openpi ``Policy`` created by ``create_trained_policy`` with a PyTorch model (``policy._model`` is a ``PI0Pytorch``). obs_dict: the same raw dict you would pass to ``policy.infer`` -- e.g. ``{"observation/image": HWC uint8, "observation/state": (D,) float32, "prompt": str, ...}`` (exact keys depend on the checkpoint's data config / input transforms). Returns: An ``openpi.models.model.Observation`` on ``policy._pytorch_device`` with a leading batch dim of 1. """ import jax import numpy as np from openpi.models import model as _model device = policy._pytorch_device inputs = jax.tree.map(lambda x: x, obs_dict) # copy (transforms mutate in place) inputs = policy._input_transform(inputs) inputs = jax.tree.map( lambda x: torch.from_numpy(np.array(x)).to(device)[None, ...], inputs ) return _model.Observation.from_dict(inputs) def build_pi0_velocity_field( policy, observation ) -> tuple[VelocityField, int, int, torch.device, dict]: """Build an RTC-convention velocity field for one observation. Mirrors :func:`rtc.vlash_pi05.build_vlash_pi05_velocity_field`, using openpi's ``PI0Pytorch`` seams: preprocess -> ``embed_prefix`` -> prefill KV cache (no-grad) -> ``velocity_fn(x, tau)`` that calls ``denoise_step`` (grads flow). Args: policy: openpi ``Policy`` (PyTorch) whose ``_model`` is a ``PI0Pytorch``. observation: a batched, on-device ``Observation`` (see :func:`build_observation`). Returns: ``(velocity_fn, horizon, action_dim, device, decode_ctx)``. ``horizon`` is the model action horizon ``H`` and ``action_dim`` the (padded, usually 32) model action dim ``A``. ``decode_ctx`` carries what :func:`decode_actions` needs to unnormalize a model-space chunk back to robot space. """ from openpi.models_pytorch.pi0_pytorch import make_att_2d_masks model = policy._model model.eval() device = policy._pytorch_device horizon = int(model.config.action_horizon) action_dim = int(model.config.action_dim) # ---- Prefill: cache the prefix (VLM) KV once, no grad (it is a constant) ---- with torch.no_grad(): ( images, img_masks, lang_tokens, lang_masks, state, sam2_tokens, ) = model._preprocess_observation(observation, train=False) prefix_embs, prefix_pad_masks, prefix_att_masks = model.embed_prefix( images, img_masks, lang_tokens, lang_masks, sam2_tokens=sam2_tokens ) prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 prefix_att_2d_masks_4d = model._prepare_attention_masks_4d(prefix_att_2d_masks) # openpi requires eager attention for the KV-cache path. model.paligemma_with_expert.paligemma.language_model.config._attn_implementation = "eager" _, past_key_values = model.paligemma_with_expert.forward( attention_mask=prefix_att_2d_masks_4d, position_ids=prefix_position_ids, past_key_values=None, inputs_embeds=[prefix_embs, None], use_cache=True, ) bsize = int(prefix_pad_masks.shape[0]) def velocity_fn(x: torch.Tensor, tau: float) -> torch.Tensor: # RTC tau -> openpi t convention, and v_rtc = -v_pi (see module docstring). t = 1.0 - tau timestep = torch.full((bsize,), t, dtype=torch.float32, device=x.device) v_pi = model.denoise_step(state, prefix_pad_masks, past_key_values, x, timestep) return -v_pi # state is the normalized model-space state; the output transform may need it # (e.g. to convert delta actions to absolute) when decoding. decode_ctx = {"state": state} return velocity_fn, horizon, action_dim, device, decode_ctx def align_previous_chunk(prev_chunk: torch.Tensor, execution_horizon: int) -> torch.Tensor: """Shift the currently-executing chunk into the new chunk's frame. Drops the first ``s`` rows (already executed) and right-pads with zeros. Padding rows get soft-mask weight 0, so their value is irrelevant. Same as :func:`rtc.vlash_pi05.align_previous_chunk`. """ if prev_chunk.ndim != 3: raise ValueError(f"prev_chunk must be (B, H, A), got {tuple(prev_chunk.shape)}") b, h, a = prev_chunk.shape s = execution_horizon remaining = prev_chunk[:, s:, :] pad = torch.zeros( (b, h - remaining.shape[1], a), dtype=prev_chunk.dtype, device=prev_chunk.device ) return torch.cat([remaining, pad], dim=1) def sample_initial_chunk( policy, observation, *, num_steps: int = 5, noise: torch.Tensor | None = None, ) -> tuple[torch.Tensor, dict]: """Plain (unguided) flow sample for the FIRST chunk ``A_init``. RTC's controller needs an initial chunk in the *same model space* as the RTC output (normalized, ``action_dim``-wide), so it must come from the same velocity field -- not from ``policy.infer`` (which is decoded to robot space). This runs :func:`rtc.sampler.sample_flow` with RTC's velocity field. Returns ``(chunk, decode_ctx)`` where ``chunk`` is ``(B, H, A)`` in model space; pass it (and later RTC outputs) through :func:`decode_actions` per step. """ from .sampler import sample_flow velocity_fn, horizon, action_dim, device, decode_ctx = build_pi0_velocity_field( policy, observation ) if noise is None: noise = policy._model.sample_noise((1, horizon, action_dim), device) chunk = sample_flow(velocity_fn, noise=noise, num_steps=num_steps) return chunk, decode_ctx def rtc_sample_actions( policy, observation, prev_chunk: torch.Tensor, inference_delay: int, execution_horizon: int, *, num_steps: int = 5, max_guidance_weight: float = 5.0, prefix_attention_schedule: PrefixAttentionSchedule = "exp", noise: torch.Tensor | None = None, already_aligned: bool = False, ) -> tuple[torch.Tensor, dict]: """RTC guided inpainting for openpi pi0.5. Same surface as :func:`rtc.vlash_pi05.rtc_sample_actions_vlash`. ``observation`` is a batched, on-device ``Observation`` (see :func:`build_observation`); ``prev_chunk`` is the currently-executing chunk in model space ``(B, H, A)``. Returns ``(new_chunk, decode_ctx)``. ``new_chunk`` is in model space; run it through :func:`decode_actions` to get robot-space actions. """ velocity_fn, horizon, action_dim, device, decode_ctx = build_pi0_velocity_field( policy, observation ) y = prev_chunk if already_aligned else align_previous_chunk(prev_chunk, execution_horizon) if y.shape[1] != horizon: raise ValueError(f"prev_chunk horizon {y.shape[1]} != model action_horizon {horizon}") if noise is None: noise = policy._model.sample_noise((y.shape[0], horizon, action_dim), device) noise = noise.to(dtype=y.dtype) new_chunk = rtc_guided_sample( velocity_fn, prev_chunk=y, inference_delay=inference_delay, execution_horizon=execution_horizon, noise=noise, num_steps=num_steps, max_guidance_weight=max_guidance_weight, prefix_attention_schedule=prefix_attention_schedule, ) return new_chunk, decode_ctx def decode_actions(policy, chunk: torch.Tensor, decode_ctx: dict): """Unnormalize a model-space chunk back to robot space. Mirrors the output path of ``Policy.infer``: run the policy's output transforms on ``{"state", "actions"}``. ``chunk`` is ``(B, H, A)`` or ``(H, A)`` in model space; a single batch element is decoded. Returns a numpy array of robot-space actions ``(H, robot_action_dim)``. """ import numpy as np arr = chunk.detach().to(torch.float32).cpu().numpy() if arr.ndim == 3: arr = arr[0] state = decode_ctx.get("state") if state is not None: state = np.asarray(state.detach().to(torch.float32).cpu().numpy()) if state.ndim == 2: state = state[0] outputs = {"state": state, "actions": arr} return policy._output_transform(outputs)["actions"]