| """FACT: Force-Aware Causal Transformer policy. |
| |
| A from-scratch block-causal policy that |
| |
| * drops the 250-point track pathway entirely (no track tokens / encoder / |
| decoder / loss), |
| * replaces future-video (Wan world-model) prediction with *future-state* |
| prediction: the 16 proprioceptive EEF/hand states that follow the 16-step |
| action chunk are denoised jointly with the actions, |
| * uses a two-stage flow schedule: a coarse joint flow |
| (tau 1 -> 0.4) on [action, future state], followed by the force-only |
| tactile refinement (tau 0.4 -> 0), which is re-used unchanged from |
| ``fact.force.ForceOnlyTransformer``. |
| |
| The coarse denoiser is a small DiT-style block-causal transformer. Frozen |
| perception components run offline: the Wan2.2 VAE as a fixed image tokenizer |
| (48-channel latent of the current 3-view frame -> 50 tokens) and the Wan |
| umT5-XXL encoder whose token embeddings enter as clean prefix tokens |
| (``use_text``). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import math |
| from dataclasses import asdict, dataclass, field |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
| from .force import ( |
| ACTION_HORIZON, |
| FORCE_OFFSETS, |
| PADDED_ACTION_DIM, |
| PHYSICAL_ACTION_DIM, |
| ForceOnlyOutput, |
| ForceOnlyTransformer, |
| TwoStageFlowSchedule, |
| euler_flow_step, |
| pad_action_62_to_64, |
| sample_trex_timestep, |
| sinusoidal_tau_embedding, |
| ) |
|
|
|
|
| |
| TYPE_OBS, TYPE_STATE, TYPE_ACTION, TYPE_FSTATE, TYPE_TASK, TYPE_TEXT = range(6) |
|
|
|
|
| @dataclass |
| class FactConfig: |
| |
| dim: int = 1024 |
| num_layers: int = 12 |
| num_heads: int = 16 |
| ffn_dim: int = 4096 |
| dropout: float = 0.0 |
| gradient_checkpointing: bool = False |
| |
| obs_channels: int = 48 |
| obs_patch: int = 2 |
| obs_grid: tuple[int, int] = (10, 20) |
| |
| action_dim: int = PADDED_ACTION_DIM |
| physical_action_dim: int = PHYSICAL_ACTION_DIM |
| action_horizon: int = ACTION_HORIZON |
| state_dim: int = PADDED_ACTION_DIM |
| physical_state_dim: int = PHYSICAL_ACTION_DIM |
| |
| predict_future_state: bool = True |
| future_state_horizon: int = ACTION_HORIZON |
| |
| max_blocks: int = 32 |
| window: int = 4 |
| use_task_token: bool = True |
| time_dim: int = 256 |
| |
| use_text: bool = False |
| text_dim: int = 4096 |
| max_text_tokens: int = 64 |
| |
| use_force: bool = True |
| force_dim: int = 768 |
| force_layers: int = 6 |
| force_heads: int = 12 |
| force_ffn: int = 3072 |
| vq_codebook_size: int = 64 |
| vq_code_dim: int = 256 |
| tactile_dropout: float = 0.1 |
| |
| coarse_steps: int = 6 |
| force_steps: int = 4 |
| tau_split: float = 0.4 |
| |
| action_loss_weight: float = 1.0 |
| state_loss_weight: float = 1.0 |
| force_loss_weight: float = 1.0 |
| vq_loss_weight: float = 1.0 |
| vq_commitment_weight: float = 0.25 |
|
|
| @property |
| def obs_tokens(self) -> int: |
| h, w = self.obs_grid |
| return (h // self.obs_patch) * (w // self.obs_patch) |
|
|
| @property |
| def obs_in_dim(self) -> int: |
| return self.obs_channels * self.obs_patch * self.obs_patch |
|
|
| @property |
| def fstate_tokens(self) -> int: |
| return self.future_state_horizon if self.predict_future_state else 0 |
|
|
| @property |
| def tokens_per_block(self) -> int: |
| return self.obs_tokens + 1 + self.action_horizon + self.fstate_tokens |
|
|
| @property |
| def prefix_tokens(self) -> int: |
| return (1 if self.use_task_token else 0) + (self.max_text_tokens if self.use_text else 0) |
|
|
| def to_json(self) -> str: |
| return json.dumps(asdict(self), indent=1) |
|
|
| @classmethod |
| def from_json(cls, text: str) -> "FactConfig": |
| data = json.loads(text) |
| if "obs_grid" in data: |
| data["obs_grid"] = tuple(data["obs_grid"]) |
| return cls(**data) |
|
|
|
|
| class FactLayout: |
| """Index bookkeeping for one packed [task | text | K x block] sequence.""" |
|
|
| def __init__(self, cfg: FactConfig, num_blocks: int) -> None: |
| self.cfg = cfg |
| self.num_blocks = num_blocks |
| self.task = 1 if cfg.use_task_token else 0 |
| self.text = cfg.max_text_tokens if cfg.use_text else 0 |
| self.prefix = self.task + self.text |
| self.obs = cfg.obs_tokens |
| self.action = cfg.action_horizon |
| self.fstate = cfg.fstate_tokens |
| self.per_block = cfg.tokens_per_block |
| self.length = self.prefix + num_blocks * self.per_block |
|
|
| @property |
| def text_slice(self) -> slice: |
| return slice(self.task, self.task + self.text) |
|
|
| def block_start(self, k: int) -> int: |
| return self.prefix + k * self.per_block |
|
|
| def obs_slice(self, k: int) -> slice: |
| s = self.block_start(k) |
| return slice(s, s + self.obs) |
|
|
| def state_slice(self, k: int) -> slice: |
| s = self.block_start(k) + self.obs |
| return slice(s, s + 1) |
|
|
| def action_slice(self, k: int) -> slice: |
| s = self.block_start(k) + self.obs + 1 |
| return slice(s, s + self.action) |
|
|
| def fstate_slice(self, k: int) -> slice: |
| s = self.block_start(k) + self.obs + 1 + self.action |
| return slice(s, s + self.fstate) |
|
|
| def token_metadata(self, device: torch.device) -> dict[str, torch.Tensor]: |
| slot = torch.zeros(self.length, dtype=torch.long) |
| block = torch.zeros(self.length, dtype=torch.long) |
| ttype = torch.zeros(self.length, dtype=torch.long) |
| if self.task: |
| ttype[0] = TYPE_TASK |
| if self.text: |
| ttype[self.text_slice] = TYPE_TEXT |
| per_slot_type = ( |
| [TYPE_OBS] * self.obs |
| + [TYPE_STATE] |
| + [TYPE_ACTION] * self.action |
| + [TYPE_FSTATE] * self.fstate |
| ) |
| per_slot_type_t = torch.tensor(per_slot_type, dtype=torch.long) |
| for k in range(self.num_blocks): |
| s = self.block_start(k) |
| slot[s : s + self.per_block] = torch.arange(self.per_block) |
| block[s : s + self.per_block] = k |
| ttype[s : s + self.per_block] = per_slot_type_t |
| return { |
| "slot": slot.to(device), |
| "block": block.to(device), |
| "type": ttype.to(device), |
| } |
|
|
| def allow_matrix( |
| self, device: torch.device, text_mask: torch.Tensor | None = None |
| ) -> torch.Tensor: |
| """Block-causal mask: block k reads blocks [k-window, k] and the prefix. |
| |
| Returns [L, L] (no text padding) or [B, 1, L, L] when ``text_mask`` |
| [B, T] marks valid text tokens. |
| """ |
|
|
| meta = self.token_metadata(device) |
| block = meta["block"] |
| is_prefix = torch.zeros(self.length, dtype=torch.bool, device=device) |
| is_prefix[: self.prefix] = True |
| qb = block[:, None] |
| kb = block[None, :] |
| allow = (kb <= qb) & (kb >= qb - self.cfg.window) |
| allow = allow | is_prefix[None, :] |
| |
| allow[is_prefix] = False |
| allow[is_prefix[:, None] & is_prefix[None, :]] = True |
| if text_mask is None or not self.text: |
| return allow |
| valid = torch.ones(text_mask.shape[0], self.length, dtype=torch.bool, device=device) |
| valid[:, self.text_slice] = text_mask.to(device, torch.bool) |
| allow_b = allow[None] & valid[:, None, :] |
| eye = torch.eye(self.length, dtype=torch.bool, device=device) |
| allow_b = allow_b | eye[None] |
| return allow_b[:, None] |
|
|
| def memory_indices(self, k: int) -> torch.Tensor: |
| first = max(0, k - self.cfg.window) |
| pieces = [] |
| if self.prefix: |
| pieces.append(torch.arange(0, self.prefix)) |
| pieces.append( |
| torch.arange(self.block_start(first), self.block_start(k) + self.per_block) |
| ) |
| return torch.cat(pieces) |
|
|
|
|
| class FactBlock(nn.Module): |
| """Pre-norm transformer block with per-token adaLN-zero modulation.""" |
|
|
| def __init__(self, dim: int, num_heads: int, ffn_dim: int, dropout: float) -> None: |
| super().__init__() |
| if dim % num_heads: |
| raise ValueError("dim must be divisible by num_heads") |
| self.dim = dim |
| self.num_heads = num_heads |
| self.head_dim = dim // num_heads |
| self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| self.qkv = nn.Linear(dim, 3 * dim) |
| self.q_norm = nn.RMSNorm(self.head_dim, eps=1e-6) |
| self.k_norm = nn.RMSNorm(self.head_dim, eps=1e-6) |
| self.proj = nn.Linear(dim, dim) |
| self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| self.ffn = nn.Sequential( |
| nn.Linear(dim, ffn_dim), |
| nn.GELU(approximate="tanh"), |
| nn.Linear(ffn_dim, dim), |
| ) |
| self.dropout = nn.Dropout(dropout) |
| self.modulation = nn.Parameter(torch.zeros(1, 1, 6, dim)) |
|
|
| def forward( |
| self, x: torch.Tensor, e: torch.Tensor, allow: torch.Tensor |
| ) -> torch.Tensor: |
| batch, length, _ = x.shape |
| m = (e + self.modulation).unbind(2) |
| shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = m |
| h = self.norm1(x) * (1 + scale_a) + shift_a |
| qkv = self.qkv(h).view(batch, length, 3, self.num_heads, self.head_dim) |
| q, k, v = qkv.unbind(2) |
| q = self.q_norm(q).transpose(1, 2) |
| k = self.k_norm(k).transpose(1, 2) |
| v = v.transpose(1, 2) |
| a = F.scaled_dot_product_attention(q, k, v, attn_mask=allow) |
| a = a.transpose(1, 2).reshape(batch, length, self.dim) |
| x = x + gate_a * self.dropout(self.proj(a)) |
| h = self.norm2(x) * (1 + scale_f) + shift_f |
| x = x + gate_f * self.dropout(self.ffn(h)) |
| return x |
|
|
|
|
| @dataclass |
| class FactCoarseOutput: |
| action_flow: torch.Tensor |
| fstate_flow: torch.Tensor | None |
| hidden: torch.Tensor |
| memory_by_block: torch.Tensor |
| memory_mask_by_block: torch.Tensor |
| layout: FactLayout |
|
|
|
|
| class FactCoarseTransformer(nn.Module): |
| """Block-causal joint denoiser for [action, future state] tokens.""" |
|
|
| def __init__(self, cfg: FactConfig) -> None: |
| super().__init__() |
| self.cfg = cfg |
| dim = cfg.dim |
| self.obs_encoder = nn.Linear(cfg.obs_in_dim, dim) |
| self.state_encoder = nn.Sequential( |
| nn.Linear(cfg.state_dim, dim), nn.SiLU(), nn.Linear(dim, dim) |
| ) |
| self.action_encoder = nn.Linear(cfg.action_dim, dim) |
| if cfg.predict_future_state: |
| self.fstate_encoder = nn.Linear(cfg.state_dim, dim) |
| self.fstate_decoder = nn.Linear(dim, cfg.state_dim) |
| nn.init.zeros_(self.fstate_decoder.weight) |
| nn.init.zeros_(self.fstate_decoder.bias) |
| self.slot_embedding = nn.Embedding(cfg.tokens_per_block, dim) |
| self.block_embedding = nn.Embedding(cfg.max_blocks, dim) |
| self.type_embedding = nn.Embedding(6, dim) |
| if cfg.use_task_token: |
| self.task_token = nn.Parameter(torch.randn(1, 1, dim) * 0.02) |
| if cfg.use_text: |
| self.text_encoder = nn.Sequential( |
| nn.LayerNorm(cfg.text_dim), nn.Linear(cfg.text_dim, dim) |
| ) |
| self.text_position = nn.Embedding(cfg.max_text_tokens, dim) |
| nn.init.normal_(self.text_position.weight, std=0.02) |
| self.time_embedding = nn.Sequential( |
| nn.Linear(cfg.time_dim, dim), nn.SiLU(), nn.Linear(dim, dim) |
| ) |
| self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim)) |
| nn.init.zeros_(self.time_projection[1].weight) |
| nn.init.zeros_(self.time_projection[1].bias) |
| self.blocks = nn.ModuleList( |
| [ |
| FactBlock(dim, cfg.num_heads, cfg.ffn_dim, cfg.dropout) |
| for _ in range(cfg.num_layers) |
| ] |
| ) |
| self.final_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| self.final_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim)) |
| nn.init.zeros_(self.final_modulation[1].weight) |
| nn.init.zeros_(self.final_modulation[1].bias) |
| self.action_decoder = nn.Linear(dim, cfg.action_dim) |
| nn.init.zeros_(self.action_decoder.weight) |
| nn.init.zeros_(self.action_decoder.bias) |
| nn.init.normal_(self.slot_embedding.weight, std=0.02) |
| nn.init.normal_(self.block_embedding.weight, std=0.02) |
| nn.init.normal_(self.type_embedding.weight, std=0.02) |
|
|
| |
| def patchify_obs(self, latents: torch.Tensor) -> torch.Tensor: |
| """[B,K,C,H,W] latents -> [B,K,obs_tokens,C*p*p].""" |
|
|
| b, k, c, h, w = latents.shape |
| p = self.cfg.obs_patch |
| if (h, w) != tuple(self.cfg.obs_grid): |
| raise ValueError(f"obs latent grid {(h, w)} != {self.cfg.obs_grid}") |
| x = latents.view(b, k, c, h // p, p, w // p, p) |
| x = x.permute(0, 1, 3, 5, 2, 4, 6).reshape(b, k, (h // p) * (w // p), c * p * p) |
| return x |
|
|
| def forward( |
| self, |
| *, |
| obs_latents: torch.Tensor, |
| state: torch.Tensor, |
| action: torch.Tensor, |
| tau: torch.Tensor, |
| future_state: torch.Tensor | None = None, |
| fstate_tau: torch.Tensor | None = None, |
| text_tokens: torch.Tensor | None = None, |
| text_mask: torch.Tensor | None = None, |
| ) -> FactCoarseOutput: |
| cfg = self.cfg |
| batch, num_blocks = action.shape[:2] |
| device = action.device |
| dtype = self.action_encoder.weight.dtype |
| layout = FactLayout(cfg, num_blocks) |
| if tau.shape != (batch, num_blocks): |
| raise ValueError("tau must be [B,K]") |
| if fstate_tau is None: |
| fstate_tau = tau |
|
|
| obs = self.obs_encoder(self.patchify_obs(obs_latents.to(device, dtype))) |
| st = self.state_encoder(state.to(device, dtype)).unsqueeze(2) |
| act = self.action_encoder(action.to(device, dtype)) |
| pieces = [obs, st, act] |
| if cfg.predict_future_state: |
| if future_state is None: |
| raise ValueError("future_state is required when predict_future_state") |
| pieces.append(self.fstate_encoder(future_state.to(device, dtype))) |
| x = torch.cat(pieces, dim=2) |
| meta = layout.token_metadata(device) |
| x = x + self.slot_embedding.weight[None, None] |
| x = x + self.block_embedding(torch.arange(num_blocks, device=device))[None, :, None] |
| x = x.reshape(batch, num_blocks * layout.per_block, cfg.dim) |
| prefix: list[torch.Tensor] = [] |
| if layout.task: |
| prefix.append(self.task_token.to(dtype).expand(batch, -1, -1)) |
| if layout.text: |
| if text_tokens is None: |
| raise ValueError("use_text=True requires text_tokens [B,T,text_dim]") |
| tt = text_tokens.to(device, dtype) |
| if tt.shape[1] != layout.text: |
| raise ValueError(f"text_tokens must have {layout.text} tokens, got {tt.shape[1]}") |
| if text_mask is None: |
| text_mask = torch.ones(batch, layout.text, dtype=torch.bool, device=device) |
| tt = self.text_encoder(tt) + self.text_position.weight[None] |
| tt = tt * text_mask.to(device, dtype)[..., None] |
| prefix.append(tt) |
| if prefix: |
| x = torch.cat(prefix + [x], dim=1) |
| x = x + self.type_embedding(meta["type"])[None] |
|
|
| |
| token_tau = torch.zeros(batch, layout.length, device=device, dtype=torch.float32) |
| for k in range(num_blocks): |
| token_tau[:, layout.action_slice(k)] = tau[:, k : k + 1].float() |
| if cfg.predict_future_state: |
| token_tau[:, layout.fstate_slice(k)] = fstate_tau[:, k : k + 1].float() |
| t = self.time_embedding( |
| sinusoidal_tau_embedding(token_tau.flatten(), cfg.time_dim).to(dtype) |
| ).view(batch, layout.length, cfg.dim) |
| e = self.time_projection(t).view(batch, layout.length, 6, cfg.dim) |
| allow = layout.allow_matrix(device, text_mask if layout.text else None) |
|
|
| for block in self.blocks: |
| if cfg.gradient_checkpointing and self.training and torch.is_grad_enabled(): |
| x = torch.utils.checkpoint.checkpoint(block, x, e, allow, use_reentrant=False) |
| else: |
| x = block(x, e, allow) |
| hidden = x |
| shift, scale = self.final_modulation(t).chunk(2, dim=-1) |
| y = self.final_norm(x) * (1 + scale) + shift |
| action_flow = torch.stack( |
| [self.action_decoder(y[:, layout.action_slice(k)]) for k in range(num_blocks)], |
| dim=1, |
| ) |
| fstate_flow = None |
| if cfg.predict_future_state: |
| fstate_flow = torch.stack( |
| [self.fstate_decoder(y[:, layout.fstate_slice(k)]) for k in range(num_blocks)], |
| dim=1, |
| ) |
| |
| token_valid = torch.ones(batch, layout.length, dtype=torch.bool, device=device) |
| if layout.text and text_mask is not None: |
| token_valid[:, layout.text_slice] = text_mask.to(device, torch.bool) |
| indices = [layout.memory_indices(k).to(device) for k in range(num_blocks)] |
| max_len = max(int(i.numel()) for i in indices) |
| memories, masks = [], [] |
| for idx in indices: |
| mem = hidden.index_select(1, idx) |
| mask = token_valid.index_select(1, idx) |
| valid = idx.numel() |
| if valid < max_len: |
| mem = F.pad(mem, (0, 0, 0, max_len - valid)) |
| mask = F.pad(mask, (0, max_len - valid), value=False) |
| memories.append(mem) |
| masks.append(mask) |
| return FactCoarseOutput( |
| action_flow=action_flow, |
| fstate_flow=fstate_flow, |
| hidden=hidden, |
| memory_by_block=torch.stack(memories, dim=1), |
| memory_mask_by_block=torch.stack(masks, dim=1), |
| layout=layout, |
| ) |
|
|
|
|
| class FactPolicy(nn.Module): |
| """Coarse joint [action, future-state] flow + tactile suffix refinement.""" |
|
|
| def __init__(self, cfg: FactConfig) -> None: |
| super().__init__() |
| self.cfg = cfg |
| self.coarse = FactCoarseTransformer(cfg) |
| self.force: ForceOnlyTransformer | None = None |
| if cfg.use_force: |
| self.force = ForceOnlyTransformer( |
| dim=cfg.force_dim, |
| num_heads=cfg.force_heads, |
| num_layers=cfg.force_layers, |
| ffn_dim=cfg.force_ffn, |
| action_horizon=cfg.action_horizon, |
| action_dim=cfg.action_dim, |
| vq_codebook_size=cfg.vq_codebook_size, |
| vq_code_dim=cfg.vq_code_dim, |
| coarse_memory_dim=cfg.dim, |
| ) |
| self.schedule = TwoStageFlowSchedule( |
| coarse_steps=cfg.coarse_steps, |
| force_steps=cfg.force_steps, |
| tau_split=cfg.tau_split, |
| action_horizon=cfg.action_horizon, |
| ) |
|
|
| |
| @property |
| def device(self) -> torch.device: |
| return self.coarse.action_encoder.weight.device |
|
|
| @property |
| def dtype(self) -> torch.dtype: |
| return self.coarse.action_encoder.weight.dtype |
|
|
| def _action_mask(self, like: torch.Tensor) -> torch.Tensor: |
| mask = torch.zeros_like(like, dtype=torch.bool) |
| mask[..., : self.cfg.physical_action_dim] = True |
| return mask |
|
|
| def _state_mask(self, like: torch.Tensor) -> torch.Tensor: |
| mask = torch.zeros_like(like, dtype=torch.bool) |
| mask[..., : self.cfg.physical_state_dim] = True |
| return mask |
|
|
| @staticmethod |
| def _noise( |
| clean: torch.Tensor, tau: torch.Tensor, noise: torch.Tensor | None = None |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| if noise is None: |
| noise = torch.randn_like(clean) |
| t = tau.to(clean.dtype) |
| while t.ndim < clean.ndim: |
| t = t.unsqueeze(-1) |
| noisy = (1 - t) * clean + t * noise |
| target = noise - clean |
| return noisy, target, noise |
|
|
| def _text_kwargs(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: |
| if not self.cfg.use_text: |
| return {} |
| if "text_tokens" not in batch: |
| raise KeyError("use_text=True but the batch has no text_tokens") |
| out = {"text_tokens": batch["text_tokens"].to(self.device, self.dtype)} |
| if "text_mask" in batch: |
| out["text_mask"] = batch["text_mask"].to(self.device, torch.bool) |
| return out |
|
|
| |
| def forward( |
| self, batch: dict[str, torch.Tensor], *, tau: torch.Tensor | None = None |
| ) -> dict[str, torch.Tensor]: |
| """nn.Module entry point (so DDP wrappers sync gradients): training losses.""" |
|
|
| return self.forward_core(batch, tau=tau) |
|
|
| def forward_core( |
| self, batch: dict[str, torch.Tensor], *, tau: torch.Tensor | None = None |
| ) -> dict[str, torch.Tensor]: |
| cfg = self.cfg |
| dev, dt = self.device, self.dtype |
| action_clean = batch["action"].to(dev, dt) |
| state = batch["state"].to(dev, dt) |
| obs = batch["obs_latents"].to(dev, dt) |
| text_kwargs = self._text_kwargs(batch) |
| B, K = action_clean.shape[:2] |
| if tau is None: |
| tau = sample_trex_timestep((B, K), device=dev, dtype=dt) |
| else: |
| tau = tau.to(dev, dt) |
| if tau.ndim == 1: |
| tau = tau[:, None].expand(-1, K) |
| amask = self._action_mask(action_clean).to(dt) |
| action_noise = torch.randn_like(action_clean) * amask |
| action_noisy, action_target, _ = self._noise(action_clean, tau, action_noise) |
|
|
| fstate_clean = fstate_noisy = fstate_target = None |
| smask = None |
| if cfg.predict_future_state: |
| fstate_clean = batch["future_state"].to(dev, dt) |
| smask = self._state_mask(fstate_clean).to(dt) |
| fstate_noisy, fstate_target, _ = self._noise( |
| fstate_clean, tau, torch.randn_like(fstate_clean) * smask |
| ) |
|
|
| out = self.coarse( |
| obs_latents=obs, |
| state=state, |
| action=action_noisy, |
| tau=tau, |
| future_state=fstate_noisy, |
| **text_kwargs, |
| ) |
| action_loss = ((out.action_flow - action_target).square() * amask).sum() / amask.sum().clamp_min(1) |
| zero = action_loss.new_zeros(()) |
| state_loss = zero |
| if cfg.predict_future_state: |
| state_loss = ((out.fstate_flow - fstate_target).square() * smask).sum() / smask.sum().clamp_min(1) |
|
|
| force_loss = vq_loss = commitment_loss = zero |
| perplexity = active_codes = zero |
| if self.force is not None and "force_history" in batch: |
| current_force = batch["current_force"].to(dev, dt) |
| history = batch["force_history"].to(dev, dt) |
| valid = batch.get("force_history_valid") |
| if valid is not None: |
| valid = valid.to(dev, torch.bool) |
| S = current_force.shape[2] |
| force_tau = sample_trex_timestep((B, K, S), device=dev, dtype=dt, upper=cfg.tau_split) |
| force_clean = action_clean[:, :, None].expand(-1, -1, S, -1, -1) |
| noise_base = torch.randn_like(action_clean) * amask |
| force_noisy, force_target, _ = self._noise( |
| force_clean, force_tau, noise_base[:, :, None].expand_as(force_clean) |
| ) |
| split_tau = torch.full_like(tau, cfg.tau_split) |
| action_at_split, _, _ = self._noise(action_clean, split_tau, noise_base) |
| fstate_at_split = None |
| if cfg.predict_future_state: |
| fstate_at_split, _, _ = self._noise( |
| fstate_clean, split_tau, torch.randn_like(fstate_clean) * smask |
| ) |
| with torch.no_grad(): |
| split_out = self.coarse( |
| obs_latents=obs, |
| state=state, |
| action=action_at_split, |
| tau=split_tau, |
| future_state=fstate_at_split, |
| **text_kwargs, |
| ) |
| memory = ( |
| split_out.memory_by_block.detach()[:, :, None] |
| .expand(-1, -1, S, -1, -1) |
| .flatten(0, 2) |
| ) |
| memory_mask = ( |
| split_out.memory_mask_by_block[:, :, None] |
| .expand(-1, -1, S, -1) |
| .flatten(0, 2) |
| ) |
| offsets = torch.as_tensor(self.schedule.offsets[:S], device=dev, dtype=torch.long) |
| offsets = offsets.view(1, 1, S).expand(B, K, -1).flatten() |
| keep = torch.rand(B * K * S, device=dev) >= cfg.tactile_dropout |
| fo = self.force( |
| force_noisy.flatten(0, 2), |
| force_tau.flatten(), |
| current_force.flatten(0, 2), |
| history.flatten(0, 2), |
| coarse_memory=memory, |
| coarse_memory_mask=memory_mask, |
| update_offset=offsets, |
| tactile_keep_mask=keep, |
| tactile_history_valid_mask=None if valid is None else valid.flatten(0, 2), |
| return_details=True, |
| ) |
| assert isinstance(fo, ForceOnlyOutput) |
| pred = fo.flow.view_as(force_target) |
| positions = torch.arange(cfg.action_horizon, device=dev).view(1, 1, 1, -1, 1) |
| suffix = positions >= offsets.view(B, K, S, 1, 1) |
| fmask = amask[:, :, None] * suffix |
| force_loss = ((pred - force_target).square() * fmask).sum() / fmask.sum().clamp_min(1) |
| vq_loss = fo.vq_loss |
| commitment_loss = fo.commitment_loss |
| if fo.tactile_perplexity is not None: |
| perplexity = fo.tactile_perplexity |
| if fo.tactile_active_codes is not None: |
| active_codes = fo.tactile_active_codes.to(dt) |
|
|
| loss = ( |
| cfg.action_loss_weight * action_loss |
| + cfg.state_loss_weight * state_loss |
| + cfg.force_loss_weight * force_loss |
| + cfg.vq_loss_weight * (vq_loss + cfg.vq_commitment_weight * commitment_loss) |
| ) |
| return { |
| "loss": loss, |
| "action_loss": action_loss, |
| "state_loss": state_loss, |
| "force_loss": force_loss, |
| "vq_loss": vq_loss, |
| "commitment_loss": commitment_loss, |
| "tactile_perplexity": perplexity, |
| "tactile_active_codes": active_codes, |
| } |
|
|
| |
| @torch.no_grad() |
| def refine_action_suffix( |
| self, |
| coarse_action_at_split: torch.Tensor, |
| coarse_memory: torch.Tensor, |
| current_force: torch.Tensor, |
| history: torch.Tensor, |
| *, |
| update_offset: int = 0, |
| history_valid: torch.Tensor | None = None, |
| coarse_memory_mask: torch.Tensor | None = None, |
| tactile_keep_mask: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| assert self.force is not None |
| action = pad_action_62_to_64(coarse_action_at_split).to(self.device, self.dtype).clone() |
| action[..., self.cfg.physical_action_dim :] = 0 |
| for step in self.schedule.iter_steps("force"): |
| flow = self.force( |
| action, |
| step.tau, |
| current_force.to(self.device, self.dtype), |
| history.to(self.device), |
| coarse_memory=coarse_memory, |
| coarse_memory_mask=coarse_memory_mask, |
| update_offset=update_offset, |
| tactile_history_valid_mask=history_valid, |
| tactile_keep_mask=tactile_keep_mask, |
| ) |
| proposed = euler_flow_step(action, flow, step.tau, step.tau_next) |
| action[:, update_offset:] = proposed[:, update_offset:] |
| action[..., self.cfg.physical_action_dim :] = 0 |
| return action |
|
|
| @torch.no_grad() |
| def sample( |
| self, |
| batch: dict[str, torch.Tensor], |
| *, |
| seed: int | None = None, |
| mode: str = "cascade", |
| ) -> dict[str, torch.Tensor]: |
| """One-block open-loop sampling. |
| |
| mode: |
| cascade 6 coarse Euler steps to tau=0.4, then 4 tactile steps. |
| coarse_only the coarse action read out at tau=0.4 (the baseline |
| eval's ``coarse_only`` semantics: no refinement). |
| coarse_to_zero coarse model carries the flow to tau=0 (10 steps). |
| """ |
|
|
| cfg = self.cfg |
| dev, dt = self.device, self.dtype |
| state = batch["state"].to(dev, dt) |
| obs = batch["obs_latents"].to(dev, dt) |
| text_kwargs = self._text_kwargs(batch) |
| if state.ndim == 2: |
| state = state[:, None] |
| if obs.ndim == 4: |
| obs = obs[:, None] |
| B = state.shape[0] |
| if state.shape[1] != 1: |
| raise ValueError("online sampling consumes exactly one block") |
| gen = None |
| if seed is not None: |
| gen = torch.Generator(device=dev).manual_seed(seed) |
| action = torch.randn(B, 1, cfg.action_horizon, cfg.action_dim, device=dev, dtype=dt, generator=gen) |
| action[..., cfg.physical_action_dim :] = 0 |
| fstate = None |
| if cfg.predict_future_state: |
| fstate = torch.randn(B, 1, cfg.future_state_horizon, cfg.state_dim, device=dev, dtype=dt, generator=gen) |
| fstate[..., cfg.physical_state_dim :] = 0 |
|
|
| def run(tau_value: float) -> FactCoarseOutput: |
| tau = torch.full((B, 1), tau_value, device=dev, dtype=dt) |
| return self.coarse( |
| obs_latents=obs, state=state, action=action, tau=tau, future_state=fstate, **text_kwargs |
| ) |
|
|
| for step in self.schedule.iter_steps("coarse"): |
| out = run(step.tau) |
| action = euler_flow_step(action, out.action_flow, step.tau, step.tau_next) |
| action[..., cfg.physical_action_dim :] = 0 |
| if fstate is not None: |
| fstate = euler_flow_step(fstate, out.fstate_flow, step.tau, step.tau_next) |
| fstate[..., cfg.physical_state_dim :] = 0 |
| out = run(cfg.tau_split) |
| coarse_action = action[:, 0] |
| memory = out.memory_by_block[:, 0] |
| memory_mask = out.memory_mask_by_block[:, 0] |
| fstate_pred = None |
| if fstate is not None: |
| fstate_pred = euler_flow_step(fstate, out.fstate_flow, cfg.tau_split, 0.0)[:, 0] |
| result: dict[str, torch.Tensor] = { |
| "coarse_action_at_split": coarse_action, |
| "coarse_memory": memory, |
| "coarse_memory_mask": memory_mask, |
| } |
| if fstate_pred is not None: |
| result["future_state_pred"] = fstate_pred |
| if mode == "cascade": |
| if self.force is None: |
| raise RuntimeError("cascade sampling requires the force stage") |
| current_force = batch["current_force"].to(dev, dt) |
| history = batch["force_history"].to(dev) |
| valid = batch.get("force_history_valid") |
| |
| while current_force.ndim > 3: |
| current_force = current_force[:, 0] |
| while history.ndim > 4: |
| history = history[:, 0] |
| if valid is not None: |
| valid = valid.to(dev, torch.bool) |
| while valid.ndim > 2: |
| valid = valid[:, 0] |
| result["action_pred"] = self.refine_action_suffix( |
| coarse_action, memory, current_force, history, |
| update_offset=0, history_valid=valid, coarse_memory_mask=memory_mask, |
| ) |
| elif mode == "coarse_only": |
| result["action_pred"] = coarse_action.clone() |
| elif mode == "coarse_to_zero": |
| for step in self.schedule.iter_steps("force"): |
| out = run(step.tau) |
| action = euler_flow_step(action, out.action_flow, step.tau, step.tau_next) |
| action[..., cfg.physical_action_dim :] = 0 |
| if fstate is not None: |
| fstate = euler_flow_step(fstate, out.fstate_flow, step.tau, step.tau_next) |
| fstate[..., cfg.physical_state_dim :] = 0 |
| result["action_pred"] = action[:, 0] |
| else: |
| raise ValueError(f"unknown sampling mode {mode!r}") |
| return result |
|
|
| |
| def save(self, directory: str | Path, extra: dict[str, Any] | None = None) -> None: |
| directory = Path(directory) |
| directory.mkdir(parents=True, exist_ok=True) |
| (directory / "fact_config.json").write_text(self.cfg.to_json()) |
| torch.save(self.state_dict(), directory / "fact_model.pt") |
| if extra: |
| (directory / "extra.json").write_text(json.dumps(extra, indent=1)) |
|
|
| @classmethod |
| def load(cls, directory: str | Path, map_location: str = "cpu") -> "FactPolicy": |
| directory = Path(directory) |
| cfg = FactConfig.from_json((directory / "fact_config.json").read_text()) |
| model = cls(cfg) |
| state = torch.load(directory / "fact_model.pt", map_location=map_location) |
| |
| |
| key = "coarse.type_embedding.weight" |
| current = model.state_dict()[key] |
| if key in state and state[key].shape[0] < current.shape[0]: |
| padded = current.clone() |
| padded[: state[key].shape[0]] = state[key] |
| state[key] = padded |
| model.load_state_dict(state, strict=True) |
| return model |
|
|
|
|
| def count_parameters(module: nn.Module, trainable_only: bool = False) -> int: |
| return sum( |
| p.numel() for p in module.parameters() if (p.requires_grad or not trainable_only) |
| ) |
|
|
|
|
| __all__ = [ |
| "FactBlock", |
| "FactCoarseOutput", |
| "FactCoarseTransformer", |
| "FactConfig", |
| "FactLayout", |
| "FactPolicy", |
| "count_parameters", |
| ] |
|
|
