| """Stateful dual-rate rollout predictor. |
| |
| Keeps the object state as float32 internally (no text round-trip loss); |
| the Qwen semantic latent is refreshed from the text context only every |
| `tick_k` steps and reused (stale) in between — that is the dual rate. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
|
|
| import numpy as np |
| import torch |
|
|
| from .normalization import wrap_angle |
|
|
|
|
| def static_features_from_header(header: dict) -> tuple[np.ndarray, np.ndarray]: |
| """Returns (static (N, 8) float32, ids (N,) int32) in header object order.""" |
| objs = header["objects"] |
| static = np.zeros((len(objs), 8), dtype=np.float32) |
| ids = np.zeros(len(objs), dtype=np.int32) |
| for i, obj in enumerate(objs): |
| ids[i] = obj["id"] |
| is_circle = 1.0 if obj.get("type") == "circle" else 0.0 |
| mat = obj.get("material", {}) |
| if "radius" in obj: |
| s0 = s1 = float(obj["radius"]) |
| else: |
| s0 = float(obj.get("width", obj.get("size", 0.0)) or 0.0) |
| s1 = float(obj.get("height", obj.get("size", 0.0)) or 0.0) |
| static[i] = [is_circle, 1.0 - is_circle, |
| math.log(max(mat.get("mass", 1.0), 1e-3)), |
| mat.get("friction", 0.0), mat.get("elasticity", 0.0), |
| s0 / 100.0, s1 / 100.0, 0.0] |
| return static, ids |
|
|
|
|
| class DualRatePredictor: |
| def __init__(self, model, tokenizer, cfg, normalizer, device="cuda", |
| denoise_steps: int | None = None, sample_seed: int | None = 0): |
| self.model = model |
| self.tokenizer = tokenizer |
| self.cfg = cfg |
| self.norm = normalizer |
| self.device = device |
| self.denoise_steps = denoise_steps or cfg.denoise_steps |
| self.sample_seed = sample_seed |
| self.latent_refreshes = 0 |
| self._reset_fields() |
|
|
| def _reset_fields(self): |
| self.state = None |
| self.static_t = None |
| self.obj_mask = None |
| self.gravity_dt = None |
| self.ids = None |
| self.frame_num = 0 |
| self.steps = 0 |
| self.latent = None |
| self.latent_refreshes = 0 |
|
|
| def reset_from_arrays(self, static: np.ndarray, ids: np.ndarray, gravity, |
| dt: float, init_state: np.ndarray, start_frame: int): |
| """Seed directly from cached arrays (static (N,8), ids (N,), gravity (2,)).""" |
| self._reset_fields() |
| n = init_state.shape[0] |
| self.state = init_state.astype(np.float32).copy() |
| self.ids = ids[:n] |
| self.static_t = torch.from_numpy(static[:n]).float()[None].to(self.device) |
| self.obj_mask = torch.ones((1, n), dtype=torch.bool, device=self.device) |
| self.gravity_dt = torch.tensor( |
| [[float(gravity[0]) / 981.0, float(gravity[1]) / 981.0, float(dt) * 60.0]], |
| dtype=torch.float32).to(self.device) |
| self.frame_num = start_frame |
|
|
| def reset(self, header: dict, init_state: np.ndarray, start_frame: int): |
| """init_state: (N, 6) raw units; start_frame: frame number to emit first.""" |
| self._reset_fields() |
| static, ids = static_features_from_header(header) |
| n = init_state.shape[0] |
| self.state = init_state.astype(np.float32).copy() |
| self.ids = ids[:n] |
| self.static_t = torch.from_numpy(static[:n]).float()[None].to(self.device) |
| self.obj_mask = torch.ones((1, n), dtype=torch.bool, device=self.device) |
| g = header.get("gravity", {"x": 0.0, "y": -981.0}) |
| dt = header.get("timestep", 1 / 60) |
| self.gravity_dt = torch.tensor( |
| [[g["x"] / 981.0, g["y"] / 981.0, dt * 60.0]], dtype=torch.float32, |
| ).to(self.device) |
| self.frame_num = start_frame |
|
|
| def _refresh_latent(self, context_text: str): |
| ids = self.tokenizer(context_text, add_special_tokens=False)["input_ids"] |
| ids = ids[-self.cfg.max_ctx_tokens:] |
| t = torch.tensor([ids], dtype=torch.long, device=self.device) |
| mask = torch.ones_like(t) |
| self.latent = self.model.compute_latent(t, mask) |
| self.latent_refreshes += 1 |
|
|
| @torch.no_grad() |
| def step(self, context_text: str) -> str: |
| """Predict the next frame; refresh the semantic latent every tick_k steps.""" |
| if self.latent is None or self.steps % self.cfg.tick_k == 0: |
| self._refresh_latent(context_text) |
|
|
| cur_feat = self.state.copy() |
| cur_feat[:, 4] = wrap_angle(cur_feat[:, 4]) |
| cur_t = torch.from_numpy(self.norm.norm_state(cur_feat))[None].float().to(self.device) |
|
|
| gen = None |
| if self.sample_seed is not None: |
| gen = torch.Generator(device=self.device) |
| gen.manual_seed(self.sample_seed + self.steps) |
| delta_norm = self.model.sample_delta( |
| self.latent, cur_t, self.static_t, self.obj_mask, self.gravity_dt, |
| steps=self.denoise_steps, generator=gen, |
| ) |
| delta = self.norm.denorm_delta(delta_norm[0].float().cpu().numpy()) |
| self.state = self.state + delta |
|
|
| text = self._render_prediction() |
| self.frame_num += 1 |
| self.steps += 1 |
| return text |
|
|
| def _render_prediction(self) -> str: |
| """Matches the baseline prediction output format (parsed by _extract_physics).""" |
| lines = [f"Frame {self.frame_num}:"] |
| for i in range(self.state.shape[0]): |
| px, py, vx, vy, a, av = (float(x) for x in self.state[i]) |
| s = f" obj_{int(self.ids[i])}: pos=({px:.4f}, {py:.4f}), vel=({vx:.4f}, {vy:.4f})" |
| if abs(a) > 0.001 or abs(av) > 0.001: |
| s += f", a={a:.4f}, av={av:.4f}" |
| lines.append(s) |
| return "\n".join(lines) |
|
|
