120 lines
4.3 KiB
Python
120 lines
4.3 KiB
Python
|
"""Wall-X Real-Time Chunking helpers.
|
||
|
|
|
||
|
|
The guided update follows LeRobot's RTC implementation, adapted to Wall-X's
|
||
|
|
flow convention: Wall-X integrates normalized actions from noise at t=0 to a
|
||
|
|
clean action at t=1 with velocity ``dx/dt``.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from dataclasses import dataclass
|
||
|
|
import math
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class WallXRTCConfig:
|
||
|
|
enabled: bool = True
|
||
|
|
execution_horizon: int = 6
|
||
|
|
max_guidance_weight: float = 10.0
|
||
|
|
prefix_attention_schedule: str = "linear"
|
||
|
|
|
||
|
|
def __post_init__(self):
|
||
|
|
if self.execution_horizon <= 0:
|
||
|
|
raise ValueError("execution_horizon must be positive")
|
||
|
|
if self.max_guidance_weight <= 0:
|
||
|
|
raise ValueError("max_guidance_weight must be positive")
|
||
|
|
if self.prefix_attention_schedule not in {"zeros", "ones", "linear", "exp"}:
|
||
|
|
raise ValueError(
|
||
|
|
"prefix_attention_schedule must be zeros, ones, linear, or exp"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class WallXRTCProcessor:
|
||
|
|
"""Inference-time RTC guidance in normalized Wall-X action space."""
|
||
|
|
|
||
|
|
def __init__(self, config: WallXRTCConfig):
|
||
|
|
self.config = config
|
||
|
|
|
||
|
|
def get_prefix_weights(self, start: int, end: int, total: int) -> torch.Tensor:
|
||
|
|
start = max(0, min(int(start), int(end), int(total)))
|
||
|
|
end = max(start, min(int(end), int(total)))
|
||
|
|
schedule = self.config.prefix_attention_schedule
|
||
|
|
if schedule == "zeros":
|
||
|
|
weights = torch.zeros(total)
|
||
|
|
weights[:start] = 1.0
|
||
|
|
return weights
|
||
|
|
if schedule == "ones":
|
||
|
|
weights = torch.zeros(total)
|
||
|
|
weights[:end] = 1.0
|
||
|
|
return weights
|
||
|
|
|
||
|
|
middle_len = end - start
|
||
|
|
if middle_len:
|
||
|
|
middle = torch.linspace(1.0, 0.0, middle_len + 2)[1:-1]
|
||
|
|
if schedule == "exp":
|
||
|
|
middle = middle * torch.expm1(middle) / (math.e - 1.0)
|
||
|
|
else:
|
||
|
|
middle = torch.empty(0)
|
||
|
|
return torch.cat(
|
||
|
|
[torch.ones(start), middle, torch.zeros(total - end)], dim=0
|
||
|
|
)
|
||
|
|
|
||
|
|
def guide_increasing_flow(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
x_t: torch.Tensor,
|
||
|
|
time: torch.Tensor | float,
|
||
|
|
predict_velocity,
|
||
|
|
prev_chunk_left_over: torch.Tensor | None,
|
||
|
|
inference_delay: int = 0,
|
||
|
|
execution_horizon: int | None = None,
|
||
|
|
) -> torch.Tensor:
|
||
|
|
"""Return RTC-guided velocity for a flow integrated from t=0 to t=1."""
|
||
|
|
if prev_chunk_left_over is None or not self.config.enabled:
|
||
|
|
return predict_velocity(x_t)
|
||
|
|
|
||
|
|
x = x_t.detach().clone().requires_grad_(True)
|
||
|
|
prefix = prev_chunk_left_over.to(device=x.device, dtype=x.dtype)
|
||
|
|
if prefix.ndim == 2:
|
||
|
|
prefix = prefix.unsqueeze(0)
|
||
|
|
if prefix.shape[0] == 1 and x.shape[0] > 1:
|
||
|
|
prefix = prefix.expand(x.shape[0], -1, -1)
|
||
|
|
|
||
|
|
padded = torch.zeros_like(x)
|
||
|
|
steps = min(prefix.shape[1], x.shape[1])
|
||
|
|
dims = min(prefix.shape[2], x.shape[2])
|
||
|
|
padded[:, :steps, :dims] = prefix[:, :steps, :dims]
|
||
|
|
|
||
|
|
horizon = execution_horizon or self.config.execution_horizon
|
||
|
|
horizon = min(int(horizon), steps, x.shape[1])
|
||
|
|
weights = self.get_prefix_weights(inference_delay, horizon, x.shape[1])
|
||
|
|
weights = weights.to(device=x.device, dtype=x.dtype).view(1, -1, 1)
|
||
|
|
|
||
|
|
with torch.enable_grad():
|
||
|
|
velocity = predict_velocity(x)
|
||
|
|
t = torch.as_tensor(time, device=x.device, dtype=x.dtype)
|
||
|
|
remaining = torch.clamp(1.0 - t, min=1e-6)
|
||
|
|
clean_estimate = x + remaining * velocity
|
||
|
|
error = (padded - clean_estimate) * weights
|
||
|
|
correction = torch.autograd.grad(
|
||
|
|
clean_estimate,
|
||
|
|
x,
|
||
|
|
grad_outputs=error.detach(),
|
||
|
|
retain_graph=False,
|
||
|
|
)[0]
|
||
|
|
|
||
|
|
# Same guidance schedule as LeRobot RTC after mapping its 1->0 time
|
||
|
|
# convention to Wall-X's 0->1 convention.
|
||
|
|
t_safe = torch.clamp(t, min=1e-6)
|
||
|
|
inv_r2 = (remaining.square() + t.square()) / remaining.square()
|
||
|
|
weight = (remaining / t_safe) * inv_r2
|
||
|
|
weight = torch.nan_to_num(
|
||
|
|
weight,
|
||
|
|
nan=self.config.max_guidance_weight,
|
||
|
|
posinf=self.config.max_guidance_weight,
|
||
|
|
).clamp(max=self.config.max_guidance_weight)
|
||
|
|
return (velocity + weight * correction).detach()
|
||
|
|
|