Kinetix: Investigating the Training of General Agents through Open-Ended Physics-Based Control Tasks

TL;DR

The paper targets the perennial challenge of generalising RL agents beyond narrow, hand-designed environments by introducing Kinetix, a vast open-ended space of 2D physics tasks implemented in a hardware-accelerated engine (Jax2D). The authors train a general agent on millions of procedurally generated levels and demonstrate strong zero-shot generalisation to unseen handcrafted tasks, while showing that fine-tuning this base agent yields significantly better sample efficiency than training from scratch. Key contributions include the Jax2D engine, the Kinetix framework with a diverse level distribution and 74 holdout levels, and experimental evidence that large-scale online pre-training can produce robust, transferable control policies. This work provides a foundation for exploring open-ended RL, large-scale pretraining, and unsupervised environment design within a public, scalable framework.

Abstract

While large models trained with self-supervised learning on offline datasets have shown remarkable capabilities in text and image domains, achieving the same generalisation for agents that act in sequential decision problems remains an open challenge. In this work, we take a step towards this goal by procedurally generating tens of millions of 2D physics-based tasks and using these to train a general reinforcement learning (RL) agent for physical control. To this end, we introduce Kinetix: an open-ended space of physics-based RL environments that can represent tasks ranging from robotic locomotion and grasping to video games and classic RL environments, all within a unified framework. Kinetix makes use of our novel hardware-accelerated physics engine Jax2D that allows us to cheaply simulate billions of environment steps during training. Our trained agent exhibits strong physical reasoning capabilities in 2D space, being able to zero-shot solve unseen human-designed environments. Furthermore, fine-tuning this general agent on tasks of interest shows significantly stronger performance than training an RL agent *tabula rasa*. This includes solving some environments that standard RL training completely fails at. We believe this demonstrates the feasibility of large scale, mixed-quality pre-training for online RL and we hope that Kinetix will serve as a useful framework to investigate this further.

Figures (32)

• Figure 1: We train a general agent on randomly generated physics tasks and assess its transfer performance on hand-designed environments. In every environment the goal is to make the green shape touch the blue shape, without touching the red shape. The agent exerts control over every motor and thruster on each task.
• Figure 2: The transformer-based architecture used for training. The scene is decomposed into its constituent entities and then passed through the network, consisting of $L$ layers of self-attention and message passing, followed by $K$ fully connected layers.
• Figure 3: Zero-shot results on the holdout levels throughout training. In each pane, the training levels are sampled from the SFL distribution of the corresponding size, and the y-axis measures the solve rate on the evaluation set of that size. The shaded area shows the standard error over 5 seeds.
• Figure 4: Heatmaps of goal $x$ position and morphology $x$ position. An ideal agent that can perfectly maneuver a morphology to under the goal position would manifest itself as a diagonal line.
• Figure 5: The performance of fine-tuned and tabula rasa agents (left) aggregated over the entire L holdout set, and (right) for four selected levels. We train a separate agent for each environment and plot mean and standard error over five seeds. We stress that the MuJoCo levels are reimplementations of the classic environments in Kinetix.