AlexWortega/dualrate-physics

DualRate Physics v1

A dual-rate 2D rigid-body physics predictor. A frozen, truncated Qwen3.5-0.8B semantic encoder (first 12 of 24 layers + 8 readout tokens) reads the scene text (header + recent frames) and produces a latent that is refreshed only every K=8 frames; between refreshes a fast 38.9M flow-matching expert (per-object tokens, cross/self-attention, 10 Euler steps) advances the numeric object state [px, py, vx, vy, angle, angular_velocity] every frame. The slow semantic path is the "slow rate", the per-frame flow head is the "fast rate".

• 🎮 Interactive demo (ZeroGPU): https://huggingface.co/spaces/AlexWortega/physics-dualrate-playground
• 💻 Code & full training/eval pipeline (GitHub): https://github.com/AlexWortega/dualrate-physics-llm
• 🏋️ Checkpoints: ckpt_v1_60k.pt (final champion, step 60k) · ckpt_v1_45k.pt (earlier) · ckpt_v1_10k_ood.pt (10k-step checkpoint — best mid-training OOD before LR-decay recovery; see the checkpoint-trajectory table in results/RESULTS.md)
• 📐 Semantic encoder: Qwen/Qwen3.5-0.8B (downloaded from the Hub; its weights are not stored in these checkpoints — the files hold only the expert, projection, and conditioning weights + config + normalizer).

Results — ckpt_v1_60k.pt (champion)

Evaluated against the PhysicsLLMEngine LFM2-350M+LoRA text baseline.

Metric	DualRate (ours)	Baseline	Factor
Single-step position error (teacher-forced val)	0.052 px	LFM2: 22.64 px	×435
Single-step, eval rollout step 1 (200 scenes)	0.255 px	—	—
PhysicsLLMEngine exact protocol, overall MSE	0.650	862.21	×1327
— pong (baseline's worst case)	0.131	5027.39	×38,377
Wide OOD (6 unseen types × 50 scenes, ctx10 / 5 steps)	1.70 MSE	—	—
Inference	0.067–0.107 s/frame (A6000)	~30–60 s / 50 frames (CPU)	—

Exact protocol = bit-identical scenes (same generator code, seeds, difficulties as the baseline's llm_mse_results.json), same prompt format, same per-step position MSE, ctx = 10 frames, 5 autoregressive steps. Per-step MSE: ours [0.01, 0.09, 0.33, 0.89, 2.16] vs baseline [80.32, 308.6, 732.58, 1274.23, 1915.31].

Dual-rate ablation (K = latent refresh period, exact protocol)

K	overall MSE
1 (refresh every frame)	0.643
8 (refresh every 8 frames)	0.650

The stale latent is nearly free: refreshing the Qwen latent 8× less often costs ~1% accuracy. The slow semantic path can tick at ⅛ rate; all per-frame precision lives in the fast flow head.

Honest caveats

• Long rollouts drift: short horizons are excellent (~7.8 px RMS at step 10), but beyond ~30 steps chaotic many-body dynamics + error compounding dominate (no rollout-aware training yet — the obvious v2 lever).
• Trained on 54k val-derived scenes (17× smaller than the baseline's 900k train split) yet wins.
• OOD generalization dipped mid-training (high LR) and recovered via cosine decay — don't early-stop on seen-type metrics.

Repository contents

• ckpt_v1_*.pt — slimmed checkpoints (expert + projection + readout + config + normalizer; no optimizer state).
• results/ — RESULTS.md plus the raw evaluation JSONs (exact-protocol, wide-OOD, rollout, and the full checkpoint trajectory).
• code/ — the dualrate/ model package and the training/eval scripts, so these weights are reproducible from this repo. Full project (data generators, baselines, tests): https://github.com/AlexWortega/dualrate-physics-llm

Loading

import torch
from transformers import AutoTokenizer
from huggingface_hub import hf_hub_download
# from your copy of src/dualrate/
from dualrate.config import DualRateConfig
from dualrate.model import DualRatePhysicsModel
from dualrate.normalization import StateNormalizer

ckpt = torch.load(hf_hub_download("AlexWortega/dualrate-physics", "ckpt_v1_60k.pt"),
                  map_location="cpu", weights_only=False)
cfg = DualRateConfig.from_dict(ckpt["config"])
norm = StateNormalizer(**{k: ckpt["normalizer"][k]
                          for k in ("state_offset", "state_scale", "delta_scale")})
tok = AutoTokenizer.from_pretrained(cfg.model_name)
model = DualRatePhysicsModel(cfg, readout_token_id=tok.pad_token_id)
missing, unexpected = model.load_state_dict(ckpt["model"], strict=False)
# `semantic.lm.*` keys are "missing" on purpose — they come from the Qwen3.5 download.
model.eval().cuda()  # requires a CUDA GPU: Qwen3.5 linear-attention uses Triton kernels

⚠️ A CUDA GPU is required for inference: the Qwen3.5 linear-attention layers use flash-linear-attention Triton kernels that have no CPU backend.