Upload folder using huggingface_hub

README.md

@@ -0,0 +1,68 @@
+---
+license: mit
+library_name: jax
+tags:
+  - reinforcement-learning
+  - kinetix
+  - jax
+  - flax
+---
+
+# Kinetix pretrained agents
+
+General reinforcement learning agents for [Kinetix](https://github.com/FLAIROx/Kinetix), a 2D physics-based control environment written in JAX. Both agents were trained with [SFL](https://arxiv.org/abs/2408.15099) on procedurally generated levels, and use the transformer architecture from the [Kinetix paper](https://arxiv.org/abs/2410.23208) (786k parameters), entity observations and multi-discrete actions. 
+The 1M parallel environment run was based on [this](https://arxiv.org/abs/2603.06009) paper. Note that this checkpoint is not from the original paper, but was trained using the insights from this work.
+
+| Checkpoint | Description | Trained on | Parallel envs | Env steps | Hand-designed (S / M / L / all) | Random (S / M / L / all) |
+|---|---|---|---|---|---|---|
+| `sfl-paper` | The generalist agent from the Kinetix paper | random L levels | 2,048 | 18B | 0.42 / 0.32 / 0.11 / **0.22** | 0.33 / 0.24 / 0.16 / **0.24** |
+| `sfl-1m-envs` | The same architecture, trained at a much larger scale | random M levels | 1,048,576 | 376B | 0.70 / 0.43 / 0.18 / **0.33** | 0.56 / 0.36 / 0.22 / **0.38** |
+
+The numbers are average solve rates on the 74 hand-designed evaluation levels (20 attempts per level) and on 512 random levels per size (5 attempts per level). "all" averages over all levels.
+
+## Files
+
+Each checkpoint directory contains:
+- `params.safetensors`: the network parameters.
+- `config.json`: the model options needed to create the network (`model`), plus training details and evaluation results.
+
+## Usage
+
+```bash
+pip install kinetix-env huggingface_hub
+hf download mbeukman/Kinetix-Checkpoints --local-dir ./checkpoints
+```
+
+```python
+from kinetix.models import make_network_from_config
+from kinetix.util import load_pretrained_checkpoint
+
+params, pretrained_config = load_pretrained_checkpoint("./checkpoints/sfl-1m-envs")
+config |= pretrained_config["model"]  # your normalised Kinetix config
+network = make_network_from_config(env, env_params, config)
+hstate, pi, value = network.apply(params, hstate, (obs, done))
+```
+
+See [`examples/example_pretrained.py`](https://github.com/FLAIROx/Kinetix/blob/main/examples/example_pretrained.py) for a complete example that evaluates a checkpoint on the hand-designed levels.
+
+The `sfl-paper` checkpoint sets `legacy_entity_id: true`, which reproduces a quirk of the code it was trained with; this is handled automatically when you use its `config.json`.
+
+## Citation
+
+```bibtex
+@article{matthews2024kinetix,
+      title={Kinetix: Investigating the Training of General Agents through Open-Ended Physics-Based Control Tasks},
+      author={Michael Matthews and Michael Beukman and Chris Lu and Jakob Foerster},
+      booktitle={The Thirteenth International Conference on Learning Representations},
+      year={2025},
+      url={https://arxiv.org/abs/2410.23208}
+}
+
+
+@inproceedings{beukman2026preventing,
+  title={Preventing Learning Stagnation in PPO by Scaling to 1 Million Parallel Environments},
+  author={Beukman, Michael and Khetarpal, Khimya and Zheng, Zeyu and Dabney, Will and Foerster, Jakob and Dennis, Michael and Lyle, Clare},
+  booktitle={Reinforcement Learning Conference},
+  year={2026},
+}
+```

sfl-1m-envs/config.json

@@ -0,0 +1,58 @@
+{
+  "name": "Kinetix SFL generalist (1M parallel envs)",
+  "description": "The same architecture as the paper, trained with SFL at a much larger scale (1M parallel environments, 376B steps) on random M levels.",
+  "model": {
+    "actor_depth": 5,
+    "critic_depth": 5,
+    "actor_width": 128,
+    "critic_width": 128,
+    "activation": "tanh",
+    "aggregate_mode": "dummy_and_mean",
+    "full_attention_mask": false,
+    "transformer_depth": 2,
+    "transformer_size": 16,
+    "transformer_encoder_size": 128,
+    "num_heads": 8,
+    "transformer_ffn": false,
+    "multilayer_joint_thruster_mixing": false,
+    "permutation_invariant_mlp": false,
+    "recurrent_model": false,
+    "legacy_entity_id": false
+  },
+  "env": {
+    "observation_type": "symbolic_entity",
+    "action_type": "multi_discrete"
+  },
+  "training": {
+    "algorithm": "SFL + PPO",
+    "trained_on": "random M levels",
+    "num_parallel_envs": 1048576,
+    "env_steps": 376000000000.0,
+    "lr": 0.0006,
+    "sfl_buffer_size": 65536,
+    "hardware": "64x NVIDIA GH200",
+    "wandb_runs": [
+      "general-ai/kinetix-bc/jil19kf1"
+    ],
+    "code": "Kinetix (multi-GPU SFL)",
+    "date": "2026-02"
+  },
+  "evaluation": {
+    "hand_designed": {
+      "s": 0.695,
+      "m": 0.429,
+      "l": 0.175,
+      "all": 0.328,
+      "num_levels": 74,
+      "attempts_per_level": 20
+    },
+    "random": {
+      "s": 0.556,
+      "m": 0.355,
+      "l": 0.217,
+      "all": 0.376,
+      "num_levels": 1536,
+      "attempts_per_level": 5
+    }
+  }
+}

sfl-1m-envs/params.safetensors

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

sfl-paper/config.json

@@ -0,0 +1,60 @@
+{
+  "name": "Kinetix SFL generalist (paper)",
+  "description": "The SFL generalist agent from the Kinetix paper (L-size 'SFL' curve), trained on random L levels.",
+  "model": {
+    "actor_depth": 5,
+    "critic_depth": 5,
+    "actor_width": 128,
+    "critic_width": 128,
+    "activation": "tanh",
+    "aggregate_mode": "dummy_and_mean",
+    "full_attention_mask": false,
+    "transformer_depth": 2,
+    "transformer_size": 16,
+    "transformer_encoder_size": 128,
+    "num_heads": 8,
+    "transformer_ffn": false,
+    "multilayer_joint_thruster_mixing": false,
+    "permutation_invariant_mlp": false,
+    "recurrent_model": false,
+    "legacy_entity_id": true
+  },
+  "env": {
+    "observation_type": "symbolic_entity",
+    "action_type": "multi_discrete"
+  },
+  "training": {
+    "algorithm": "SFL + PPO",
+    "trained_on": "random L levels",
+    "num_parallel_envs": 2048,
+    "env_steps": 18000000000.0,
+    "lr": 5e-05,
+    "sfl_buffer_size": 1024,
+    "hardware": "1x NVIDIA L40S",
+    "wandb_runs": [
+      "general-ai/kinetix-experiments/3iihjaa2",
+      "general-ai/kinetix-experiments/rtyg0xvu"
+    ],
+    "code": "Jaxcredibots @ 5e3a0936 (pre-release Kinetix)",
+    "date": "2024-09",
+    "notes": "3iihjaa2 crashed at 6B steps and was resumed by rtyg0xvu for another 12B."
+  },
+  "evaluation": {
+    "hand_designed": {
+      "s": 0.415,
+      "m": 0.319,
+      "l": 0.111,
+      "all": 0.22,
+      "num_levels": 74,
+      "attempts_per_level": 20
+    },
+    "random": {
+      "s": 0.333,
+      "m": 0.236,
+      "l": 0.156,
+      "all": 0.242,
+      "num_levels": 1536,
+      "attempts_per_level": 5
+    }
+  }
+}

sfl-paper/params.safetensors

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