Diffusion Pong EMA step 650 + live play card

README.md

@@ -1,3 +1,57 @@
----
-license: mit
----
+---
+license: apache-2.0
+library_name: pytorch
+pipeline_tag: other
+tags:
+  - world-model
+  - diffusion
+  - pong
+  - video
+  - causal
+  - game
+---
+
+# Diffusion Pong
+
+A small action-conditioned world model that learned to keep a Pong match going. You hold **W** / **S**, it invents the next frame. That's the whole demo.
+
+![demo](demo.gif)
+
+Weights are the EMA from training step **650** (~59M DiT params). Resolution is 128×128 at 6 FPS with a 12-frame latent history. Codec is a frozen SDXL VAE (`madebyollin/sdxl-vae-fp16-fix`). Sampling is 1-step instantaneous velocity — don't crank the step count expecting magic; this checkpoint is still in the teacher stage.
+
+## Play it
+
+From the livediffusion source tree (after `pip install -e .`):
+
+```bash
+python scripts/live_infer.py
+```
+
+First run pulls this repo into the HF cache. Needs a CUDA GPU with BF16.
+
+```bash
+# optional knobs
+python scripts/live_infer.py --steps 1 --window-scale 7 --seed 42
+```
+
+Controls: **W** up, **S** down. Click the window once so it has focus. Clicking the canvas drops a yellow ball into the latent history for a few frames if you want to poke it.
+
+## What's in the files
+
+| file | what |
+|---|---|
+| `ema.safetensors` | playable weights |
+| `config.json` | video / model / codec settings used at train time |
+| `demo.gif` | live rollout clip |
+
+Action vector is `[W, S]` as floats in `{0,1}`. Hold both and it treats that as neutral.
+
+## Notes / honesty
+
+- The right paddle is baked into the learned dynamics (training used a noisy scripted opponent). You're only steering the left one.
+- Autoregressive rollouts drift. Short play sessions look good; leave it running forever and the court starts inventing its own physics fanfic.
+- Not a RL agent. No reward head in use. Just next-frame generation conditioned on your button state.
+
+Trained with the Livediffusion prefix-transformer stack (schema v3): clean history prefix, noisy suffix, 3D RoPE, multiscale context tokens.
+
+If something breaks, it's probably the VAE download or a GPU that can't do BF16.

config.json

@@ -0,0 +1,108 @@
+{
+  "codec": {
+    "enable_slicing": true,
+    "enable_tiling": false,
+    "frame_batch_size": 256,
+    "kind": "frame_vae",
+    "latent_channels": 4,
+    "model_id": "madebyollin/sdxl-vae-fp16-fix",
+    "revision": "main",
+    "sample_posterior": false,
+    "spatial_compression": 8,
+    "subfolder": "",
+    "temporal_compression": 1
+  },
+  "data": {
+    "batch_size": 128,
+    "data_path": "",
+    "dataset_kind": "pong",
+    "frame_stride": 1,
+    "hf_repo_id": "",
+    "hf_revision": "main",
+    "latent_cache_path": "",
+    "latent_prefetch_batches": 8,
+    "network_clip_prefetch": 2,
+    "num_workers": 8,
+    "persistent_workers": true,
+    "pin_memory": true,
+    "prefetch_batches": 2,
+    "sample_shuffle": 64,
+    "shard_shuffle": 256,
+    "train_samples": 50000,
+    "validation_samples": 1024,
+    "windows_per_video": 4
+  },
+  "inference": {
+    "frames": 120,
+    "steps": 1
+  },
+  "model": {
+    "action_dim": 2,
+    "context_dim": 0,
+    "depth": 12,
+    "fine_history_frames": 2,
+    "fine_memory_patch_size": 2,
+    "generation_frames": 1,
+    "hidden_size": 384,
+    "memory_depth": 3,
+    "memory_patch_size": 4,
+    "mlp_ratio": 4.0,
+    "num_heads": 8,
+    "qk_norm": true,
+    "target_patch_size": 2
+  },
+  "schema_version": 3,
+  "simulation": {
+    "max_shapes": 2,
+    "max_size": 6.0,
+    "max_speed": 20.0,
+    "min_shapes": 1,
+    "min_size": 3.0,
+    "min_speed": 8.0,
+    "supersample": 2,
+    "warmup_frames": 8
+  },
+  "training": {
+    "activation_checkpointing": true,
+    "critic_learning_rate": 0.0001,
+    "distribution_matching_start_step": 30000,
+    "distribution_matching_weight": 0.05,
+    "ema_decay": 0.9999,
+    "epochs": 100,
+    "flow_matching_ratio": 0.25,
+    "full_interval_ratio": 0.25,
+    "gradient_accumulation_steps": 1,
+    "history_noise_std": 0.01,
+    "keep_checkpoints": 5,
+    "learning_rate": 0.0003,
+    "log_every": 20,
+    "max_grad_norm": 1.0,
+    "min_learning_rate": 1e-05,
+    "on_policy_max_horizon": 4,
+    "on_policy_ratio": 0.15,
+    "on_policy_sampling_steps": 4,
+    "on_policy_start_step": 15000,
+    "output_dir": "runs/pong-action-world-model",
+    "reward_loss_weight": 0.0,
+    "save_every": 50,
+    "seed": 2026,
+    "teacher_endpoint_ratio": 0.1,
+    "teacher_integration_steps": 8,
+    "teacher_logit_mean": 1.0,
+    "teacher_logit_std": 1.0,
+    "teacher_steps": 10000,
+    "teacher_x0_weight": 0.25,
+    "terminal_loss_weight": 0.0,
+    "validate_every": 50,
+    "validation_batches": 1,
+    "warmup_steps": 500,
+    "weight_decay": 0.0
+  },
+  "video": {
+    "fps": 6.0,
+    "future_frames": 5,
+    "height": 128,
+    "history_frames": 12,
+    "width": 128
+  }
+}

ema.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1f9210f3f6a2f9cebde3c1a89fae69665a51d81cc53b922e3c9fec1625446007
+size 235185584