PerceptronAI/Isaac-0.5

1"""Import-light runtime contracts for real-time chunking (RTC)."""2 3from __future__ import annotations4 5from collections.abc import Callable, Sequence6from dataclasses import dataclass7 8import torch9 10 11def rtc_is_enabled(*, max_delay_steps: int, probability: float | None) -> bool:12    """Return whether training can produce a non-empty RTC prefix."""13    return int(max_delay_steps) > 0 and probability != 0.014 15 16def effective_rtc_max_prefix_steps(17    *,18    max_delay_steps: int,19    probability: float | None,20    action_horizon: int,21) -> int:22    """Return the largest prefix length in the RTC training support."""23    max_delay_steps = int(max_delay_steps)24    action_horizon = int(action_horizon)25    if max_delay_steps < 0:26        raise ValueError(f"max_delay_steps must be >= 0; got {max_delay_steps}.")27    if probability is not None and not 0.0 <= float(probability) <= 1.0:28        raise ValueError(f"probability must be None or in [0, 1]; got {probability}.")29    if action_horizon < 1:30        raise ValueError(f"action_horizon must be >= 1; got {action_horizon}.")31    if not rtc_is_enabled(max_delay_steps=max_delay_steps, probability=probability):32        return 033    return min(max_delay_steps, action_horizon - 1)34 35 36DIT_ACTION_EXPERT_CONFIG_SCHEMA_VERSION = 137DIT_ACTION_EXPERT_CONFIG_V1_FIELDS = (38    "action_dim",39    "action_horizon",40    "num_layers",41    "hidden_dim",42    "num_heads",43    "mlp_ratio",44    "num_inference_steps",45    "timestep_sampling_alpha",46    "timestep_sampling_beta",47    "timestep_sampling_scale",48    "timestep_sampling_offset",49    "train_samples_per_chunk",50    "timestep_embed_dim",51    "rtc_max_delay_steps",52    "rtc_probability",53    "rtc_delay_sampling",54    "rtc_poisson_mean",55    "mask_padded_action_rows",56    "drop_action_dim_overflow",57    "ffn_multiple_of",58    "qk_norm",59    "qk_norm_eps",60    "rope",61    "context_layer_norm",62    "causal_attn",63    "k_batched_cross_attn",64    "k_batched_cross_attn_backend",65)66 67 68@dataclass(frozen=True, eq=False)69class ResolvedRTCActionPrefix:70    """Validated RTC prefix geometry shared by native and HF sampling."""71 72    source_dim: int | None73    output_dim: int74    lengths: torch.Tensor | None75    max_length: int76 77 78@dataclass(eq=False)79class ActionExpertStepModulation:80    """Precomputed per-row AdaLN values for one action-expert step."""81 82    conditioning: torch.Tensor83    block_modulations: Sequence[tuple[torch.Tensor, ...]]84    final_modulation: tuple[torch.Tensor, torch.Tensor]85 86 87def prepare_rtc_conditioning(88    base_timesteps: torch.Tensor,89    prefix_mask: torch.Tensor,90    *,91    time_conditioning: Callable[[torch.Tensor], torch.Tensor],92) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:93    """Build compact suffix/prefix conditioning for checkpointed DiT blocks.94 95    RTC has a sampled suffix timestep per chunk and the fixed clean timestep 196    for prefix rows. Blocks project these two values inside their activation-97    checkpointed forward, then select them per row.98    """99    if base_timesteps.dim() != 1:100        raise ValueError(f"base_timesteps must have shape [B]; got {tuple(base_timesteps.shape)}.")101    if prefix_mask.dim() != 2 or prefix_mask.shape[0] != base_timesteps.shape[0]:102        raise ValueError(103            f"prefix_mask must have shape [B,H] with B={base_timesteps.shape[0]}; got {tuple(prefix_mask.shape)}."104        )105    suffix_conditioning = time_conditioning(base_timesteps)106    # The clean RTC prefix always uses flow time 1, independent of the chunk or107    # flow draw. Compute it once; block projections broadcast this single row.108    prefix_conditioning = time_conditioning(torch.ones(1, device=base_timesteps.device, dtype=base_timesteps.dtype))109    row_mask = prefix_mask.to(device=base_timesteps.device, dtype=torch.bool).unsqueeze(-1)110    return suffix_conditioning, prefix_conditioning, row_mask111 112 113def project_rtc_modulation(114    suffix_conditioning: torch.Tensor,115    prefix_conditioning: torch.Tensor,116    prefix_mask: torch.Tensor,117    *,118    modulation: Callable[[torch.Tensor], torch.Tensor],119    chunks: int,120) -> tuple[torch.Tensor, ...]:121    """Project two compact RTC values and select the result per action row."""122    suffix = modulation(suffix_conditioning)123    prefix = modulation(prefix_conditioning)124    selected = torch.where(prefix_mask, prefix.unsqueeze(1), suffix.unsqueeze(1))125    return tuple(selected.chunk(chunks, dim=-1))126 127 128def validate_rtc_prefix_capability(129    *,130    prefix_length: int,131    max_delay_steps: int,132    probability: float | None,133    action_horizon: int,134    allow_ood: bool = False,135) -> int:136    """Validate a requested prefix against the RTC training support."""137    prefix_length = int(prefix_length)138    effective_max = effective_rtc_max_prefix_steps(139        max_delay_steps=max_delay_steps,140        probability=probability,141        action_horizon=action_horizon,142    )143    if prefix_length < 0:144        raise ValueError(f"prefix_length must be >= 0; got {prefix_length}.")145    if prefix_length > effective_max and not allow_ood:146        raise ValueError(147            f"prefix_length={prefix_length} exceeds the maximum supported RTC prefix {effective_max} "148            f"(configured max_delay_steps={int(max_delay_steps)}, probability={probability}, "149            f"action_horizon={int(action_horizon)}). Pass allow_ood_rtc_prefix=True only for an explicit "150            "out-of-distribution research experiment."151        )152    return effective_max153 154 155def resolve_rtc_action_prefix(156    *,157    action_prefix: torch.Tensor | None,158    prefix_length: int | Sequence[int] | torch.Tensor | None,159    action_dim: int | None,160    batch_size: int,161    action_horizon: int,162    expert_action_dim: int,163    rtc_max_delay_steps: int,164    rtc_probability: float | None,165    device: torch.device,166    allow_ood: bool = False,167) -> ResolvedRTCActionPrefix:168    """Validate and normalize an RTC action-prefix request."""169    action_horizon = int(action_horizon)170    expert_action_dim = int(expert_action_dim)171    output_dim = expert_action_dim if action_dim is None else int(action_dim)172    if output_dim < 1 or output_dim > expert_action_dim:173        raise ValueError(f"action_dim must be in [1, {expert_action_dim}]; got {output_dim}.")174 175    if action_prefix is None:176        if isinstance(prefix_length, torch.Tensor):177            has_prefix_length = bool(prefix_length.detach().to(device="cpu").any().item())178        elif prefix_length is None:179            has_prefix_length = False180        elif isinstance(prefix_length, int):181            has_prefix_length = prefix_length != 0182        else:183            has_prefix_length = any(int(value) != 0 for value in prefix_length)184        if has_prefix_length:185            raise ValueError("action_prefix is required when prefix_length is non-zero.")186        return ResolvedRTCActionPrefix(None, output_dim, None, 0)187 188    if action_prefix.dim() != 3 or action_prefix.shape[0] != batch_size:189        raise ValueError(190            f"action_prefix must have shape [B, P, D] with B={batch_size}; got {tuple(action_prefix.shape)}."191        )192    source_dim = int(action_prefix.shape[2])193    if source_dim < 1 or source_dim > expert_action_dim:194        raise ValueError(f"action_prefix last dim must be in [1, {expert_action_dim}]; got {source_dim}.")195 196    if prefix_length is None:197        lengths = torch.full((batch_size,), action_prefix.shape[1], device=device, dtype=torch.long)198    elif isinstance(prefix_length, torch.Tensor):199        lengths = prefix_length.to(device=device, dtype=torch.long)200        if lengths.dim() == 0:201            lengths = lengths.expand(batch_size)202        elif tuple(lengths.shape) != (batch_size,):203            raise ValueError(f"prefix_length tensor must have shape [] or [{batch_size}]; got {tuple(lengths.shape)}.")204    elif isinstance(prefix_length, int):205        lengths = torch.full((batch_size,), prefix_length, device=device, dtype=torch.long)206    else:207        lengths = torch.as_tensor(list(prefix_length), device=device, dtype=torch.long)208        if tuple(lengths.shape) != (batch_size,):209            raise ValueError(f"prefix_length sequence must have length {batch_size}; got {tuple(lengths.shape)}.")210 211    max_allowed = action_horizon - 1212    # One host transfer for every check below: this runs per /predict and per chunk in the213    # inference-MSE eval, so each extra `.item()` on a device tensor is a blocking sync.214    lengths_cpu = lengths.detach().to("cpu")215    if bool(((lengths_cpu < 0) | (lengths_cpu > max_allowed)).any().item()):216        raise ValueError(217            f"prefix_length values must be in [0, {max_allowed}] so RTC leaves at least one model-generated "218            f"suffix row; got {lengths_cpu.tolist()}."219        )220    max_length = int(lengths_cpu.max().item()) if lengths_cpu.numel() else 0221    validate_rtc_prefix_capability(222        prefix_length=max_length,223        max_delay_steps=rtc_max_delay_steps,224        probability=rtc_probability,225        action_horizon=action_horizon,226        allow_ood=allow_ood,227    )228    if action_prefix.shape[1] < max_length:229        raise ValueError(f"action_prefix has only {action_prefix.shape[1]} rows but max prefix_length={max_length}.")230    if max_length > 0 and action_dim is None and source_dim != expert_action_dim:231        raise ValueError(232            "action_dim is required when action_prefix is narrower than the expert action width; "233            f"got prefix dim {source_dim} and expert width {expert_action_dim}."234        )235    if max_length > 0 and source_dim not in (output_dim, expert_action_dim):236        raise ValueError(237            "action_prefix last dim must match action_dim or the expert action width; "238            f"got prefix dim {source_dim}, action_dim {output_dim}, expert width {expert_action_dim}."239        )240    return ResolvedRTCActionPrefix(source_dim, output_dim, lengths, max_length)241 242 243def materialize_rtc_action_prefix(244    resolved: ResolvedRTCActionPrefix,245    action_prefix: torch.Tensor | None,246    *,247    batch_size: int,248    action_horizon: int,249    expert_action_dim: int,250    device: torch.device,251    dtype: torch.dtype,252    dim_mask: torch.Tensor | None,253) -> tuple[torch.Tensor | None, torch.Tensor | None]:254    """Materialize the fixed prefix values and their row mask."""255    if action_prefix is None or resolved.max_length == 0:256        return None, None257    assert resolved.lengths is not None258    assert resolved.source_dim is not None259    prefix_tensor = torch.zeros(batch_size, action_horizon, expert_action_dim, dtype=dtype, device=device)260    prefix_tensor[:, : resolved.max_length, : resolved.source_dim] = action_prefix[:, : resolved.max_length].to(261        device=device,262        dtype=dtype,263    )264    if dim_mask is not None:265        prefix_tensor = prefix_tensor * dim_mask266    prefix_mask = torch.arange(action_horizon, device=device).view(1, action_horizon, 1) < resolved.lengths.view(267        batch_size,268        1,269        1,270    )271    return prefix_tensor, prefix_mask272 273 274def integrate_rtc_euler(275    initial_state: torch.Tensor,276    *,277    num_steps: int,278    velocity_fn: Callable[[torch.Tensor, float], torch.Tensor],279    prefix_tensor: torch.Tensor | None,280    prefix_mask: torch.Tensor | None,281    dim_mask: torch.Tensor | None,282) -> torch.Tensor:283    """Integrate a flow velocity while pinning an optional RTC prefix."""284 285    def apply_masks(state: torch.Tensor) -> torch.Tensor:286        if dim_mask is not None:287            state = state * dim_mask288        if prefix_mask is not None and prefix_tensor is not None:289            state = torch.where(prefix_mask, prefix_tensor, state)290        return state291 292    state = apply_masks(initial_state)293    dt = 1.0 / num_steps294    for step in range(num_steps):295        flow_time = step / num_steps296        velocity = velocity_fn(state, flow_time)297        if dim_mask is not None:298            velocity = velocity * dim_mask299        if prefix_mask is not None:300            velocity = torch.where(prefix_mask, torch.zeros_like(velocity), velocity)301        state = apply_masks(state + dt * velocity)302    return state303