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

## Abstract

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, 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. 1 1 1 We provide full code and models at https://kinetix-env.github.io .

## 1 Introduction

Figure: 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.
Refer to caption: extracted/5966380/images/fig1_v6.png

The development of a general agent, capable of performing competently in unseen domains, has been a long-standing goal in machine learning (Newell et al., ; Minsky, ; Lake et al., ). One perspective is that large transformers, trained on vast amounts of offline text and video data, will ultimately achieve this goal (Brown et al., ; Bubeck et al., ; Mirchandani et al., ). However, applying these techniques in an offline reinforcement learning (RL) setting often constrains agent capabilities to those found within the dataset (Levine et al., ; Kumar et al., ).
An alternative approach is to use online RL, where the agent gathers its own data through interaction with an environment. However, with some notable exceptions (Team et al., ; ), most RL environments represent a narrow and homogeneous set of scenarios (Todorov et al., ; Bellemare et al., ; Brockman et al., ; Cobbe et al., ), limiting the generalisation ability of the trained agents (Kirk et al., ).

In this paper, we aim to address this limitation by introducing Kinetix: a framework for representing the vast, open-ended space of 2D physics-based environments, and using it to train a general agent. Kinetix is broad enough to represent robotics tasks like grasping (Rajeswaran et al., ) and locomotion (Todorov et al., ), classic RL environments such as Cartpole (Barto et al., ), Acrobot (DeJong & Spong, ) and Lunar Lander (Brockman et al., ), as well as video games like Pinball (Bellemare et al., ), along with the multitude of tasks that lie in the intervening space (see Figure 1).
To run the backend of Kinetix we developed Jax2D, a hardware-accelerated physics engine that allows us to efficiently simulate the billions of environment interactions required to train this agent.(^2^22[https://github.com/MichaelTMatthews/Jax2D](https://github.com/MichaelTMatthews/Jax2D))

Through sampling random Kinetix environments from the space of representable 2D physics problems, we can produce a virtually unlimited supply of meaningfully diverse tasks for training. Since these levels are programmatically sampled, many are not useful for learning—indeed most are either trivial or unsolvable. Training on this large, diverse set of mixed-quality levels mirrors the pretraining stage of a language model (Devlin et al., ; Brown et al., ; Dubey et al., ).

We find that an RL agent trained on these environments exhibits understanding of general mechanical properties, with the ability to zero-shot solve unseen handmade environments (Section 5).
We further analyse the benefits of fine-tuning this general agent on specific hard environments and find that it greatly reduces the number of samples required to learn a particular task, when comparing against a tabula rasa agent. Fine-tuning also affords new capabilities, including solving tasks for which an agent specifically trained does not make progress (Section 6).

In summary, our contributions are:

- 1.
We introduce Jax2D, a fast hardware-accelerated 2D physics engine.
- 2.
We introduce Kinetix, an open-ended space of RL environments within a unified framework. We provide the capability to sample random levels from the vast space of possible physics tasks, as well as providing a large set of 74 interpretable handmade levels.
- 3.
We demonstrate the zero-shot generalisation ability of an agent trained on Kinetix.
- 4.
We show that fine-tuning this general agent on difficult tasks leads to significantly improved sample efficiency and new capabilities.

## 2 Background

### 2.1 Reinforcement Learning

We model the decision-making process as a Markov Decision Process (MDP), which is defined as a tuple $\langle\mathcal{S},\mathcal{A},\mathcal{R},\mathcal{T}\rangle$, where $\mathcal{S}$ is the set of states; $\mathcal{A}$ is the set of actions; $\mathcal{T}:\mathcal{S}\times\mathcal{A}\to\Delta\mathcal{S}$ is the transition function, defining the distribution over next states $\mathcal{T}(s,a)$ given a current state $s$ and action $a$; and $\mathcal{R}:\mathcal{S}\to\mathbb{R}$ is the reward function.
We consider finite-horizon MDPs, with a maximum number of timesteps $T$. The goal of an agent in RL is to maximise its discounted sum of rewards, $G_{t}\dot{=}\sum_{t=0}^{T}\gamma^{t}R_{t}$, where $R_{t}\dot{=}\mathcal{R}(s_{t})$ is the reward at timestep $t$ and $\gamma$ is the discount factor.

### 2.2 Unsupervised Environment Design

Unsupervised Environment Design (UED) is a paradigm where learning is phrased as a two-player game between a teacher and a student.
The student maximises its expected discounted return as in the standard RL formulation, while the teacher chooses levels to maximise some utility function, effectively inducing a curriculum of tasks through training (Oudeyer et al., ; Florensa et al., ; Matiisen et al., ; Narvekar et al., ; Dennis et al., ; Parker-Holder et al., ).
In this paper, these tasks (we also refer to these as levels or environments) are particular initial states, $s_{0}\in\mathcal{S}$.
One common approach sets a level’s utility as the negative of the agent’s return (Pinto et al., ), and another class of approaches instead uses regret (Dennis et al., ).
Domain Randomisation (Jakobi, ; Tobin et al., , DR), where levels are sampled from an uninformed distribution, can be considered a degenerate form of this paradigm, where a constant utility is assigned to each level. More recently, Tzannetos et al. () and Rutherford et al. () sample levels in binary-outcome domains using learnability, defined as $p(1-p)$, with $p$ being the success rate of the agent on the particular level. In this way, learnability disincentivises the teacher from sampling levels that the agent cannot solve at all (where $p=0$) or where the agent can already perfectly solve them ($p=1$), meaning that the agent trains on levels with a high learning potential.

### 2.3 RL in JAX

JAX (Bradbury et al., ) is a Python library for writing parallelisable code for hardware accelerators. While deep RL has traditionally been divided between environments on the CPU and models on the GPU (Mnih et al., ; Espeholt et al., ), JAX has facilitated the development of GPU-based environments (Lange, ; Rutherford et al., ; Nikulin et al., ; Matthews et al., ; Kazemkhani et al., ), allowing the entire RL pipeline to run on a hardware accelerator (Hessel et al., ). Through massive parallelisation and elimination of CPU-GPU transfer, this gives tremendous speed benefits (Lu et al., ).
While UED has also followed this trend (Jiang et al., ; Coward et al., ), experiments have largely been confined to simple gridworlds, due to the lack of any suitable alternative (Garcin et al., ; Rutherford et al., ).

### 2.4 Transformers and Permutation Invariant Representations

Transformers and Attention
Transformers (Vaswani et al., ) use the attention mechanism (Bahdanau et al., ) to model interactions within a set. Given $N$ embeddings, ${x_{i}}_{1}^{N}\in\mathbb{R}^{n}$, self-attention computes queries $q_{i}$, keys $k_{i}$, and values $v_{i}$ for each element through linear projections. Weights for each element $i$ relative to element $j$ are calculated as $w_{i,j}\dot{=}q_{i}\cdot k_{j}$ and normalised via softmax to get $\tilde{w}_{i,j}$. The new embedding for element $i$ is a weighted sum of the values: $x^{\text{new}}_{i}\dot{=}\sum_{j=1}^{N}\tilde{w}_{i,j}v_{j}$, allowing each element to attend to others. The common practice of adding positional embeddings to encode sequence order (Vaswani et al., ) may obfuscate the fact that transformers are permutation invariant and naturally operate on sets.

Transformers in RL
While recurrent policies have been long popular in deep RL to help deal with partial observability, sequence models like transformers are gaining traction as an alternate solution (Lu et al., ; Bousmalis et al., ; Team et al., ; Raparthy et al., ).
A less common use of transformers in RL is for processing inherently permutation-invariant observations, such as entities in Starcraft II (Vinyals et al., ).
Although graphs are traditionally processed with graph neural networks (Wang et al., ; Battaglia et al., ), transformers are also now being applied to this domain (Sferrazza et al., ; Buterez et al., ), with attention masks set to a graph’s adjacency matrix to restrict attention to neighboring nodes (Sferrazza et al., ).

## 3 Kinetix

In this section, we introduce Kinetix, a large and open-ended environment for RL, implemented entirely in JAX. We describe our underlying physics engine (Section 3.1), the RL environment (Section 3.2), and finally propose Kinetix as a novel challenge for open-endedness (Section 3.3).

### 3.1 Jax2D

Jax2D is our deterministic, impulse-based, 2D rigid-body physics engine, written entirely in JAX, that forms the foundation of the Kinetix benchmark. We designed Jax2D to be as expressive as possible through simulation of only a few fundamental components. To this end, a Jax2D scene contains only 4 unique entities: circles, (convex) polygons, joints and thrusters. From these simple building blocks, a huge diversity of different physical tasks can be represented.

Jax2D simulates discrete Euler steps for rotational and positional velocities and then applies instantaneous impulses and higher order corrections to solve constraints.
The notion of a constraint encompasses collisions (two objects cannot be inside each other) and joint constraints (two objects connected by a joint cannot separate at the point of connection). Constraints are pairwise, meaning that it may be necessary to apply multiple steps of constraint solving for a stable simulation, especially when simulating systems of many interacting bodies. The number of solver steps therefore serves as a tradeoff between accuracy and speed.
An agent (human or artificial) can act on the scene by applying torque through motors attached to revolute joints or by applying force through thrusters.

Jax2D is based on Box2D (Catto, ) and can be thought of as a minimalist rewrite of the C library in JAX. Appendix B shows the benefit of this reimplementation, with hardware acceleration allowing Jax2D to easily scale to thousands of parallel environments on a single GPU, outperforming Box2D by a factor of $4\times$ when comparing just the engines and $30\times$ when training an RL agent (this difference is due to Jax2D natively integrating with RL pipelines that exist entirely on the GPU).

The key differentiator of Jax2D from other JAX-based physics simulators such as Brax (Freeman et al., ), is that Jax2D scenes are almost entirely dynamically specified, meaning that the same underlying computation graphs are run for every simulation. For example, this means that running Half-Cheetah, Pinball and Grasper (Figure 1) involves executing the exact same instructions. This allows us to parallelise across different tasks with the JAX vmap operation—a crucial component of harnessing the power of hardware acceleration in a multi-task RL setting. Brax, by contrast, is almost entirely statically specified meaning it is impossible to vmap across, for instance, different morphologies. Further Jax2D implementation details are discussed in Appendix A.

### 3.2 Kinetix : RL Environment Specification

Kinetix builds on Jax2D to create an environment for RL, which we now briefly outline.

Action Space
Kinetix supports both multi-discrete and continuous action spaces.
In the multi-discrete action space, each motor and thruster can either be inactive, or activated at maximum power each timestep, with motors being able to be run either forwards or backwards.
In the continuous action space, motors can be powered in the range $[-1,1]$ and thrusters in the range $[0,1]$.

Observation Space
We use a symbolic observation where each entity (shape, joint or thruster) is defined by an array of values of physical properties including position, rotation and velocity. The observation is then defined as the set of these entities, allowing the use of permutation-invariant network architectures such as transformers. This observation space makes the environment fully observable, removing the need for a policy with memory. We also provide the option for pixel-based observations and a symbolic observation that simply concatenates and flattens the entity information.

Reward
To facilitate our goal of a general agent, we choose a simple yet highly expressive reward function that remains fixed across all environments. Each scene must contain a green shape and a blue shape—the goal is simply to make these two shapes collide, upon which the episode terminates with a reward of $+1$.
Scenes can also contain red shapes, which, if they collide with the green shape, will terminate the episode with $-1$ reward. As demonstrated in Figure 1, these simple and interpretable rules allow for a large number of semantically diverse environments to be represented. To improve learning, we augment this sparse reward with an auxiliary dense reward signal, defined as $R^{d}_{t}=\kappa\left(d_{t}-d_{t+1}\right)$, where $d_{t}$ is the distance between the green and blue objects at timestep $t$ and $\kappa$ is a coefficient that we tune to ensure the dense signal does not dominate.

### 3.3 Kinetix : A Benchmark for Investigating Open-Endedness

The expressivity, diversity, and speed of Kinetix makes it an ideal environment for studying open-endedness, including generalist agents, UED, and lifelong learning. In order to make it maximally effective for agent training and evaluation, we provide a heuristic environment generator, a set of hand-designed levels, and an environment taxonomy describing the complexity of environments.

Environment Generator
The strength of Kinetix lies in the diversity of environments it can represent. However, this environment set contains many degenerate cases, which can dominate the distribution if sampled from naïvely. For this reason, we provide a random level generator that is designed to be maximally expressive, while minimising the number of degenerate levels.
We ensure that every level has exactly one green and blue shape, and at least one controllable aspect (either a motor or a thruster). Furthermore, we follow Team et al. () and perform rejection sampling on levels solved with a no-op policy (defined as the policy that activates no motors or thrusters), thus eliminating trivial levels.
The remaining pathology is unsolvable levels, which are largely intractable to determine and for which we will rely on automatic curriculum methods to filter out.

Each level is built up iteratively from an empty base by adding shapes either freely or connected to an already existing shape. We perform rejection sampling on proposed shape additions to try and ensure that no collisions are active in the initial level state. These methods to add shapes (along with analogous methods for editing and removing) can also serve as mutators for automatic level editing algorithms like ACCEL (Parker-Holder et al., ). We also provide functionality to generate levels using RL (Dennis et al., ) and generative models (Garcin et al., ).

Hand-Designed Levels
Along with the capability to sample random levels, Kinetix contains a suite of 74 hand designed levels (Appendix E), as well as a powerful graphical editor to facilitate the creation of new levels. Some of these levels are inspired by other RL benchmarks, such as L-MuJoCo-Walker, L-MuJoCo-Hopper, L-MuJoCo-Half-Cheetah, L-MuJoCo-Swimmer (Todorov et al., ) and L-Lunar-Lander, L-Swing-Up, L-Cartpole-Wheels-Hard (Brockman et al., ). We made other levels, like L-Pinball, L-Lorry and L-Catapult, specifically for Kinetix. These levels tests agent capabilities including fine-grained motor control, navigation, planning and physical reasoning.

Environment Taxonomy
Kinetix has the useful characteristic of containing a controllable and interpretable axis of complexity—the number of each type of entity in a scene.
While not a strict rule, scenes with less entities tend to represent simpler problems.
We therefore quantise our experiments and handmade levels into one of three distinct sizes: small (S), medium (M), and large (L). A convenient feature of the entity-based observation space is that an agent trained on one level size can also meaningfully operate in other sizes, just as a language model can condition on a variable number of tokens, allowing us to interoperate between the sizes.

## 4 Experimental Setup

We train on programatically generated Kinetix levels drawn from the statically defined distribution. We refer to training on sampled levels from this distribution as DR.
Our main metric of assessment is the solve rate on the set of handmade holdout levels.
The agent does not train on these levels but they do exist inside the support of the training distribution.
Since all levels follow the same underlying structure and are fully observable, it is theoretically possible to learn a policy that can perform optimally on all levels inside the distribution.

To select levels to train on, we use SFL (Rutherford et al., ), a state-of-the-art UED algorithm that regularly performs a large number of rollouts on randomly generated levels. It then selects a subset of these with high learnability and trains on them for a fixed duration before again selecting new levels.
The main limitation of SFL, that it is only applicable to deterministic settings with a binary success reward, does not constrain us, as Kinetix satisfies both of these assumptions.
We ran preliminary experiments using PLR (Jiang et al., ; ) and ACCEL (Parker-Holder et al., ), but found that these approaches provided no improvements over DR (see Appendix L).

For all experiments, we use PPO (Schulman et al., ) with multi-discrete actions.
We allot each method 5 billion environment interactions and periodically evaluate performance on the holdout levels. Hyperparameters are detailed in Appendix H.

### 4.1 Architecture

Figure: 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.
Refer to caption: extracted/5966380/images/architecture3.png

The architecture we use is summarised in Figure 2.
To process the observation in a permutation-invariant way, we represent each entity as a vector $v$, containing information about its physical properties, such as friction, mass and rotation.
We separately encode (using a set of small feedforward networks) polygons, circles, joints and thrusters into initial embeddings $x^{T}_{i}$, where $T\in\{p,c,j,t\}$.
We perform self-attention (Bahdanau et al., ; Vaswani et al., ) over the set of shapes (i.e., polygons and circles) without positional embeddings to obtain new shape embeddings $\tilde{x}^{S}_{i}$. To incorporate joint information, we take each joint feature $x^{j}_{i}$, and its two connected shapes $\tilde{x}^{T}_{\text{from}}$ and $\tilde{x}^{S}_{\text{to}}$, and pass the concatenation through a feedforward network $f$, and add it to the embedding for $\tilde{x}^{S}_{\text{from}}$. We have two feature vectors for each joint, with the from and to shape swapped. This layer is reminiscent of message passing in graph neural networks (Gilmer et al., ; Bronstein et al., ).
Similarly, for each thruster $x^{t}_{i}$ and associated shape $\tilde{x}^{S}_{o}$, we process these using a message-passing layer and add the result back to $\tilde{x}^{S}_{o}$.
This entire process constitutes one transformer layer, which we apply multiple times. We use multi-headed attention, with a different attention mask for each head.
The first mask represents a fully-connected graph and contains all shapes; the second allows shapes to attend to those that are connected by a joint (Sferrazza et al., ; Buterez et al., ); the third allows attention to shapes that are joined by any n-step connection; and the final mask allows shapes to attend to those that they are currently colliding with.
Finally, following Parisotto et al. (), we use a gated transformer, and perform layernorm (Lei Ba et al., ) before the attention block.

## 5 Zero-Shot Results

Figure: 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.
Refer to caption: x5.png

In Figure 3, we run SFL on the S, M and L environment sizes, respectively (see Appendix J for a per-level breakdown). In each case, the (random) training environments are of the corresponding size, and we use the corresponding holdout set (see Appendix E for a full listing) to evaluate the agent’s generalisation capabilities.
We see that, in every case, the agent’s performance increases throughout training, indicating that it is learning a general policy that it can apply to unseen environments.
For S, the agent very quickly learns a policy superior to the random policy, and is able to solve most of the hold out levels zero-shot.
While the solve rate is lower on M, the agent can still zero-shot a number of unseen hand-designed environments.
On the L environments, in which the agent is assessed on the most challenging holdout tasks, we see a very slow, and non-monotonic, performance increase. As well as being trained and tested on more complex levels, it seems that as the complexity increases, randomly generated levels are more likely to be unsolvable, reducing the proportion of useful data the agent can learn on.
Overall, this result demonstrates that training an agent on a large set of mixed-quality levels can lead to general behaviour on unseen tasks.

### 5.1 Analysis: Zero-shot Locomotion of an Arbitrary Morphology

Figure: 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.
Refer to caption: x6.png

In this section, we take a closer look at the zero-shot capabilities of the learned general agent by probing its behaviour in a constrained goal-following setup. Specifically, we create levels with a single morphology (a set of shapes connected with motors and containing the green shape) in the centre of the level, with a goal (the blue shape) fixed at the top of the level with a random $x$ position. Since the goal is made to be unreachable, the optimal behaviour of the agent is to maximise the dense auxiliary reward and move as close as possible to the goal (i.e., directly underneath it). We evaluate three hand-designed morphologies: Car, Snake and Thruster, as well as Morphology-Random, which selects from one of 2000 randomly generated 3-shape morphologies (Appendix F).

We measure how the $x$ position of the goal correlates with the $x$ position of the controllable morphology (Figure 4). The behaviour of an optimal agent would manifest itself as a high correlation and would therefore show high incidence along the diagonal. We evaluate both a random agent and a general agent trained on random M levels for 5 billion timesteps. Each plot is aggregated over 2000 randomly sampled levels, each of which is run for 64 timesteps to allow the agent to maneuver into position and then run for a further 64 timesteps for data collection.

As would be expected, the random agent shows no correlation between the position of the controllable morphology and the goal. By contrast, the trained agent shows positive correlation, indicating it is able to maneuver the morphology towards the goal location. We see a variety of outcomes across the different morphologies, with the agent showing very strong results on Car and Thrust, with a slightly weaker performance on Snake. When evaluating on Morphology-Random, we do see some positive correlation, although not as strong as the hand-designed levels.

The positive results on these constrained ‘goal-conditioned’ environments show that the agent has indeed learned a general policy that encompasses purposeful locomotion of an arbitrary morphology.

## 6 Fine-Tuning Results

In this section we leave the zero-shot paradigm and investigate the performance of the general agent when given a limited number of samples to fine-tune on the holdout tasks.
In particular, in Figure 5 we train a separate specialist agent for each level in the L holdout set, and compare this to fine-tuning a general agent (the same one used for Section 5.1). We plot the learning curves for four selected environments, as well as the aggregate performance over the entire holdout set.
On three of these levels, fine-tuning the agent drastically outperforms training from scratch. In particular, for Mujoco-Hopper-Hard and Mujoco-Walker-Hard, the fine-tuned agent is able to competently complete these levels, whereas the tabula rasa agent cannot do so consistently. Notably, this is despite the fact that the pre-trained agent cannot solve these environments zero-shot.
While the general trend is that fine-tuning beats training from scratch, we do see one case: Thruster-Large-Obstacles, where fine-tuning learns slower.

Figure: 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.
Refer to caption: x7.png

### 6.1 Analysis: General Pretraining can beat Training on the Target Task

We now further investigate the case of Car-Ramp (Figure 6(a)) where RL, even with a large sample budget, fails to solve but that our fine-tuned general agent can complete (note that this behaviour is also shown in MuJoCo-Walker-Hard). Car-Ramp is an example of a deceptive problem (Goldberg, ; Liepins & Vose, ; Lehman & Stanley, ) that requires the agent to first move away from the goal (and incur a negative reward) to obtain enough momentum to jump the gap.

An agent trained tabula rasa with PPO for 1 billion timesteps fails to reach the target a single time.
By contrast, our general agent (which has never seen the task before) solves it zero-shot around 5% of the time. This proves to be enough traction that, with a small amount of fine-tuning, the agent can reliably solve this task (Figure 6(b)).
We do stress that, while impressive, this behaviour is the exception rather than the rule, only occurring on 2 of 74 handmade levels. We see this as a promising sign for a trained general agent in Kinetix to serve as a strong base model.

Figure: (a) Initial State
Refer to caption: x8.png

## 7 Related Work

Hardware-Accelerated Physics Engines
Jax2D joins a thriving ecosystem of hardware-accelerated physics engines used in RL tasks. Brax (Freeman et al., ), MJX (Todorov et al., ) and Isaac-Gym (Makoviychuk et al., ) have all been been widely used in the RL community, particularly for robotics tasks. While superficially similar, we believe Jax2D is useful for an entirely different set of problems. Firstly, Jax2D only operates in two dimensions, so training on robotics tasks for transfer to the real world is not a goal of the engine. Jax2D instead aims to be able to represent a hugely diverse range of physics problems and, most crucially, can do so with the same computation graph, allowing work across multiple heterogeneous environments to be parallelised.

Physical Reasoning
PHYRE (Bakhtin et al., ) also uses 2D rigid-body physics by tasking agents with placing a ball to achieve some goal state. Li et al. () extend this bandit-like problem, allowing the agent to take actions throughout the episode. A crucial difference is that we train on a large automatically generated set of tasks rather than a small set of handmade ones.

Hardware-Accelerated RL
Our work follows the recent trend of using hardware-accelerated RL environments to run significantly larger-scale experiments than would be possible with CPU-based environments (Lu et al., ; Jackson et al., ; ; Goldie et al., ; Rutherford et al., ; Nikulin et al., ; Kazemkhani et al., ). By leveraging Kinetix’s speed, we can train for billions of timesteps and, as we show, general capability does only emerge after such a long time.

Generalist Robotics Agents Recent work has strived to learn a generalist foundation model for robotics (Reed et al., ; Bousmalis et al., ; Team et al., ; Nasiriany et al., ; O’Neill et al., ). While most of these approaches perform behaviour cloning on a large dataset from a variety of robot morphologies and tasks, Nasiriany et al. () develop a large-scale simulation environment, with an initial focus on kitchen environments.
By contrast, Kinetix aims to train an online agent tabula rasa, without using external data, and further has a large variety of different tasks.

Open-Ended Learning
Kinetix also ties into the paradigm of open-ended learning (Soros & Stanley, ; Stanley, ; Sigaud et al., ; Hughes et al., ), in which a system continually generates new and novel artifacts. In the context of RL, this often means training inside a large, diverse distribution and having some method to adapt this distribution over time, with fields such as UED specifically focusing on the latter. While these methods hold the promise of generating novel and useful levels in an open-ended manner, the environments used in their experiments are often very constrained in what they can represent (Wang et al., ; Dennis et al., ; Wang et al., ; Jiang et al., ; ; Parker-Holder et al., ). As we have shown, in a significantly more diverse task space, these methods tend to fail.

A recent work with a similar vision to Kinetix is Autoverse (Earle & Togelius, ), where an agent acts inside a cellular automata (CA) based gridworld. By changing the underlying rules of the CA, a large diversity of different levels can be represented in this system.
Relatedly, Sun et al. () use prior knowledge in the form of Large Language Models to generate code for video games and robotic simulations, to train RL agents on.
Powderworld (Frans & Isola, ) instead creates a expressive environment based on different types of elements interacting in a sandbox environment. Other notable work that aims to use open-ended discovery to train generalist agents include Voyager (Wang et al., ), Jarvis-1 (Wang et al., ) and Optimus-1 (Li et al., ). These are more focused on long-horizon planning, the self discovery of new tasks to perform, and use Minecraft as their domain with prior knowledge in the form of a large language model.

Perhaps the works most similar to ours is the highly impactful XLand line of research (Team et al., ; ). XLand defines a large and diverse distribution of levels inside a 3D physics simulation, with an embodied agent (or set of agents) required to fulfill some specified goal in the environment. Similar to us, XLand agents train on a procedurally generated set of levels and are then assessed on human-designed holdout levels. We see the main differences to Kinetix being the expressivity of the tasks and the public state of the work. In particular, we would subjectively claim that Kinetix, through representation of almost any conceivable 2D rigid-body physics problem, has a more expressive universe of tasks. While XLand also employs a physics engine, all the tasks are constrained to homogeneous agents acting in the world, potentially limiting the scope of tasks—it is not clear, for instance, how one would represent any of the holdout environments in Figure 1 in XLand. Lastly, we note that the XLand environment has not been made open-source, limiting its use for future research.
XLand-Minigrid (Nikulin et al., ) provides a fast, open-source version of XLand, however in doing so it simplifies the environment into a gridworld.

## 8 Discussion and Future Work

We believe Kinetix is a uniquely diverse, fast and open-ended environment, placing it well as a foundation to study open-ended RL, including large-scale online pre-training for general RL agents.
In stark contrast to many other benchmarks used for open-ended learning (Wang et al., ; ; Parker-Holder et al., ; Chevalier-Boisvert et al., ; Rutherford et al., ), Kinetix represents a large space of semantically diverse tasks, instead of just variations on a single task. This presents a challenge for future environment design research that can intelligently generate levels (Dennis et al., ), rather than just filtering from a predefined distribution.
We also believe Kinetix is an excellent framework for investigating issues in agent training like network capacity (Obando-Ceron et al., ), plasticity loss (Igl et al., ; Berariu et al., ; Dohare et al., ; Sokar et al., ), lifelong learning (Kirkpatrick et al., ) and multi-task learning (Sodhani et al., ; Hafner, ; Benjamins et al., ).

Requiring billions of online environment interactions is impractical for real-world applications. However, we see three primary ways to leverage the cheap samples of simulations for sample-constrained tasks. One approach is to meta-learn parts of the RL process, for instance the algorithm (Oh et al., ; Lu et al., ; Jackson et al., ), optimiser (Goldie et al., ) or loss function (Bechtle et al., ). Alternatively, the emerging capabilities of large world models (Bruce et al., ; Valevski et al., ) hint at a new paradigm of online training entirely in imagination (Ha & Schmidhuber, ; Yu et al., ; Hafner et al., ; ; ), where the only bottleneck to environment samples is compute. Finally, we may find that, with enough scale, we can fine-tune an agent trained in simulation on real world tasks.

## 9 Conclusion

In this work, we first introduce Jax2D, a hardware-accelerated 2D physics engine. Using Jax2D, we build Kinetix, a vast and open-ended physics-based RL environment. We illustrate the diversity of Kinetix by hand-designing a comprehensive holdout set of environments that test various skills, such as navigation, planning and physical reasoning. We train an agent on billions of environment interactions from randomly generated tasks, and show that it can zero-shot generalise to many human-designed tasks, as well as function as a strong base model for fine-tuning.
We hope that this work can serve as a foundation for future research in open-endedness, large-scale online pre-training of general RL agents and unsupervised environment design.

## Appendix

We structure the appendix as follows:

- •
Appendix A describes the mathematical and computational logic behind Jax2D and Appendix B performs speed tests on it.
- •
Appendix C provides further details of the Kinetix RL environment, while Appendix D shows examples of randomly generated levels.
- •
Appendix E lists the hand-designed holdout levels and Appendix F shows example morphologies used in Figure 4.
- •
Appendix G describes the different network architectures in further detail and Appendix H lists the associated hyperparameters used.
- •
Appendix I investigates training agents directly on the holdout levels.
- •
Appendix J provides a de-aggregated view of the main generalist agent results, split out by every environment.
- •
Appendix K provides additional generalist agent results, while Appendix L compares UED methods.
- •
Appendix M performs a small ablations study where we try removing aspects of our general agent training pipeline.

## Appendix A Jax2D

This section provides an in-depth look into the logic behind Jax2D. Jax2D largely owes its heritage to Box2D (Catto, ) and ImpulseEngine (Gaul, ), with most of the underlying framework being lifted from these engines and adapted for JAX. For a more thorough account of some of the concepts behind rigid-body physics, we recommend Erin Catto’s talks.(^3^33[https://box2d.org/publications/](https://box2d.org/publications/))

### A.1 Core Engine

The main loop of Jax2D is summarised in LABEL:alg:jax2d. Each part of the engine is subsequently explained as referenced.

### A.2 Impulse Resolution and Constraint Solving

The core of Jax2D is impulse resolution, in which an equal and opposite impulse is applied to a pair of shapes in order to satisfy some constraint. For a given impulse $\bm{j}$, the positional and angular velocities of a shape are affected as follows.

$$ $\displaystyle\bm{v}$ $\displaystyle\leftarrow\bm{v}+\frac{\bm{j}}{m}$ (1) $\displaystyle\omega$ $\displaystyle\leftarrow\omega+\frac{\bm{r}\times\bm{j}}{I}$ (2) $$

where $v$ is positional velocity, $m$ is mass, $\omega$ is angular velocity, $r$ is the displacement from the centre of mass of the shape to the position the impulse is being applied at and $I$ is the rotational inertia.

We use $\times$ to represent either the scalar-vector or the vector-vector cross product (the choice should be inferable from the operands).

### A.3 Collisions

The first type of constraint we consider is the collision constraint, which prevents objects from moving inside of each other.

#### A.3.1 Collision Manifolds

The notion of a collision between to shapes is reduced to the concept of a collision manifold, containing the information shown in Table 1.

**Table 1: Collision Manifold Specification**
| Attribute | Symbol | Data Type | Description |
| --- | --- | --- | --- |
| Position | $\bm{p}$ | [float, float] | Global position of the collision. |
| Normal | $\bm{\hat{n}}$ | [float, float] | Normalised vector along which the collision occurs. |
| Penetration | $p$ | float | Positive penetration indicates an active collision. |

The resolution of a collision takes place in two steps. First a collision manifold is generated. This is dependent on the exact shapes that are colliding (e.g. the logic for deriving a collision manifold between two circles is different than for two polygons). Once the collision manifold is generated, the exact nature of the colliding shapes are no longer relevant and only their common attributes (mass, inertia, etc.) are used for the subsequent collision resolution. In this way, while the generation of the collision manifolds is heterogeneous, the resolution of these occurs homogeneously.

#### A.3.2 Circle-Circle Collision Manifolds

Generating a collision manifold between two circles is relatively simple, and is calculated as follows:

$$ $\displaystyle\bm{p}$ $\displaystyle\leftarrow\bm{p_{a}}+r_{a}\cdot\bm{\hat{n}}$ (3) $\displaystyle\bm{\hat{n}}$ $\displaystyle\leftarrow\frac{\bm{p_{b}}-\bm{p_{a}}}{|\bm{p_{b}}-\bm{p_{a}}|}$ (4) $\displaystyle p$ $\displaystyle\leftarrow r_{a}+r_{b}-|\bm{p_{b}}-\bm{p_{a}}|$ (5) $$

#### A.3.3 Polygon-Circle Collision Manifolds

The collision between a polygon $a$ and a circle $b$ is calculated by first determining the closest point on any edge to the circle. For each edge, the centre of the circle is clipped to perpendicular lines extending from both corners, before being projected onto the edge to find the closest point for that particular edge. The clipping ensures that the point doesn’t end up off the end of an edge - it will instead be clipped to a corner. Once this closest point $\bm{p}$ has been found, the collision manifold can be calculated.

$$ $\displaystyle\bm{\hat{n}}$ $\displaystyle\leftarrow\frac{\bm{p_{b}}-\bm{p}}{|\bm{p_{b}}-\bm{p}|}$ (6) $\displaystyle p$ $\displaystyle\leftarrow r_{b}-|\bm{p}|$ (7) $$

#### A.3.4 Polygon-Polygon Collision Manifolds

Collisions between two convex polygons are the most complex. The underlying stratgey is defined by the separating axis theorem: any two convex polygons that are not colliding will have an axis upon which, when the vertices of both shapes are projected onto, there will be no overlap. Furthermore, it can be shown that if this axis exists, it must run perpendicular to one of the edges of one of the polygons. Intuitively, one can imagine drawing a straight line (perpendicular to the separating axis and thus parallel with an edge) that separates the two convex polygons.

If there is no separating axis then the two polygons are colliding. Finding the point of collision involves pinpointing the axis of least penetration, that is the axis that when projected upon causes the least amount of overlap. The face that the axis of least penetration is derived from is termed the reference face, and the face (on the other shape) of which the corners have the least penetration is termed the incident face. Similar to the polygon-circle collision, the incident face is then clipped to the boundaries of the reference face. Each of the (clipped) vertices of the incident face can then produce their own collision manifolds (if they are indeed penetrating the reference face). The normal of the collision is that of the reference face and the penetration can be easily calculated by projecting the clipped incident face onto this normal.

The decision to (sometimes) produce two collision manifolds for polygon-polygon collisions is one of stability. When two edges rest on each other a single collision manifold will cause the polygon to oscillate as the collision manifold flips from side to side.

#### A.3.5 Collision Resolution

Once a collision manifold has been created, it is then turned into an impulse that affects the two shapes. When two objects are deemed to have collided (i.e. a collision manifold with positive penetration is found), the collision constraint specifies that the new relative velocity at the point of collision should be equal to $-e\bm{v_{r}}$, where $e$ is the restitution of the collision and $\bm{v_{r}}$ is the relative velocity at the point of collision. If $e=0$ we see an inelastic collision where the collision points on both shapes should have zero relative velocity. Conversely, if $e=1$ we would see a perfectly elastic collision and the conservation of kinetic energy.

We first note that the velocity of a point on an object can be calculated by

$$ $\displaystyle\bm{v_{r}}=\bm{v}+\omega\times\bm{r}$ (8) $$

where $\bm{v}$ is the velocity of the objects centre of mass, $\omega$ is the angular velocity and $r$ is the point on the object relative to the centre of mass. Given this, we can derive the required impulse to resolve a collision between objects $a$ and $b$ is

$$ $\displaystyle\bm{j_{n}}=\frac{-(1+e)(\bm{\hat{n}}\cdot(\bm{v_{a}}+(\omega_{a} \times\bm{r_{a}})-\bm{v_{b}}-(\omega_{b}\times\bm{r_{b}})))}{m_{a}^{-1}+m_{b}^ {-1}+\frac{(\bm{r_{a}}\times\bm{\hat{n}})^{2}}{I_{a}}+\frac{(\bm{r_{b}}\times \bm{\hat{n}})^{2}}{I_{b}}}\cdot\bm{\hat{n}}$ (9) $$

where $e$ is the restitution, $\bm{\hat{n}}$ is the collision normal, $\bm{v_{a}}$ and $\bm{v_{b}}$ are the respective positional velocities, $\omega_{a}$ and $\omega_{b}$ are the respective angular velocities, $\bm{r_{a}}$ and $\bm{r_{b}}$ are the respective relative positions of the collision from the centre of masses, $m_{a}$ and $m_{b}$ are the respective masses and $I_{a}$ and $I_{b}$ are the respective rotational inertias.

Intuitively, the numerator represents the change in speed we wish to occur between the collision points along the axis of the collision normal. The denominator then scales this value by the mass and inertia of the colliding objects so that the resultant impulse will cause this change in speed.

In Jax2D every shape has an associated restitution, with the restitution of a collision defined as the minimum of the restitutions of the colliding shapes $e=\min(e_{a},e_{b})$.

#### A.3.6 Friction in Collisions

As well as the collision impulse which acts along the collision normal, we calculate a friction impulse which acts perpendicular to it against the relative movement. This follows Couloumb’s Law:

$$ $\displaystyle|\bm{j_{f}}|\leq\mu|\bm{j_{n}}|$ (10) $$

where $\bm{j_{f}}$ is the friction impulse, $\bm{j_{n}}$ is the normal impulse and $\mu$ is the coefficient of friction. $\bm{j_{f}}$ is therefore defined, similarly to Equation 12, as

$$ $\displaystyle\bm{j_{f}}=\text{clip}\left(\frac{-(\bm{\hat{t}}\cdot(\bm{v_{a}}+ (\omega_{a}\times\bm{r_{a}})-\bm{v_{b}}-(\omega_{b}\times\bm{r_{b}})))}{m_{a}^ {-1}+m_{b}^{-1}+\frac{(\bm{r_{a}}\times\bm{\hat{t}})^{2}}{I_{a}}+\frac{(\bm{r_ {b}}\times\bm{\hat{t}})^{2}}{I_{b}}},-\mu|\bm{j_{n}}|,\mu|\bm{j_{n}}|\right) \cdot\bm{\hat{t}}$ (11) $$

where $\bm{\hat{t}}$ is the normalised vector perpendicular to the normal of the collision.

Similar to restitution, every shape has its own coefficient of friction, with the coefficient for a collision defined as $\mu=\sqrt{\mu_{a}^{2}+\mu_{b}^{2}}$.

#### A.3.7 Positional and Velocity Corrections

In a simulation of infinite temporal granularity, impulses would be enough to guarantee reliable behaviour. However, since in practice we must quantise our simulation into discrete timesteps, only using impulses to solve constraints causes compounding errors to emerge in the simulation. In the case of collision constraints, this manifests itself as resting objects slowly sinking into each other.

To deal with this, we first introduce a velocity correction. Decomposing Equation 12 we can see that the numerator defines the change in speed that will occur along the collision normal between the two collision points. Since our velocity correction will also operate along the collision normal, we can simply add the desired speed bias to the numerator. We calculate this bias as $\alpha p$ where $p$ is the penetration and $\alpha$ is a coefficient in units of inverse time. Since this bias a function of the penetration, it will prevent bodies from sinking into each other, even if they have low velocity. It should be noted that this practice introduces some ‘bounce’ into the simulation, which can in effect slightly increase the restitution of collisions.

We also introduce a positional correction, which directly moves colliding shapes when they overlap. We similarly define this as $\beta p$, where $\beta$ is a unitless coefficient.

### A.4 Joints

As well as collision constraints, Jax2D also represents the concept of joint constraints. These in their most basic form fix two relative points on two separate objects together such that they must always occupy the same global position. It should be noted that (assuming the relative positions are inside the shapes), this is directly at odds with the collision constraint. Therefore, when we connect two shapes with a joint, we disable their respective collision constraint.

#### A.4.1 Revolute Joints

The most basic type of joint constraint is the revolute joint. This simply specifies that the two positions on each of the shape occupy the same position and have zero relative velocity to each other. Note that they are allowed to have non-zero relative angular velocity, which allows the shapes to spin around the joint (hence revolute).

This is achieved in effect by applying a constant collision with no restitution at the point of joining, with the collision normal pushing the joined positions back towards each other. As with collisions, we also apply velocity and positional corrections.

#### A.4.2 Fixed Joints

Jax2D also faciliates a ‘fixed’ joint, in which an additional rotational constraint enforces that the relative angle between two shapes remains constant, fixing them together effectively into a single rigid body.

The rotational constraint applies an angular impulse around the fixed joint, defined as

$$ $\displaystyle j_{r}=\frac{\omega_{a}-\omega_{b}}{I_{a}^{-1}+I_{b}^{-1}}$ (12) $$

This will cause the relative angular velocity of the two shapes to become zero.

We also apply corrections directly to the angular velocities defined as $\gamma(\theta_{a}-\theta_{b}-\theta_{f})$, where $\theta_{a}$ and $\theta_{b}$ are the respective rotations of the two shapes, $\theta_{f}$ is the target rotation at which they have been fixed at and $\gamma$ is a coefficient in units of inverse time. This is analogous to the velocity correction, with the angular difference from the target taking the place of the penetration.

#### A.4.3 Joint Limits

In order to allow for Jax2D to represent environments like the MuJoCo inspired tasks, revolute joints can have rotational limits applied to them, meaning they can only rotate within a given range. When the relative rotation between two shapes connected with a limited revoloute joint exceeds either the minimum or maximum rotation, an angular impulse is applied to correct this. This is applied similarly to that for a fixed joint, except that the angular velocity correction is not applied if the relative angular velocity of the two shapes is already bringing them back into within their limits. This is to allow motors to push joints back within limits potentially faster than the angular velocity correction would do.

### A.5 Motors

A revolute joint can have a motor attached to it, which can apply a torque around the joint. Each motor has a target angular velocity and a strength to which it will apply a torque to achieve it. For stability, as the angular velocity approaches the target, the motor applies less torque. If the angular velocity exceeds the motors target then it will apply a torque in the opposite direction. The applied angular impulse is calculated as

$$ $\displaystyle j_{r}=p\cdot\tanh{((\omega_{a}-\omega_{b}-s\cdot A)\cdot\rho)}$ (13) $$

where $s$ is the target speed of the motor, $A$ is the action being applied on the motor (by a human or artificial agent), $p$ is the motor power and $\rho$ is a coefficient to control to what degree the power wanes as it approaches the target angular velocity.

It should be noted that the angular impulse applied by a motor is not a constraint to be solved but a true impulse being applied to the scene, similar to gravity. For this reason it is applied once, before the main constraint solving loop.

### A.6 Thrusters

Thrusters can be attached to shapes and can apply a force in the direction they are facing. The force applied is defined as $p\cdot A$, where $p$ is the power of the thruster and $A$ is the action taken on the thruster. As with motors, the thruster impulse is applied before constraint solving begins.

### A.7 Impulse Accumulation and Warm Starting

For a stable simulation, we simulate multiple solver steps for every simulation timestep. This is because solving one pairwise constraint can often affect other constraints. For instance, imagine a stack of rectangles resting on top of each other – solving the collision constraint of the bottom rectangle with the floor might push this rectangle further into the one above it (especially with the velocity and positional corrections). This same problem would then propagate its way up the entire stack (and back down again), necessitating multiple solver steps for stability (each solver step iteratively solves each constraint).

One interesting observation to make is that solver steps from previous timesteps can provide useful information for the current timestep. In particular, the aggregate impulse applied at each manifold last timestep serves as good ‘first guess’ for the impulse to apply at the current timestep, especially when bodies are mostly static. In this way, we can effectively solve constraints not only over multiple solver steps but also over multiple timesteps, with little extra cost. This technique is referred to as ‘warm starting’.

Warm starting requires us to record accumulated impulses throughout the solver steps and also to match collision manifolds across timesteps. Jax2D takes the simple approach of naïvely matching collision manifolds across adjacent timesteps – if a collision does not occur between two bodies on a timestep then all accumulated impulses are wiped. Jax2D by default warm starts collisions, joint positional constraints and fixed joint rotational constraints. Efforts to apply warm starting to the joint limits of revolute joints caused instability.

### A.8 Parallelised Computation and Batched Impulse Resolution

As well as being able to easily parallelise multiple Jax2D environments with the Jax vmap operation, we also parallelise many of the calculations within a single environment, providing further speed increases. The calculation of collision manifolds is easily parallelised, as they have no side effects. The application of motors and thrusters is also parallelised. A more nuanced parallelisation is the constraint solving.

As discussed in Section A.7, solving one constraint can affect (and even unsolve) other constraints. For this reason, solving constraints sequentially provides a greater efficiency in terms of solver steps, as each constraint can in effect take into account the effects of already solved constraints. In testing, we found that fully parallelising constraint solving did indeed noticeably reduce the stability of the simulation.

Due to the way the vmap operation works, everything in the parallelised function must run the same compute graph – there can be no branching. For us, this means that every collision constraint between every pair of shapes must be solved every solver step, as we can’t know a priori which shapes will collide. This means that, in most cases, the vast majority of computed collision resolutions are inactive.

We want to parallelise collision constraints for speed reasons, but it makes the solution unstable, however we also find that the majority of collision constraints are actually inactive. This naturally leads to the solution of partially parallelising the collision constraints by solving them in batches, which we vmap across. By spreading out the active collision manifolds across as many batches as possible, we gain the speed advantages of parallelisation without the negative effects on stability (except in the cases where many shapes are colliding with each other). The solver batch size therefore also arises as a tuneable parameter that trades off between simulation speed and accuracy. We use a value of 16 by default.

We do not parallelise joint constraint solving, as there are far less joints than possible collisions (as collisions grows quadratically with the number of shapes), so the potential for speed improvements is significantly less.

## Appendix B Jax2D Speed Results

Here we investigate the runtime speed of both Jax2D and Kinetix. For all comparisons we use a single NVIDIA L40S GPU, on a server with two AMD EPYC 9554 64-Core CPUs. We first compare Jax2D against Box2D (Catto, ). We implement environments in Box2D and Jax2D that are comparable in size (notably, the Box2D environment has three polygons and two joints, whereas the Jax2D environment uses the M size, with 6 polygons, 3 circles, 2 joints and 2 thrusters).
We then use two different approaches of comparing speed: The first is by simply running the engines, and applying fixed actions, giving us a raw speed measure of each engine. In the second approach, we compare speed when running the RL training loop, to have a more realistic estimate for speed during training. We use PureJaxRL-style training for Jax2D (Lu et al., ) and Stable Baselines 3 (Raffin et al., ) for Box2D.
We use the flattened symbolic representation for Jax2D and use comparably-sized networks for both Box2D and Jax2D.

The results are presented in Figures 7 and 2. First, inside an RL loop, Jax2D always outperforms Box2D, and shows improved scaling once the number of parallel processes greatly exceeds the number of physical CPU cores. When comparing just the engine, Box2D outperforms Jax2D when using fewer than 1024 environments, at which point Jax2D overtakes Box2D.

Figure: Figure 7: Comparing Box2D vs Jax2D’s speed in two scenarios. The first, on the left, includes RL training, whereas the rightmost plot corresponds to raw engine performance.
Refer to caption: x10.png

**Table 2: The best-case steps per second for both Jax2D and Box2D, in an RL loop and outside. In raw performance, Jax2D’s best case is approximately $4.5\times$ faster than Box2D, and this increases to more than $30\times$ inside an RL training pipeline.**
| Approach | Steps Per Second (Best case) | Environment Workers (Best Case) |
| --- | --- | --- |
| Jax2D (RL) | 824K | 32768 |
| Jax2D (Engine Only) | 9049K | 16384 |
| Box2D (RL) | 24K | 32768 |
| Box2D (Engine Only) | 1982K | 8192 |

In Figure 8, we compare the three different level sizes in Kinetix (S, M and L), as well as the different observation spaces.
Speed predictably decreases as we increase the environment size. Using the pixel-based observation requires more memory, so we cannot run as many parallel environments as with the other observation spaces. Symbolic-Entity does not scale as well as Symbolic-Flat, likely due to saturating memory bandwidth.

Figure: Figure 8: The number of steps per second (SPS) in Kinetix for a variety of observation spaces. Symbolic-Entity is what we use in our experiments, while Symbolic-Flat is a flattened (and therefore not permutation invariant) representation.
Refer to caption: x11.png

## Appendix C Kinetix : Further Details

### C.1 Environment Class Sizes

The environment sizes we use are detailed in Table 3. Note that every level in Kinetix contains 4 large fixated polygons (floor, ceiling, left wall, right wall).

**Table 3: The size of each environment class.**
| Entity | Small | Medium | Large |
| --- | --- | --- | --- |
| Polygons | 5 | 6 | 12 |
| Circles | 2 | 3 | 4 |
| Joints | 1 | 2 | 6 |
| Thrusters | 1 | 2 | 2 |
| Thruster Joint | 4 | 4 | 4 |
| Thruster Bindings | 2 | 2 | 2 |

### C.2 Observation Spaces

Kinetix allows for three observation spaces: Symbolic-Entity, Symbolic-Flat and Pixels. Both the symbolic observations use a common representation for shapes Table 4, joints Table 6 and thrusters Table 6

For use in Symbolic-Entity, we construct 2 entities per joint: a to and from version of each joint. Given two shapes, we first set one as the from shape and the second as the to shape to construct the first feature vector for this joint. The second feature vector is obtained by the same process, just with from and to swapped. This allows each joint to affect both its attached shapes in the message passing layer.

**Table 4: Information provided for shapes**
| Name | Dimensions |
| --- | --- |
| Position | 2 |
| Velocity | 2 |
| Inverse Mass | 1 |
| Inverse Inertia | 1 |
| Density | 1 |
| $\tanh($Angular Velocity$/10)$ | 1 |
| $\texttt{OneHot}(\text{Role})$ | $n_{\text{roles}}$ |
| $\sin(\text{Rotation})$ | 1 |
| $\cos(\text{Rotation})$ | 1 |
| Friction | 1 |
| Restitution | 1 |
| $\texttt{OneHot}(\text{ShapeType})$ | $n_{\text{types}}$ |
| Radius (only for circle) | 1 |
| Vertices (only for polygons) | 8 |
| TriangleOrRectangle (only for polygons) | 2 |

**Table 5: Information provided for thrusters**
| Name | Dimensions |
| --- | --- |
| Active | 1 |
| Relative Position | 2 |
| Power | 1 |
| $\sin(\text{Rotation})$ | 1 |
| $\cos(\text{Rotation})$ | 1 |

## Appendix D Randomly Generated Levels

We show 24 example random levels for size S (Figure 9), M (Figure 10) and L (Figure 11).

Figure: Figure 9: Randomly generated filtered levels from the DR distribution (S).
Refer to caption: x12.png

Figure: Figure 10: Randomly generated filtered levels from the DR distribution (M).
Refer to caption: x13.png

Figure: Figure 11: Randomly generated filtered levels from the DR distribution (L).
Refer to caption: x14.png

## Appendix E Hand-Designed Levels Listing

In this section we provide plots of the handmade levels. Figures 12, 13 and 14 contain the full holdout sets for each environment size, respectively. We note that a darker colour indicates that a shape is fixated, i.e., that it has infinite mass and cannot move.

Figure: Figure 12: Handmade levels (S).
Refer to caption: x15.png

Figure: Figure 13: Handmade levels (M).
Refer to caption: x16.png

Figure: Figure 14: Handmade levels (L).
Refer to caption: x17.png

## Appendix F Randomly Generated 3-Shape Morphologies

Figure 15 shows a sample of the randomly-generated morphologies used for the analysis in Section 5.1.

Figure: Figure 15: Randomly generated 3-shape morphologies.
Refer to caption: x18.png

## Appendix G Further Network Architecture Details

We use the same actor-critic architecture for each observation space, consisting of five fully connected layers, of width 128, and a tanh activation. However, how the input to this network is obtained differs for each observation space. Since the environment is fully observable (except in the case of Pixels), we do not use a recurrent network.

Pixels
Inspired by the IMPALA architecture (Espeholt et al., ), we use two convolutional layers to process the $125\times 125$ observation. The first has 16 channels, a size of $8\times 8$ and a stride of $4\times 4$ while the second has 32 channels, a size of $4\times 4$ and a stride of $2\times 2$. The result of these layers is flattened before being passed to the main actor-critic network.

Symbolic-Flat
The Symbolic-Flat encoder is simply a feed forward network with width of 512.

## Appendix H Hyperparameters

Table 7 contains a listing of the hyperparameters we use for experimentation.

**Table 7: Learning Hyperparameters.**
| Parameter | Value |
| --- | --- |
| Env |  |
| Frame Skip | 2 |
| PPO |  |
| $\gamma$ | 0.995 |
| $\lambda_{\text{GAE}}$ | 0.9 |
| PPO number of steps | 256 |
| PPO epochs | 8 |
| PPO minibatches per epoch | 32 |
| PPO clip range | 0.02 |
| PPO # parallel environments | 2048 |
| Adam learning rate | 5e-5 |
| Anneal LR | no |
| PPO max gradient norm | 0.5 |
| PPO value clipping | yes |
| return normalisation | no |
| value loss coefficient | 0.5 |
| entropy coefficient | 0.01 |
| Model |  |
| Fully-connected dimension size | 128 |
| Fully-connected layers | 5 |
| Transformer layers | 2 |
| Transformer Encoder Size | 128 |
| Transformer Size | 16 |
| Number of heads | 8 |
| SFL |  |
| Batch Size $N$ | 12288 |
| Rollout Length $L$ | 512 |
| Update Period $T$ | 128 |
| Buffer Size $K$ | 1024 |
| Sample Ratio $\rho$ | 0.5 |

## Appendix I Specialist Results

In this section, we investigate the performance of agents directly trained on the holdout levels. We consider two paradigms here: An agent trained on tabula rasa, and one fine-tuned from a general agent. The results in this section are a different way to present the findings in Section 6, as well as including results for S and M.
In Figures 16, 17 and 18, we plot the performance of the agents trained for Figure 5 on each individual holdout level. We note that the fine-tuning base model is one trained on M for 5B timesteps.

Figure: Figure 16: Specialist Agents on S.
Refer to caption: x19.png

Figure: Figure 17: Specialist Agents on M.
Refer to caption: x20.png

Figure: Figure 18: Specialist Agents on L.
Refer to caption: x21.png

## Appendix J General Agent Results by Holdout Level

Next, we plot the performance of SFL and DR on individual levels, with results in Figures 19, 20 and 21. We see that, generally, there is an upwards trend in the performance on most levels, but this is not monotonic. Additionally, on some levels (e.g. M-Thrust-Control), performance decreases over training, potentially indicating a bias in the levels trained on.

Figure: Figure 19: DR vs SFL on the full set of S levels.
Refer to caption: x22.png

Figure: Figure 20: DR vs SFL on the full set of M levels.
Refer to caption: x23.png

Figure: Figure 21: DR vs SFL on the full set of L levels.
Refer to caption: x24.png

## Appendix K Further General Agent Results

Figure 22 contains the performance of DR and SFL on each environment size.
We can see that, in every case, the agent’s performance increases throughout training, indicating that it is learning a general policy that it can apply to unseen environments.
In all cases, SFL is superior to DR, but the performance of both methods deteriorates as the environment size increases.
Interestingly, DR on L, which trains on random levels, performs worse than a random policy.

In Figure 23, we plot the performance of the models trained for Figure 3 on the other holdout sets. Here, we can see that when training on M and L, the agent is still able to zero-shot a number of the S levels.

Next, in Figure 24, we evaluate on a fixed set of randomly-generated levels of the appropriate size. This is to evaluate whether the agents are indeed learning useful behaviour on tasks that are in-distribution. Despite selecting potentially impossible levels, we find that the solve rate steadily increases over time.

Figures 25 and 26 show the performance of an agent trained on the L distribution for a longer time than the main results. Performance on random levels steadily increases, and there is also an upward trend in the solve rate on the holdout sets, indicating that we could expect further improvements by training this agent for longer.

Figure: Figure 22: Results for DR and SFL for S, M and L, respectively. In each pane, the training levels are sampled from the DR distribution of the corresponding size, and the y-axis measures the solve rate on the evaluation set of that same size. SFL outperforms DR, but both methods suffer as the environment size increases.
Refer to caption: x25.png

Figure: Figure 23: Each row corresponds to the same agents, evaluated on each holdout set.
Refer to caption: x26.png

Figure: Figure 24: Performance of the agents trained for Figure 3 on fixed sets of 1000 randomly-generated levels for each size.
Refer to caption: x27.png

Figure: Figure 25: Performance of a single seed, trained on L, on random levels.
Refer to caption: x28.png

Figure: Figure 26: Performance of a single seed, trained on L, on the holdout set of levels.
Refer to caption: x29.png

## Appendix L UED Results

In Figure 27, we present results for two popular UED methods, PLR (Jiang et al., ; ) and ACCEL (Parker-Holder et al., ), with the hyperparamters listed in Table 8. These results show neither method significantly outperformed DR, leading us to focus solely on SFL in the main text.

Figure: Figure 27: Solve rate on (left) M and (right) L evaluation sets for PLR, ACCEL and DR.
Refer to caption: x30.png

**Table 8: UED Hyperparameters.**
| Parameter | Value |
| --- | --- |
| PLR |  |
| Replay rate, $p$ | 0.5 |
| Buffer size, $K$ | 8000 |
| Scoring function | MaxMC |
| Prioritisation | Rank |
| Temperature, $\beta$ | 1.0 |
| staleness coefficient | 0.3 |
| Duplicate check | no |
| ACCEL |  |
| Number of Edits | 3 |

## Appendix M Ablations

Figure: Figure 28: The solve rate on M for different ablations. SFL denotes the training regime we used in the main text.
Refer to caption: x31.png

We perform ablations to investigate which factors played into the success of training the agent (Figure 28). All of these experiments are on the size M environments.
We first consider removing the filtering, and find that there is no large difference in performance for SFL.
Secondly, we run DR instead of SFL: as before, we find that DR performs significantly worse than SFL, indicating that prioritising levels based on learnability is important.
Finally, we consider using the Symbolic-Flat representation, and find that the performance of this method is significantly worse than Symbolic-Entity, likely due to the large number of symmetries inherent in the environment. Despite this, SFL with Symbolic-Flat does outperform DR with Symbolic-Entity.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
questions: See Weaknesses section. Additionally,
- The range of the y axis for the four plots on the right in Figure 5 are missing. Are these also from 0 to 1?
- How long does training take for training on 1B Kinetix environments? Would be good to see if training on such a large number of environments itself would be a bottleneck for learning generalist agents.
soundness: 3
strengths: - Kinetix provides 66 hand-designed levels while having the option to edit tasks with a graphical editor or to randomly generate more levels with rejection sampling. 
- The unified goal and dynamics within all environments encourage policies to have physical reasoning capabilities instead of merely memorizing the solution for some particular task, which is a valuable objective for researchers to pursue. 
- Kinetix provides a way to generate unlimited environments and tasks with a unified goal, objects, and dynamics, which could be of interest to multiple research communities like generalist RL policy learning, meta-learning, world modeling, spatial understanding and physics reasoning, and so on.
summary: The paper introduces Kinetix, a new 2D simulated benchmark designed for training generalist agents with capabilities in fine-grained motor control, navigation, planning, and physical reasoning. The benchmark is built on a novel hardware-accelerated physics engine called Jax2D. Kinetix enables the procedural generation of a vast amount of environments using simple shapes, joints, and thrusters, allowing for tasks that include robot locomotion, object manipulation, simple video games, and classic reinforcement learning scenarios. Each environment shares a unified objective: “make the green shape touch the blue shape without touching the red shape,” enabling zero-shot generalization to new environments within the same distribution. Experimental results show that policies trained across a wide range of environments generalize better to unseen tasks, and fine-tuning these generalist policies on new tasks yields improved performance over training from scratch. This work contributes to the development of generalist RL agents and open-ended physics-based control tasks.
weaknesses: - The paper notes that as the generated environments increase in complexity, they may become unsolvable, which could contribute to the lower performance observed in the Large-level environments. If so, how does this impact the usability and interpretability of the benchmark results? To what extent does this affect the performance results reported in Figure 3?
- It is unclear whether the proposed benchmark supports visual observations, which are essential for training generalist policies and building agents that can operate in real-world settings.
- Although Kinetix can generate a vast range of environments, it is unclear how this benchmark would generalize to tasks or environments outside of its defined task distribution.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Kinetix enables a wide distribution of morphologies and initial environment scene configurations, but it doesn’t deviate beyond the single unified goal. How can it be expanded to also cover a wide distribution of goals and potentially even task specifications?

2. Does Kinetix support parallel pixel rendering, such as for vectorized image-based experiments?

3. Is SFL only choosing between the S, M, and L levels?

4. Are the inputs into the policy purely one-hots? (One-hot encoding for each of the polygons, thrusters, and shapes.)

5. Is the (x, y) 2D position of each polygon an input into the network? It would seem that positional embeddings not of the ordering of polygons, but their actual spatial positions, would matter a great deal in this problem.

6. In section 4, authors write that Kinetix is a deterministic setting (thus satisfying one of the conditions for using SFL). How is Kinetix deterministic, given that environments are randomized?

7. How many environments were used during training, and how many environments were held-out for zero-shot evaluation?

8. What was the generalist policy in Figure 5 trained on? A distribution of environments over all 4 tasks at their hard level?

9. Say we train an agent only trained on L levels. How does it perform on held-out levels of a different difficulty (M and S)?

10. Heatmaps were an interesting way to convey the agent’s performance. Why not simply graph the agent’s x distance from the goal, with x-axis being the training iterations?
soundness: 2
strengths: 1. Introduces a physics engine that provides “almost entirely dynamically specified” scenes, where environments with different robot morphologies can be vmap-ed and run in parallel, which is not doable with prior Jax-based sim frameworks like Brax.

2. Paper is clearly written.
summary: This paper introduces Jax2D physics engine, which is a reimplementation of Box2D but written in Jax, and introduce the Kinetix environment on top of Jax2D. Kinetix allows for procedurally generated open-ended environments with different robot morphologies. Authors create a self-attention-based policy and demonstrate performance zero-shot, with pretraining, and with finetuning on the target environments.
weaknesses: 1. All environments in benchmark must fall under the goal of making green shape touch blue shape without touching red shape. This seems to mainly constrain the problem to single-step tasks, where the reward of minimizing the distance from green to blue always incentivizes progress. Was this unified goal constraint purposefully imposed by design, or was it a constraint of Jax implementation, where the reward function for all environments must be the same to be parallelizable?

2. Authors emphasized that parallelism and speed were big advantages of Jax2D. Since it is a reimplementation of Box2D, and this is a critical contribution of the paper, what are the performance gain metrics over Box2D?

3. Experiments were on multi-discrete action space with binary rewards. However, it would strengthen the argument of the paper to do experiments on more of the important features of Kinetix, such as pixel-based observations and continuous action space.

4. The state representation of the policy is very specific to the Kinetix environment suite and not very generalizable to other 2D RL problems. For instance, each entity is encoded separately and there is no scene-level encoding that is passed in as observation for the policy. Often, it is essential for a policy to understand the entire scene when predicting an action.

5. There were no supplementary materials submitted, which would have been a good opportunity to show video rollouts of the trained agent in action.

6. Experiments were mainly limited to the improvement of finetuned policies over pretrained and task-specific, trained-from-scratch policies. However, I would have liked to see more experiments that provide additional insights beyond “finetuning is mostly good” and “zero-shot mostly doesn’t work.” For instance, using Kinetix for lifelong learning, transfer learning, and cross-embodiment learning.

7. Abstract sentence seems like an oversell, given the results. “Our trained agent exhibits strong physical reasoning capabilities, being able to zero-shot solve unseen human-designed environments.” Most would also disagree with the 2D learned behaviors as “strong physical reasoning capabilities.”

8. Minor: I think the wrong citation was provided for MJX in Section 7 (that work seems to be Mujoco).

9. Minor: Experiments would benefit from some comparison to prior approaches/architectures, though this is less important given this is mainly a systems/framework paper.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: - For the environment generator, it is mentioned that there may exist unsolvable levels which automatic curriculum methods can filter out. Could you clarify what was done here?
- How does the choice of algorithm affect the performance in your benchmark. Do you anticipate releasing a dataset of transitions from Kinetics which can be used for offline 2 online RL?
soundness: 3
strengths: - Provide a highly efficient 2D rigid-body physics engine, leveraging JAX for scalable computation, with speedups of up to 30x when training an RL agent, allowing for 
- The learnt agent is highly effective at Zero-Shot transfer in the S and M levels that are held out, indicating the efficacy of pre-training on a wide set of procedural generation tasks. Additionally, show faster convergence/higher performance with this initialization
- Have interpretable/handmade levels to understand the performance on different sizes/difficulties of tasks.
summary: This paper provides a framework for procedurally generated 2D physics-based tasks to learn an agent for physical control that can effectively transfer to tasks that involve physical control. They provide a hardware-accelerated physics engine that allows for cheap/efficient simulation to generate a mixed-quality pre-training dataset for online RL. The authors additionally provide human interpretable handmade levels in Kinetix to understand the type of tasks and interactions the agent must do. The authors show the efficacy of both zero-shot transfer to new tasks and the agent when fine-tuned on a new task. The authors evaluate their agent on classic RL environments such as Cartpole and video games like Pinball.
weaknesses: - The JAX2D environment seems to be somewhat limited in its expressivities, modeling only 4 unique entities, which may not transfer to a wide set of domains/tasks outside of the ones studied.  
- The task/reward function seems to be fixed across all environments to collide the green and blue shaped objects, while avoiding red shapes. Additional reward shaping seems to be needed for effective training, leading to some limited applicability of generating this data at scale for any set of tasks.

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 4
questions: They are mentioned in the Weaknesses.
soundness: 3
strengths: - The paper is well-written, organized, and straightforward.
- Extensive testing across various task complexities validates its robustness in diverse 2D environments.
- This paper has strong potential to serve as a valuable benchmark for future research.
summary: This paper introduces Kinetix, a 2D physics-based RL environment aimed at enhancing generalization in RL agents. Leveraging the simulation, the agent is pre-trained for billions of steps, enabling zero-shot evaluation on novel tasks. Fine-tuning further improves performance, surpassing traditional RL methods. The approach integrates a transformer architecture with PPO as the core RL algorithm.
weaknesses: - **Real World Tasks:** While this paper provides a strong foundation in 2D simulations, expanding its scope to assess the agent’s adaptability to real-world tasks, such as 3D simulations or complex dynamics as seen in [1,2], would enhance its practical relevance. Bridging this gap could amplify the study’s contributions, offering broader insights into real-world generalization and scalability.

- **Filtering out:** The authors mention that trivial and unsolvable levels are filtered out. What quantitative metrics were used to determine this filtering.

- **Generalizability:** The claims of generalizability might be overstated given that the tasks remain in controlled simulations. Could the authors clarify the expected limitations of deploying such an agent in real-world scenarios with unpredictable environmental factors?

Meta-review: The paper introduces Kinetix, a 2D physics-based RL benchmark built on their new Jax2D physics engine. The main idea is to use procgen to create diverse control/physics based tasks. The framework enables pre-training of generalist RL agents that demonstrate strong zero-shot performance and fine-tuning capabilities. Jax2D is 30x more efficient than Box2D and is open sourced, along with pre-trained models.

The experiments validate the framework's effectiveness for transfer learning, particularly on diverse hand-designed tasks. During rebuttal, the authors ran extensive additional experiments covering pixel observations, continuous actions, cross-embodiment learning and lifelong learning scenarios. The benchmark's open-ended generation capabilities make it valuable for multiple RL research directions.

The main limitations stem from the constrained goal specification. The work is limited to 2D environments, so generalization claims may not extend to 3D or other realistic settings. But I think the authors make a fair point about the complexity of this domain itself and that current agents can't yet efficiently solve this domain fully. So it is a good test bed to make progress and learn about physics generalization. All major concerns were adequately addressed during the review process. So this is a very interesting direction to pursue further and the paper should be accepted.

Decision: Accept (Oral)