| """Dual-rate physics model: flow-matching loss and Euler sampler. |
| |
| Flow formulation mirrors tinyvla (modeling_tinyvla.py:229-245 / 266-279): |
| t ~ Beta(alpha, beta) scaled into (0.001, 1.0] |
| x_t = t * noise + (1 - t) * data (noise at t=1, data at t=0) |
| velocity target = noise - data |
| sampling: x ~ N(0,1) at t=1, Euler x <- x - dt * v down to t=0 |
| """ |
|
|
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from .expert import PhysicsFlowExpert |
| from .semantic import SemanticEncoder |
|
|
|
|
| class DualRatePhysicsModel(nn.Module): |
| def __init__(self, cfg, load_pretrained: bool = True, lm: nn.Module | None = None, |
| readout_token_id: int = 0): |
| super().__init__() |
| self.cfg = cfg |
| self.semantic = SemanticEncoder(cfg, load_pretrained=load_pretrained, lm=lm, |
| readout_token_id=readout_token_id) |
| self.expert = PhysicsFlowExpert( |
| state_dim=cfg.state_dim, |
| static_dim=cfg.static_dim, |
| max_objects=cfg.max_objects, |
| d_model=cfg.expert_dim, |
| n_layers=cfg.expert_layers, |
| n_heads=cfg.expert_heads, |
| ) |
| |
| self.global_cond = nn.Linear(3, cfg.expert_dim) |
| self.obj_cond = nn.Linear(cfg.state_dim + cfg.static_dim, cfg.expert_dim) |
|
|
| def trainable_parameters(self): |
| return [p for p in self.parameters() if p.requires_grad] |
|
|
| def compute_latent(self, input_ids, attention_mask): |
| return self.semantic(input_ids, attention_mask) |
|
|
| def build_cond(self, latent, cur_state, static, obj_mask, gravity_dt): |
| """latent (B,R,D), cur_state (B,N,S) normalized, static (B,N,F), |
| obj_mask (B,N) bool, gravity_dt (B,3). Returns cond (B,R+1+N,D), mask.""" |
| b, n = obj_mask.shape |
| g = self.global_cond(gravity_dt)[:, None, :] |
| o = self.obj_cond(torch.cat([cur_state, static], dim=-1)) |
| o = o + self.expert.obj_emb[:, :n] |
| cond = torch.cat([latent, g, o], dim=1) |
| ones = torch.ones((b, latent.shape[1] + 1), dtype=torch.bool, device=obj_mask.device) |
| cond_mask = torch.cat([ones, obj_mask], dim=1) |
| return cond, cond_mask |
|
|
| def flow_loss(self, latent, cur_state, static, obj_mask, target_delta, gravity_dt, |
| generator: torch.Generator | None = None): |
| """All state tensors normalized. Returns scalar masked-MSE flow loss.""" |
| b = target_delta.shape[0] |
| device = target_delta.device |
| alpha = torch.full((b,), self.cfg.flow_beta_alpha, device=device) |
| beta = torch.full((b,), self.cfg.flow_beta_beta, device=device) |
| t = torch.distributions.Beta(alpha, beta).sample() * 0.999 + 0.001 |
| noise = torch.randn(target_delta.shape, device=device, generator=generator, |
| dtype=target_delta.dtype) |
| t_ = t[:, None, None] |
| x_t = t_ * noise + (1 - t_) * target_delta |
| v_target = noise - target_delta |
|
|
| cond, cond_mask = self.build_cond(latent, cur_state, static, obj_mask, gravity_dt) |
| v_pred = self.expert(x_t, t, cur_state, static, obj_mask, cond, cond_mask) |
|
|
| se = (v_pred - v_target) ** 2 |
| mask = obj_mask[..., None].to(se.dtype) |
| return (se * mask).sum() / mask.sum().clamp(min=1) / se.shape[-1] |
|
|
| @torch.no_grad() |
| def sample_delta(self, latent, cur_state, static, obj_mask, gravity_dt, |
| steps: int | None = None, generator: torch.Generator | None = None): |
| """Euler-integrate the flow from noise to a normalized delta prediction.""" |
| steps = steps or self.cfg.denoise_steps |
| b, n = obj_mask.shape |
| device = cur_state.device |
| cond, cond_mask = self.build_cond(latent, cur_state, static, obj_mask, gravity_dt) |
| x = torch.randn((b, n, self.cfg.state_dim), device=device, generator=generator, |
| dtype=cur_state.dtype) |
| dt_f = 1.0 / steps |
| for i in range(steps): |
| t = torch.full((b,), 1.0 - i * dt_f, device=device, dtype=cur_state.dtype) |
| v = self.expert(x, t, cur_state, static, obj_mask, cond, cond_mask) |
| x = x - dt_f * v |
| return x |
|
|
