README.md

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