| """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) |
| 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) |
|
|
| |
| 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) |
| |
| 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: |
| |
| 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 |
|
|
| |
| |
| 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"] |
|
|
