mbeukman/Kinetix-Checkpoints

Kinetix pretrained agents

General reinforcement learning agents for Kinetix, a 2D physics-based control environment written in JAX. Both agents were trained with SFL on procedurally generated levels, and use the transformer architecture from the Kinetix paper (786k parameters), entity observations and multi-discrete actions. The 1M parallel environment run was based on this 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

pip install kinetix-env huggingface_hub
hf download mbeukman/Kinetix-Checkpoints --local-dir ./checkpoints

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 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

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