Skip to content
BranchesTags

Latest commit

 

History

14 Commits

Folders and files

NameName
Last commit message
Last commit date
assets
assets
 
 
browser_demo
browser_demo
 
 
data_samples/planetary_rotation
data_samples/planetary_rotation
 
 
demo_checkpoint
demo_checkpoint
 
 
evaluation_results
evaluation_results
 
 
inference
inference
 
 
scripts
scripts
 
 
src
src
 
 
GENERALIZATION_REPORT.md
GENERALIZATION_REPORT.md
 
 
HUGGINGFACE_UPLOAD.md
HUGGINGFACE_UPLOAD.md
 
 
MODEL_CARD.md
MODEL_CARD.md
 
 
README.md
README.md
 
 
comparison_bowling.gif
comparison_bowling.gif
 
 
comparison_chain_reaction.gif
comparison_chain_reaction.gif
 
 
comparison_particle_explosion.gif
comparison_particle_explosion.gif
 
 
comparison_solar_system.gif
comparison_solar_system.gif
 
 
demo_comparison.py
demo_comparison.py
 
 
evaluate_llm_mse.py
evaluate_llm_mse.py
 
 
generalization_test_results.json
generalization_test_results.json
 
 
llm_mse_results.json
llm_mse_results.json
 
 
llm_vs_baseline_comparison.png
llm_vs_baseline_comparison.png
 
 
mse_baseline_results.json
mse_baseline_results.json
 
 
mse_baselines_average.png
mse_baselines_average.png
 
 
mse_baselines_visualization.png
mse_baselines_visualization.png
 
 
test_generalization.py
test_generalization.py
 
 
test_mse_baseline.py
test_mse_baseline.py
 
 
visualize_mse.py
visualize_mse.py
 
 

Repository files navigation

βš™οΈ PhysicsLLMEngine

Teaching a 350M-parameter language model to simulate 2D rigid body physics β€” from text alone.

All 30 physics scenario types

πŸš€ Key Numbers

30 scenario types  Β·  900K training scenes  Β·  180M frames  Β·  ~582 GB of physics data
Sub-pixel positional accuracy (<0.2%) on zero-shot held-out scenarios

✨ Highlights
🎯 30 diverse scenarios Billiards, towers, pendulums, angry birds, orbit, hourglass and more
πŸ“ Text-only input No vision encoder, no GNN β€” pure next-token prediction
πŸŽ“ Curriculum learning 5 difficulty stages: 2-obj ballistics β†’ 50-obj chaos
πŸ”¬ Zero-shot generalization 6 held-out scenario types never seen during training
🌐 Runs in the browser WebGPU/WASM via transformers.js β€” no server needed
⚑ 350M params Fine-tuned LiquidAI/LFM2-350M with LoRA

🎬 Scenario Zoo

All 30 Scenarios β€” click to collapse

πŸ’₯ Collision & Ballistics

Billiards Bowling Explosion Head-on Projectile

πŸ—οΈ Stacking & Structural

Bridge Dominos Jenga Pyramid Tower

🎿 Ramps & Terrain

Funnel Marble Run Plinko Ramp Roll Ski Jump

πŸ”— Pendulums & Constraints

Chain Newton's Cradle Pendulum Seesaw Wrecking Ball

πŸ•ΉοΈ Mini-game Physics

Angry Birds Basketball Breakout Pinball Pong

πŸŒ€ Complex & Chaotic

Avalanche Conveyor Hourglass Orbit Wind

πŸ—ƒοΈ Dataset

AlexWortega/physics-scenarios-packed Β· AlexWortega/physics-scenarios-raw

Split Scenes Scenario Types Purpose
Train 900,000 24 (seen) Supervised fine-tuning
Val 100,020 30 (all) In-dist + zero-shot generalization

Six scenario types are held out from training entirely:

pong Β· bowling Β· ramp_roll Β· angry_birds Β· hourglass Β· newtons_cradle

Each scene = 1 JSONL header + 200 frame lines. The model receives frames 1..N and must predict frame N+1 in plain text:

Scene: Billiards: cue ball strikes a triangle of 21 balls.
Gravity: (0.0, 0.0)
Timestep: 0.01667
Type: billiards  Difficulty: 3

Frame 1: All objects at rest.
  obj_0: pos=(500.0000, 300.0000), vel=(800.0000, 0.0000)
  obj_1: pos=(522.8631, 286.8000), vel=(0.0000, 0.0000)
  ...

Frame 2: Objects in motion.
  obj_0: pos=(513.3340, 300.0000), vel=(800.0000, 0.0000)
  ...

Predict next frame:

🧠 Model & Training

Downloads

Format Link Size Use case
Merged weights AlexWortega/lfm2-scenarios ~700 MB Python / inference
ONNX q4 AlexWortega/lfm2-scenarios-ONNX 458 MB Browser (WebGPU/WASM)
GGUF Q4_K_M AlexWortega/lfm2-scenarios-GGUF 216 MB llama.cpp / local
Base model LiquidAI/LFM2-350M β€” Reference

Architecture

Component Details
Base LiquidAI/LFM2-350M (Liquid Foundation Model)
Fine-tuning LoRA r=32 Ξ±=64 via Unsloth
Context 8 192 tokens
Precision bfloat16
Optimizer AdamW 8-bit, lr=2e-4
Batch 4 Γ— grad-accum 8 = effective 32

Curriculum Learning

Stage 0  difficulty 1  2–5 obj     50 000 examples  β†’ loss 0.562 βœ…
Stage 1  difficulty 2  5–15 obj    50 000 examples  β†’ loss 0.609 βœ…
Stage 2  difficulty 3  10–30 obj   50 000 examples  β†’ loss 0.622 βœ…
Stage 3  difficulty 4  20–40 obj   50 000 examples  β†’ (training)
Stage 4  difficulty 5  30–50+ obj  50 000 examples  β†’ (pending)

Infrastructure

  • GPU β€” NVIDIA RTX A6000 (48 GB VRAM)
  • Training time β€” ~7–14 h per curriculum stage
  • Data generation β€” 22-core CPU, Pymunk, ~29 min for 900K scenes
  • Logging β€” Weights & Biases

πŸ“Š Results: Zero-Shot Physics Prediction

Evaluated on 6 held-out scenario types the model never saw during training.

Pong β€” Sub-pixel accuracy on ballistic motion

GROUND TRUTH   obj_0: pos=(339.8344, 143.2513)  vel=(-531.0383, -319.9974)
PREDICTION     obj_0: pos=(339.7855, 142.0493)  vel=(-531.0383, -319.9974)
                            Ξ”x = 0.05  Ξ”y = 1.20          velocity: exact βœ“

< 0.2% positional error on an 800 Γ— 600 canvas.

Bowling β€” Multi-body collision cascade (11 objects)

GROUND TRUTH   obj_3:  pos=(600.8966, 25.1189)  vel=(  0.0000, -179.8500)
PREDICTION     obj_3:  pos=(600.8966, 25.7189)  vel=(  0.0000, -179.8500)
                            Ξ”x = 0.00  Ξ”y = 0.60          velocity: exact βœ“

GROUND TRUTH   obj_10: pos=(160.2332, 71.8437)  vel=(264.8960,    0.0000)
PREDICTION     obj_10: pos=(160.2332, 71.8594)  vel=(264.8960,    0.0000)
                            Ξ”x = 0.00  Ξ”y = 0.02          velocity: exact βœ“

Key observations

  1. Velocities recovered near-perfectly β€” model learns the linear position-velocity relationship
  2. Sub-pixel positional errors even on multi-body collisions in unseen scenarios
  3. Generalizes compositionally β€” applies learned primitives (gravity, impulse, constraints) to new scenario types

🌐 Browser Demo

Run the model entirely in your browser β€” no Python, no GPU, no server:

cd browser_demo
npm install
npm run dev     # β†’ http://localhost:5173

The demo fetches lfm2-scenarios-ONNX (458 MB, q4) from HF CDN, runs autoregressive rollout via WebGPU (or WASM fallback), and renders each predicted frame on a Konva canvas with a live token-stream sidebar.

Source: browser_demo/

πŸ–₯️ Local Inference

pip install llama-cpp-python matplotlib pillow

cd inference

# Single-step repro (mirrors browser prompt format exactly)
python repro.py billiards 5

# Multi-step rollout with drift/parse diagnostics
python multistep.py orbit 50

# Benchmark 1-frame vs 4-frame prompt speed (ONNX)
python bench.py

# Render all 30 demo scenarios Γ— 200 frames β†’ animated GIFs
python make_gifs.py --frames 200 --out ./gifs

Source: inference/

πŸƒ Training & Evaluation

# Generate dataset
python scripts/generate_scenarios_dataset.py \
  --output-dir data_scenarios/train \
  --num-scenes-per-type 37500 \
  --num-workers 22

# Train with curriculum
python scripts/train_finetune.py \
  --data-dir data_scenarios/train \
  --output-dir checkpoints/lfm2 \
  --curriculum-stages 5 \
  --epochs-per-stage 1 \
  --batch-size 4 --grad-accum 8 --lr 2e-4

# Evaluate on held-out scenarios
python scripts/run_evaluation.py \
  --model finetune \
  --checkpoint checkpoints/lfm2/stage4/adapter \
  --output-dir evaluation_results

πŸ“ Project Structure

PhysicsLLMEngine/
β”œβ”€β”€ src/
β”‚   β”œβ”€β”€ physics/
β”‚   β”‚   β”œβ”€β”€ scenario_generator.py   # All 30 scenario generators
β”‚   β”‚   β”œβ”€β”€ scenario_registry.py    # @register_scenario decorator
β”‚   β”‚   β”œβ”€β”€ simulation.py           # Pymunk wrapper
β”‚   β”‚   └── objects.py              # Body/shape factories
β”‚   β”œβ”€β”€ data/
β”‚   β”‚   β”œβ”€β”€ formats.py              # JSONL text serialization
β”‚   β”‚   └── exporter.py             # Scene β†’ file pipeline
β”‚   β”œβ”€β”€ training/
β”‚   β”‚   β”œβ”€β”€ curriculum.py           # Difficulty-based curriculum scanner
β”‚   β”‚   └── data_loader.py          # Physics-aware data loading
β”‚   └── evaluation/
β”‚       β”œβ”€β”€ rollout.py              # Autoregressive multi-step evaluator
β”‚       └── runner.py               # Full evaluation pipeline CLI
β”œβ”€β”€ scripts/
β”‚   β”œβ”€β”€ train_finetune.py           # LFM2 + LoRA entry point
β”‚   β”œβ”€β”€ train_scratch.py            # GPT from-scratch baseline
β”‚   β”œβ”€β”€ generate_scenarios_dataset.py
β”‚   └── run_evaluation.py
β”œβ”€β”€ browser_demo/                   # 🌐 WebGPU browser demo
β”‚   β”œβ”€β”€ src/
β”‚   β”‚   β”œβ”€β”€ transformersEngine.ts   # transformers.js ONNX inference
β”‚   β”‚   β”œβ”€β”€ streamClient.ts         # Autoregressive rollout + fitPrompt
β”‚   β”‚   β”œβ”€β”€ App.tsx                 # React UI (canvas + token stream)
β”‚   β”‚   └── promptFormat.ts         # Text serialization (mirrors training)
β”‚   └── backend/
β”‚       β”œβ”€β”€ server.py               # FastAPI scenario server
β”‚       └── examples/               # 30 bundled JSONL demo scenarios
β”œβ”€β”€ inference/                      # πŸ–₯️ Local Python inference tools
β”‚   β”œβ”€β”€ repro.py                    # Single-step repro
β”‚   β”œβ”€β”€ multistep.py                # Multi-step rollout diagnostics
β”‚   β”œβ”€β”€ bench.py                    # ONNX speed benchmark
β”‚   β”œβ”€β”€ make_gifs.py                # Batch rollout β†’ GIFs
β”‚   └── patch_onnx_v2.py            # Patch LoRA weights into ONNX graph
└── assets/
    β”œβ”€β”€ gallery.png                 # Scenario gallery overview
    └── gifs/                       # 30 scenario demo GIFs

πŸ—ΊοΈ What's Next

  • Complete curriculum stages 3–4 (hard + extreme difficulty)
  • Full 200-step autoregressive rollout evaluation on all 30 scenarios
  • GPT-from-scratch baseline with muP scaling
  • Energy & momentum conservation analysis
  • q4f16 ONNX for 2Γ— faster WebGPU inference

πŸ“„ Citation

@software{physicslmengine2026,
  title   = {PhysicsLLMEngine: Learning Rigid Body Dynamics via Next-Token Prediction},
  author  = {Wortega, Alex},
  year    = {2026},
  url     = {https://github.com/AlexWortega/PhysicsLLMEngine}
}

πŸ”— Links
πŸ€— Merged model AlexWortega/lfm2-scenarios
πŸ€— ONNX (WebGPU) AlexWortega/lfm2-scenarios-ONNX
πŸ€— GGUF (llama.cpp) AlexWortega/lfm2-scenarios-GGUF
πŸ€— Dataset (packed) AlexWortega/physics-scenarios-packed
πŸ€— Dataset (raw) AlexWortega/physics-scenarios-raw
🌐 Browser demo browser_demo/
πŸ–₯️ Inference scripts inference/
⚑ Base model LiquidAI/LFM2-350M
πŸ› οΈ Training framework Unsloth

MIT License  Β·  ICML 2026

About

Training LLMs to predict 2D rigid body physics β€” ICML 2026

Resources

Readme
Activity

Stars

4 stars

Watchers

0 watching

Forks

1 fork
Report repository

Releases

Packages

Contributors

Languages