| --- |
| 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}, |
| } |
| ``` |
|
|
