rtc/pi0.py

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