Kinetix: Investigating the Training of General Agents through Open-Ended Physics-Based Control Tasks
Authors: Michael Matthews
Michael Beukman
Chris Lu
Jakob Nicolaus Foerster
Venue: ICLR 2025 Oral (2025)
Keywords: reinforcement learning
open-endedness
unsupervised environment design
automatic curriculum learning
benchmark
Primary area: reinforcement learning
TL;DR: Training with reinforcement learning on a vast open-ended distribution of physics-based tasks leads to an agent that can zero-shot solve human-designed problems.

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.

Published as a conference paper at ICLR 2025
KINETIX: INVESTIGATING THE TRAINING OF GENERAL AGENTS THROUGH OPEN-ENDED PHYSICS PHYSICS-BASED CONTROL TASKS
Michael Matthews
Michael Beukman" Chris Lu
Jakob Foerster
FLAIR, University of Oxford
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, take step towards this goal by procedurally generat- ing 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 intro- duce Kinetix: an open-ended space of physics-based RL environments that can represent tasks ranging from robotic locomotion and grasping video games and classic RL environments, all within a unified framework. Kinetix makes use of our novel 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 RL agent tabula rasa. This includes solving some environments that standard RL training completely fails at. We believe this demonstrates the feasi- bility of large scale, mixed-quality pre-training for online RL and hope that Kinet ix will serve as a useful framework to investigate this further.
physics engine Jax2D that allows us
INTRODUCTION
The development of a general agent, capable of performing competently in unseen domains, has been long-standing goal in machine learning (Newell et al., 1959; Minsky, 1961; Lake et al., 2017). One perspective is that large transformers, trained on vast amounts of offline text and video data, will ultimately achieve this goal (Brown et al., 2020; Bubeck et al., 2023; Mirchandaní al., 2023). However, applying these techniques in an offline reinforcement learning (RL) setting often constrains agent capabilities those found within the dataset (Levine al., 2020; Kumar et al., 2020). 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 al., 2021: 2023), most RL environments represent a narrow and homogeneous set of scenarios (Todorov et al., 2012; Bellemare al., 2013; Brockman al., 2016; Cobbe al., 2019), limiting the generalisation ability of the trained agents (Kirk et al., 2023)-
In this paper, we aim address this limitation by introducing Kinetix: a framework for repre- senting the vast, open-ended space of 2D physics-basee based environments, and using it to train a gen- eral agent. Kinetix is broad enough to represent robotics tasks like grasping (Rajeswaran et al., 2017) and locomotion (Todorov et al., 2012), classic RL environments such as Cartpole (Barto et al., 1983), Acrobot (DeJong & Spong, 1994) and Lunar Lander (Brockman et al., 2016), as well video games like Pinball (Bellemare et al., 2013), along with the multitude of tasks that lie in the intervening space (see Figure 1). To run the backend of Ki nintis we developed Ja Jax2D, a handware-acceleraated physies engine that allows us to efficiently simulate the billions of environ- ment interactions required to train this agent.
* "Equal Contribution
We provide full code and models at https: /etteetiiitieeeeeeeeeeeeeeeeeeee ixineexeeeve...mm x--nv.giii git..b.t..... ic. httpps: //// /ii thub. com/M chael 'IMatthews/Jaxb)

---

Published as a conference paper at ICLR 2025
Train on random environments
Test on handmade environments
Grasper
Half- Cheethh
Aunobot
Piabali
Catapult
Lunar Lander
Unified Goal
Make the green shape touch the blue shape, without it touching the red shape
Figure 1: We train general agent 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.
Through sampling random Kinetix environments from the space of representable 2D physics problems, we can produce virtually unlimited supply of meaningfully diverse tasks for training. Since these levels are programmatically sampled, many are not useful for learning-indecs most 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., 2019; Brown et al., 2020; Dubey et al., 2024).
We find that RL agent trained on these environments exhibits understanding of general mechan- ical properties, with the ability to zero-shot solve unseen handmade environments (Section 5). We further analyse the benefits of fine-tuning this general agent specific hard environments and find that it greatly reduces the number of samples required to learn particular task, when comparing against tabula rasa agent. Fine-tuning also affords capabilities, including solving tasks for which an agent specifically trained does not make progress (Section 6).
In summary, contributions are:
1. We introduce Jax2D, a fast hardware-acceleratte 2D physics engine.
2. We introduce Kinetix iinettii open-ended space of RL environments within unified frame- work. We provide the capability sample random levels from the 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 significantly improved sample efficiency and new capabilities.
2 BACKGROUND
2. REINFORCEMENT LEARNING
We model the decision makiin process as a Markov Decision Process (MDP), which is defined as tuple (S, A. R, T), where S is the set of states; A is the set of actions; T S A - AS is the transition function, defining the distribution over next states T(s, a) given a current state s and action a; and R: S R is the reward function. We consider finite-horizon MDPs, with maximum
2

---

Published as a conference paper at ICLR 2025
number of timesteps T. The goal of an agent in RL is to maximise its discounted sum of rewards, G-EE 2 R2, where Rs:)( is the reward timestep t and 7/ 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 student. The student maximises its expected discounted return as in the standard RI. formulation, while the teacher chooses levels to maximise some utility func- tion, effectively inducing curriculum of tasks through training (Oudeyer et al., 2007; Florensa et al., 2018; Matisen et al., 2020; Narvekar et al., 2020; Dennis et al., 2020; Parker-Holder et al., 2022). In this paper, these tasks (we also refer these as levels or environments) are particular initial states, So E S. One common approach sets a level's utility as the negative of the agent's s return (Pinto et al., 2017), and another class of approaches instead uses regret (Dennis et al., 2020). Domain Randomisation (Jakobi, 1997; Tobin et al., 2017, 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. (2023) and Rutherford et al. (2024) sample levels in binary-outcome domains using learnability, defined as p(L - p), with p being the success rate of the agent the particular level. In this way, learnability disincentivises the teacher from sampling levels that the agent cannot solve at all (where = 0) or where the agent can already perfectly solve them (p- 1), meaning that the agent trains on levels with high learning potential.
2.3
RL IN JAX
JAX (Bradbury et al., 2018) is a Python library for writing parallelisable code for hardware acceler- ators. While deep RL has traditionally been divided between environments the CPU and models on the GPU (Mnih et al., 2015; Espeholt et al., 2018), JAX has facilitated the development of GPU- based environments (Lange, 2022; Koyamada al., 2023; Rutherford al., 2023; Nikulin al., 2023; Matthews et al., 2024; Kazemkhani et al., 2024; Bonnet et al.. 2024; Pignatelli et al., 2024), allowing the entire RL pipeline run on hardware accelerator (Hessel al., 2021). Through massive parallelisation and elimination of CPU-GPU transfer, this gives tremendous speed bene- fits (Lu et al., 2022). While UED has also followed this trend (Jiang et al., 2023; Coward et al., 2024), experiments have largely been confined to simple gridworlds, due to the lack of any suitable alternative (Garcin et al., 2024; Rutherford et al., 2024).
2.4 TRANSFORMERS AND PERMUTATION INVARIANT REPRESENTATIONS
Transformers and Attention Transformers (Vaswani et al., 2017) use the attention mecha- nism (Bahdanau et al., 2015) to model interactions within a set. Given N embeddings, x; R." self-attention computes queries 4., keys kú and values for each element through linear projec- tions. Weights for each element relative to element j are calculated as w ki and nor- malised via softmax to get ū. The new embedding for element i is weighted sum of the values: x. new Σ> "y. allowing each element to attend to others. The common practice of adding positional embeddings to encode sequence order (Vaswani et al., 2017) 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 an alternate solution (Lu et al., 2023; Bousmalis et al., 2023; Team et al., 2023; Raparthy et al., 2024). A less observations, common use of transformers in RL is for processing inherently permutation-invarian such as entities in Starcraft I/ (Vinyals et al., 2019). Although graphs are traditionally processed also now
with graph neural networks (Wang et al., 2018: Battaglia et al., 2018). transformers being applied to this domain (Sferrazza et al., 2024; Buterez et al., 2024), with attention masks set to graph's adjacency matrix restrict attention neighboring nodes (Sferrazza al., 2024).
3

---

Published as a conference paper at ICLR 2025
3 KINETIX
In this section, we introduce Kinettix a large and open-ended environment for RI., implemented entirely in JAX. We describe our underlying physics engine (Section 3.1), , the RL environment (Sec- tion 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, huge diversity of different physical tasks can be represented.
Jax2D simulates discrete Euler steps for rotational and positional velocities and then applies stantaneous 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 joint cannot separate the point of connection). Constraints are pairwise, meaning that it may be necessary to apply multiple steps of constraint solving for a stable simulation, espe- cially when simulating systems of many interacting bodies. The number of solver steps therefore serves as tradcoff between accuracy and speed. An agent (human or artificial) can act on the scene by applying torque through motors attached revolute joints by applying force through thrusters.
Jax2D is based on Box2D (Catto, 2007) and can be thought of as a minimalist rewrite of the li- brary in JAX. Appendix shows the benefit of this reimplementation, with hardware acceleration al- lowing Jax2D to easily scale to thousands of parallel environments on a single GPU, outperforming Box2D by factor of 4x when comparing just the engines and x0x when training an RL agent (this difference is due to Jax2D natively integrating with RI pipelines that exist entirely on the GPU).
The key differentiator of Jax2D from other ---bdaed physics simulators such Brax (Freeman et al,, 2021), ), 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 vmаp operation-4 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 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. See Appendix C for further information.
and continuous action spaces. In the multi- discrete action space, each motor and thruster can either be inactive, 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].
Action Space Kinetix supports both
Observation Space We use a symbolic observation where each entity (shape, joint or thruster) is defined by array of values of physical properties including position, rotation and velocity. The observation then defined as the set of these entities, allowing the use of crrmutation-invariin network architectures such as transformers. This observation space makes the environment fully observable, removing the need for policy with memory. We also provide the option for pixel-basee observations and symbolic observation that simply concatenates and flattens the entity information.

---

Published as a conference paper at ICLR 2025
Reward To facilitate our goal of general agent, we choose simple yet highly expressive reward function that remains fixed across all environments. Each scene must contain a green shape and a blue shape-the 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 Rd = K (d. dı+1), where de is the distance between the green and blue objects at timestep and is a coefficient that we tune ensure the dense signal does not dominate. We note that Kinetix could be run with many other reward formulations (Andrychowicz et al., 2017; Frans et al., 2024), which we leave to future work.
3.3 KINETIX: A BENCHMARK FOR INVESTIGATING OPEN-ENDEDNESS
The expressivity, diversity, and speed of Kinet i 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 heuristic environment generator, 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 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 motor or a thruster). Furthermore, we follow Team et al. (2021) and perform rejection sampling on levels solved with 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 al., 2022). We also provide functionality to generate levels using RL (Dennis et al., 2020) and generative models (Garcin et al., 2024).
algorithms like ACCEL (Parker-Holdes
Hland-Deiignn 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 levels. Some of these levels inspired by other RL bench- marks, such as L-MuJoCo-Wlkker, L-MuJoCo-Hopper: L-MuJoCo-HalE-Cheee ah, U-MuJoCo-Swimme Swirmer (Todorov al., 2012) and I--unna-ITander, L-Swi Dariinn-Uppp Up, L-Cartpole e Whee s-Hard (Brockman et al., 2016). We made other levels, like I.-P nba 1, L-Lorry and L-Catapult, specifically for Kinetix. These levels tests agent capabilities in- cluding 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 scene. While not strict rule. scenes with less entities tend represent simpler problems. We therefore quantise our exper- iments and handmade levels into one of three distinct sizes: small (S), medium (M), and large (I.). 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 language model can condition on 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 distri- bution. We refer to training 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
5

---

Published as a conference paper at ICLR 2025
same underlying structure and fully observable, it is theoretically possible to lear policy that can perform optimally on all levels inside the distribution.
To select levels train that regularly performs a large number of rollouts on randomly generated levels. then selects a subset of these with high learnability and trains them for fixed duration before again selecting new levels. SFL. filters out all unsolvable levels, as the success rate (and therefore also learnability) zero. The main limitation of SFL, that it is only applicable settings with deterministic transition dynamics and binary rewards, does not constrain us, as Kineti iix satisfies both of these assump- tions. We ran preliminary experiments using PLR (Jiang et al., 2021a;b) and ACCEL (Parker Holder et al., 2022), but found that these approaches provided no improvements over DR (see Appendix L).
SFL (Rutherford et al., 2024), state-of-the-an UED algorithm
For all experiments, we use PPO (Schulman al., 2017) with multi-discrete actions. We allot each method 5 billion environment interactions and periodically evaluate performance on the holdout levels. Hyperparameters are detailed Appendix H.
4.1 ARCHITECTURE
Message Passing (Thruster))
Message Passing (Jaints)
Fully Connected
Self Attention
V(8)
D
The transformer based architecture used for training. The scene is decomposed into its constituent entities and then passed through the network, consisting of layers of self-attention and message passing, followed by K fully connected layers.
Figure 2:
The architecture we use is summarised in Figure 2. To process the observation in a permutation- invariant way, we represent each entity as vector v, containing information about its physical prop- erties, such as friction, mass and rotation. We separately encode (using a set small feedforward networks) polygons, circles, joints and thrusters into initial embeddings x:!" where T {p, c, j, (). We perform self-attention (Bahdanau et al., 2015; Vaswani et al., 2017) over the set of shapes (i.e.. polygons and circles) without positional embeddings to obtain new shape embeddings a:S. To incor- porate joint information, we take each joint feature x!. and its two connected shapes / and z and pass the concatenation through a feedforward network. f, and add it to the embedding a: ahom- We have two feature vectors for each joint, with the from and to shape swapped. This layer is rem- iniscent of message passing in graph neural networks (Gilmer et al., 2017; Bronstein et al., 2021). Similarly, for each thruster af and associated shape 25. process these using message-passing layer and add the result back to .:S. 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 fully-connected graph and contains all shapes; the second allows shapes to attend to those that are connected by a joint (Sferrazza et al., 2024; Buterez et al., 2024); the third allows attention to shapes that 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. (2020), we gated transformer, and perform layernorm (Lei Ba al., 2016) before the attention block.

---

Published as a conference paper at ICLR 2025
S
M
0.8
Zero -hot Random
Solve Rate
0.6
0.4
0.2
0
1B
2B
3B
2B 3B
4B
5B
1B
Environment Interactions
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 on the evaluation of that size. The shaded area shows the standard error over 5 seeds.
ZERO-SHOT RESULTS
In Figure 3, we run SFL on the S, M and L environment sizes, respectively (see Appendix for a per- level breakdown). In each case, train randomly-generatee environments of the corresponding size, and we use the corresponding holdout set (see Appendix for a full listing) to evaluate the agent's generalisation capabilities. We 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 envi- ronments. For the agent very quickly learns policy superior 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 number of unseen hand-designed environments. On the I. 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 being trained and tested on more complex levels, 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 large set of mixed-quality levels can lead to general behaviour unseen tasks. See Appendix K for more detailed results.
5.1 ANALYSIS: ZERO-SHOT LOCOMOTION OF AN ARBITRARY MORPHOLOGG
In this section, take closer look at the zero-shot capabilities of the learned general agent by prob- ing its behaviour in a constrained goal-folle 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 maximise the dense auxil- iary 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, well Morphologyn wrrphhooy--andooo Ranndom, which selects from one of 2000 randomly generated 3-shape morphologies (Appendix F).
We measure how the position of the goal correlates with the æ position of the controllable mor- phology (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 general agent trained on random M levels for billion timesteps. Each plot is aggregated over 2000 randomly sampled levels, each of which is run for timesteps to allow the agent to maneuver into position and then run for further 64 timesteps for data collection.
As would be expected, the random agent shows no correlation between the position of the control- lable morphology and the goal. By contrast, the trained agent shows positive correlation, indicating able to mancuver the morphology towards the goal location. We see a variety of outcomes across the different morphologies, with the agent showing very strong results with a slightly weaker performance on Snake. When evaluating on Mo corrhhllee hatphhoooy-aaaa ogy-Random, we do some positive correlation, although not strong as the hand-designed levels.
Car and Thrust,
The positive results on these constrained "goal-conditioned" environments show that the agent has indeed learned general policy that encompasses purposeful locomotion of an arbitrary morphology.

---

Published as a conference paper at ICLR 2025
Car
Snake
T Thrust
Morphology Random
Example Level
General Agent
Body
Random Agent
Target
Figure 4: Heatmaps of goal 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.
FINE-TUNING RESULTS
In this section we leave the zero-shot paradigm and investigate the performance the general agent when given limited number of samples to fine-tune the holdout tasks. In particular, in Figure we train a separate specialist agent for each level in the holdout set, and compare this to fine-tuning general agent (the same used for Section 5.1, trained for 5B timesteps We plot the learning curves for four selected environments, as well as the aggregate performance On three of these levels, fine-tuning the agent drastically outperforms
random M levels.).
the entire holdout training from scratch. In particular, for 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 that fine-tuning beats training from scratch, we do see one case: Thruster-Large cecrrgg--bbtatllle where fine-tuning learns slower.
Hopper Hard and Mujoco-
All L tasks
Galoot--Half-Cheeeee Cheetah-Hard
MuJaCooHHopper-ee
1.0
1.0
Tabula Rasa
Fine-Tuned
0.5
Solve Rate
0.0
0.5
NoJoCo MuJloCoNNlkker Valker Hard
Thruster Large Obstacles
0.5
0.0
0.0
100M
100M
100M
Environment Interactions
Figure 5: The performance of fine-tuned and tabula rasa agents (left) aggregated over the entire holdout set, and (right) for four selected levels. We train separate agent for each environment and plot mean and standard error over five seeds. We stress that the MuJoCo levels are reimplementa- tions of the classic environments in Kine Kinetix.
8

---

Published as a conference paper at ICLR 2025
6.1 ANALYSIS: GENERAL PRETRAINING CAN BEAT TRAINING ON THE TARGET TASK
We now further investigate the case of Far-Ramp (Figure 6a) where RL, even with large sample budget, fails solve but that our fine-tuned general agent can complete (note that this behaviour is also shown in MuJoCo- Wa ker-Hard), Car-Ramp is an example of a deceptive problem (Gold- berg, 1987; Liepins & Vose, 1991; Lehman Stanley, 2011) I) 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 I billion timesteps fails 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 small amount of fine-tuning, the agent can reliably solve this task (Figure 6b). We do stress that, while impressive, this behaviour the exception rather than the rule, only occurring on of 74 handmade levels. We see this as a promising sign for trained general agent in Kinet ix to serve strong basc model.
(a) Initial State
(b) Performance
Figure 6: The Car-Ramp Environment. We use single seed for the pre-trained agent (trained on for 5B timesteps), while averaging over seeds for the others. Error bars indicate standard error.
RELATED WORK
Hardware-Accelerater Physics Engines Jax2D joins a thriving ecosystem of hardware- accelerated physics engines used in RL tasks. Brax (Freeman et al., 2021), MJX (Todorov et al., 2012) and Isaac-Gym (Makoviychuk al., 2021) have all been been widely used in the RL commu- nity, 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 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 heterogencous environments to be parallelised.
Physical Reasoning PHYRE (Bakhtin et al., 2019) also uses 2D rigid-body physics by tasking agents with placing ball to achieve problem, allowing the agent to take actions throughout the episode. A crucial difference is that we train on large automatically generated set of tasks rather than small set of handmade ones.
goal state. Li al. (2024a) extend this bandit-lik
Berdware-Accelerattee RL Our work follows the recent trend of using hardware-acceleratte RL environments to significantly larger-scale experiments than would be possible with CPU-based environments (Lu et al., 2022; Jackson et al., 2023; 2024; Goldie et al., 2024; Rutherford et al., 2024: Nikulin al., 2024; Kazemkhani et al., 2024). By leveraging Kinet ix'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 learn a generalist foundation model for robotics (Reed et al., 2022; Bousmalis et al., 2023; Team et al., 2024; Nasiriany et al., 2024; O' Neill al., 2024). While most of these approaches perform behaviour cloning on a large dataset from variety of robot morphologies and tasks, Nasiriany et al. (2024) develop large-scale simulation en- vironment, with an initial focus on kitchen environments, By contrast, K nit aims to train an on- line agent tabula rasa, without using external data, and further has large variety of different tasks.
Open-Ended Learning Kinetix also ties into the paradigm of open-ended learning (Soros & Stanley, 2014; Stanley, 2019; Sigaud et al., 2023; Hughes et al., 2024), in which system contin- ually generates new and novel artifacts. In the context of RL, this often means training within a large and diverse distribution and applying some method (e.g., UED) to adapt this distribution over time. While these methods hold the promise of generating novel and useful levels an open- open-ended manner, the environments used in their experiments often very constrained in what they can rep-

---

Published as a conference paper at ICLR 2025
resent (Wang al., 2019; Dennis et al., 2020; Jiang al., 2021b;a; Parker-Holder et al., 2022). As we have shown, in a significantly more diverse task space, these approaches tend to fail.
A recent work with similar vision to Kinetix is Autoverse (Earle & Togelius, 2024), where an agent acts inside a cellular automata based gridworld, where changing the underlying rules can lead to many diverse levels. Relatedly, Sun et al. (2024) use prior knowledge in the form of large language models to generate simulation code to train RL agents in. Powderworld (Frans & Isola, 2023) instead creates an 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., 2024), Jarvis-1 (Wang al., 2023) and Optimus- I (Li al., 2024b). These are more focused on mong-horzoon planning, the self discovery of new tasks to perform, and use Minecraft as their domain with prior knowledge in the form of large language model.
Perhaps the work most similar to ours is the highly impactful XLand line of research (Team et al., 2021; 2023). XLand defines a large and diverse distribution of levels inside a 3D physics simulation, with embodied agent (or set of agents) required to fulfil some specified goal Similar to us, agents train on procedurally generated levels and are 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 subjectively claim that Kinetix, through the representation of almost any conceivable 2D rigid-body physics problem, has more expressive universe of tasks. While Laa also employs a physics engine, all the tasks are constrained to homogeneous agents acting in the world, potentially limiting its scope-it is not clear, for instance, how one would represent any of the holdout environments in Figure in XLand. Lastly, we note that XLand's source code unavailable, limiting its use for future research. Although XLand-Minigrid (Nikulin et al., 2023) provides a fast, open-source version of XLand, it simplifies the environment into a gridworld.
8 DISCUSSION AND FUTURE WORK
We believe Kinetix x is a uniquely diverse, fast and open-ended environment, placing it well as 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., 2019; Chevalier-Boisver et al., 2023; Rutherford et al., 2024), K Kinetis ix x 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., 2020), rather than just filtering from predefined distribution. We also believe Kinetix x is an excellent frame- work for investigating issues in agent training such as network capacity (Obando-Ceron et al., 2024), plasticity loss (Igl et al., 2020; Berariu et al., 2021: Sokar et al., 2023), lifelong learning (Kirkpatrick et al., 2017) and multi-task learning (Sodhani et al., 2021; Hafner, 2021; Benjamins et al., 2023).
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 meta-learn parts of the RL process, for instance the algo- rithm (Oh et al., 2020; Lu et al., 2022; Jackson et al., 2023), optimiser (Goldie et al., 2024) or loss function (Bechtle al., 2021). Alternatively, the emerging capabilities of large world mod- els (Bruce et al., 2024; Valevski et al., 2024) hint at a new paradigm of online training entirely in imagination (Ha Schmidhuber, 2018; Yu et al., 2020; Hafner et al., 2020; 2021; 2023), 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.
CONCLUSION
In this work, we first introduce Jax2D, hardware-accelerann 2D physies engine. Using Jax2D, we build Kineti sity of Kinetix ix 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 envi- ronment interactions from randomly generated tasks, and show that it can zero-shot generalise to many human-designed tasks, well as function as strong base model for fine-tuning. We hope that this work can serve as a foundation for future research in ppen-endedness, large-scale online pre-training of general RI. agents and unsupervised environment design.
vast and open-ended physics-based RL environment. We illustrate the diver-
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
We would like to thank Thomas Foster, Alex Goldie, Matthew Jackson, Sebastian Towers, Andrei Lupu and our anonymous reviewers for insightful discussions and valuable feedback that aided the development of this project and the production of the manuscript. This work was supported by UK Research and Innovation and the European Research Council, selected by the ERC, and funded by the UKRI [grant number EP/Y022481/1] We also thank the authors of the game Incredibots, which served as an initial inspiration for the development of Kinetix.
REFERENCES
Marcin Andrychowicz. Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Picter Abbeel, and Wojciech Zaremba. Hindsight experience re- play. Advances in neural information processing systems, 30, 2017.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning align and translate. In 3rd International Conference Learning Representations, 2015. URL htttt::/aaxxv.orrrr org abs 110///101...00 C473.
Anton Bakhtin, Laurens van der Maaten, Justin Johnson, Laura Gustafson, and Ross B. Girshick. PHYRE: A new benchmark for physical reasoning. In Advances in Neural Information Process- ing Systems, pp. 5083-5094, 2019. URL https: httpt:tttt///oteeeec http://totceeen.sceeeepepoo..sse nww.ips....... cc/paper/ 201
Andrew Barto, Richard Sutton, and Charles W. Anderson. Neuronlike adaptive elements that can solve difficult learning control problems. IEEE Trunsactions on Systems, Man, and Cybernetics, 2MC-13(5):834-846 1983. doi: 10.1109/TSMC..193.3333333
Peter W. Battaglia, Jessica B. Hamrick. Victor Bapst. Alvaro Sanchez-Gonzale Vinícius Flores Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, Çaglar Gülçehrre, H. Francis Song, Andrew J. Ballard, Justin Gilmer, George E. Dahl, Ashish Vaswani, Kelsey R. Allen, Charles Nash, Victoria Langston, Chris Dyer, Nicolas Heess, Daan Wierstra, Pushmeet Kohli, Matthew M. Botvinick, Oriol Vinyals, Yujia Li, and Razvan Pascanu. Relational inductive biases, deep learning, and graph networks. CoRR, 0os//806.0122 2018.
URL http: /ttp/vaaiiiv//./////v///e../aaaa org/ /10b..///...... 01261.
Sarah Bechtle, Artem Molchanov, Yevgen Chebotar, Edward Grefenstette, Ludovic Righetti, Gaurav Sukhatme, and Franziska Meier. Meta learning via learned loss. In 2020 25th International Conference on Patlern Recognition (ICPR), pp. 4161-4168. IEEE, 2021.
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
Carolin Benjamins, Theresa Eimer, Frederik Schubert, Aditya Mohan, Sebastian Döhler, André Biedenkapp, Bodo Rosenhahn, Frank Hutter, and Marius Lindauer. Contextualize me the case for context in reinforcement learning. Transactions Machine Learning Research, 2023. ISSN 2835-8856. URL https:://///o /ttpr///ee//peeeie...ooom net forum? id -oeidd-444400 Y42xVBQusn.
Tudor Berariu, Wojeiech Czarnecki, Soham De, Jorg Bornschein, Samuel Smith, Razvan Pas- canu, and Claudia Clopath. A study on the plasticity of neural networks. arXiv preprint arXiv:2/06. 00042, 2021.
Clément Bonnet, Daniel Luo, Donal Byrne, Shikha Surana, Sasha Abramowitz, Paul Duckworth, Vincent Coyette, Laurence I. Midgley, Elshadai Tegegn, Tristan Kalloniatis, Omayma Mahjoub, Matthew Macfarlane, Andries P. Smit, Nathan Grinsztajn, Raphael Boige, Cemlyn N. Waters, Mohamed A. Mimouni, UIrich A. Mbou Sob, Ruan de Kock, Siddarth Singh, Daniel Furelos- Blanco, Victor Le, Amu Pretorius, and Alexandre Laterre. Jumanji: diverse suite of scalable reinforcement learning environments in jax, 2024. URL httpa: : /arxiv org/ abs/abos000 09884.
11

---

Published as a conference paper at ICLR 2025
Konstantinos Bousmalis, Giulia Vezzani, Dushyant Rao, Coline Devin, Alex X. Lee, Maria Bauzá, Todor Davchev, Yuxiang Zhou, Agrim Gupta, Akhil Raju, Antoine Laurens, Claudio Fantacci, Valentin Dalibard, Martina Zambelli, Murilo F. Martins, Rugile Pevceviciute, Michiel Blokzijl, Misha Denil, Nathan Batchelor, Thomas Lampe, Emilio Parisotto, Konrad Zolna, Scott F. Reed, Sergio Gómez Colmenarejo, Jon Scholz, Abbas Abdolmaleki, Oliver Groth, lean-Baptiste Regli, Oleg Sushkov, Thomas Rothörl, Jose Enrique Chen, Yusuf Aytar, Dave Barker, Joy Ortiz, Mar- tin A. Riedmiller, Jost Tobias Springenberg, Raia Hadsell, Francesco Nori, and Nicolas Heess. Robocat: A self-immroo improving foundation agent for robotic manipulation. CoRR, abs/2306. 1706, 2023. doi: 10.48550/arXiv.23006 1706. URL https: /dii. http:///..o.rr...////// 485.0//555501...
2306. 11706.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
// //github. com/ googles e/ jax.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. CoRR, abs/1606.01 0440, 2016. URL http::/axxxi.......//// abs/1606 01540.
Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Velickovie. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. CoRR, abs/2104. 13478, 2021. URL https: //aaxxiv.rrg////eee.ee....ao.lat htt.o/////../..........
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Advances in Neu- ral Information Processing Systems, 2020. URL https: http:///ooceee/eeeeeeoooocc neurips.coo cc/ 9674
hatper/2220/haah//// 45t.5///000000...4444.00.0.2626. A html.
Jake Bruce, Michael Dennis, Ashley Edwards, Jack Parker-Holder, Yuge Shi, Edward Hughes, Matthew Lai, Aditi Mavalankar, Richie Steigerwald, Chris Apps, et al. Genie: Generative inter- active environments. In Forty-first International Conference on Machine Learning, 2024.
Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Ka- mar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott M. Lundberg, Harsha Nori, Hamid Palangi, Marco Túlio Ribeiro, and Yi Zhang. Sparks of artificial general intelligence: Early experi- ments with GPT-4. CoRR, abs/2303. 12712, 2023. doi: 0.44850/ARXX1.233.1227 12712.
URL
https:://oi.... ort.:...rg/11... 48550/8550//// arxiv. 2372.12712.
David Buterez, Jon Paul Janet, Dino Oglic, and Pietro Lio. Masked attention is all you need for graphs. CoRR, abs/2402. 2024. doi: 1044550/RXXI.2400.00777 URL https:// //
doi. doi.org/111 000.44455///00
Erin Catto. Box2d,
2007.
Maxime Chevalier-Boisvert, Bolun Dai, Mark Towers, Rodrigo de Lazcano, Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry. Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks. CoRR, abs/2306.13831 2023. doi: 10.48550/ARXIV2236666 13831. URL httpp:://oi.... org/10. 4 48550/arxiss 2.36 13831.
Karl Cobbe, Oleg Klimov, Christopher Hesse, Taehoon Kim, and John Schulman. Quantifying generalization in reinforcement learning. Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 1282-1289. PMLR, 2019. URL http: httt:://prrceed:: ngs .. mll.......... mir.press// v97 / cobbe 9a html.
Samuel Coward, Michael Beukman, and Jakob Foerster. Jaxued: A simple and useable ued library in jax. arXiv preprint, 2024.
12

---

Published as a conference paper at ICLR 2025
Gerald Delong and Mark W Spong. Swinging up the acrobot: An example of intelligent control. In Proceedings of 1994 American Control Conference-ACCC994 volume 2, pp. 2158 2162.. IEEE, 1994.
Michael Dennis, Natasha Jaques, Eugene Vinitsky, Alexandre M. Bayen, Stuart Rus- sell, Andrew Critch, and Sergey Levine. Emergent complexity and zero-shot trans- fer via unsupervised environment design. In Advances in Neural Information Process- ing Systems, 2020. URL 985e9a46e10000 (10.00055500001999922296539 Abstract html.
neurips. ccc cccpppper/222/haass/
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chupter of the Association for Computational Linguistics: Human Language Technologies. pp. 4171-4186, Association for Computational Linguistics, 2019. doi: 10.18653/
v1/n19- 1423. URL ht httpp: /oi. wwg./ogg///.oooo.ooo 18 653
1423.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Archi Mitra, al. The llama 3 herd of models. CoRR, abs/2407.21783, 2024. doi: 10-46550/ARXXV222722227333 URL https: /tl:://do /doii aee.i.or///...oo org/10. 48550/arXis wtw.ar.wwi............ .100722118..
Sam Earle and Julian Togelius. Autoverse: An evolvable game language for learning robust embod- ied agents. CoRR, 2024. URL https: /tttps::/dci.oro//
485500 arXiv. 2407. 04221.
Lasse Espeholt, Hubert Soyer, Rémi Munos, Karen Simonyan, Volodymyr Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, Shane Legg, and Koray Kavukcuoglu. IMPALA: scalable distributed deep-rl with importance weighted actor-learner architectures. In Proceed- ings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmüsssn, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings Machine Learning Re- search, pp. 1406-1415. PMLR, 2018. URL http: http::/pproceedinge mlr. w....peess/vB. espeholt188 html.
Carlos Florensa, David Held, Xinyang Geng, and Pieter Abbeel. Automatic goal generation for reinforcement learning agents. In Proceedings of the 35th International Conference on Ma- chine Learning, ICMI. 2018, Stockholmsmässann Stockholm, Sweden, July 15-15, 2018, vol- ume 80 of Proceedings of Machine Learning Research, pp. 1514-1523, PMLR, 2018. URL http: // /ptppppeeeenngasee press /v80/ Sazzes/v00/Eooreaaaaa3. www.ensall...t.........
Kevin Frans and Phillip Isola. Powderworld: platform for understanding generalization via rich task distributions. In The Eleventh International Conference on Learning Representations, ICIR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL htRps:/ /ttp//////eeeeeeeeeeeeeeee net / worffformmiiiomom Awd.dwZgGGss....
Kevin Frans, Seohong Park, Pieter Abbeel, and Sergey Levine. Unsupervised zero-shot reinforce- ment learning via functional reward encodings. arXiv preprint arXiv:2402.17135, 2024.
C. Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax- differentiable physics engine for large scale rigid body simulation, 2021. URL http: //githuse com/ google e/erax.....
Samuel Garcin, James Doran, Shangmin Guo, Christopher G. Lucas, and Stefano V. Albrecht. Dred: Zero-shot transfer in reinforcement learning via data-regularised environment design. 2024. URL
//tts::/doi.... wwwooor//o0 wtt.or////.......... 48550/ arXiv. 2402 03479.
Randy Gaul. Impulse engine, 2013. ImpulseEngine.
URL https: http:///gggtuub.cmo/Roaccce
Justin Gilmer, Samuel Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp. 1263-1272. PMLR, 2017.
David E Goldberg. Simple genetic algorithms and the minimal, deceptive problem. Genetic algo- rithms and simulated annealing, pp. 74-88, 1987.
13

---

Published as a conference paper at ICLR 2025
Alexander David Goldie, Chris Lu, Matthew Thomas Jackson, Shimon Whiteson, and Jakob Nico- laus Foerster. Can learned optimization make reinforcement learmin less difficult? arXiv preprint arXiv:2407.07082, 2024.
David Ha and Jürgen Schmidhuber. World models. CoRR, abs/1 803.10122, 2018. URL http:
//aaxii..rr/aas///eee...e. 10122.
Danijar Hafner. Benchmarking the spectrum of agent capabilities. arXiv preprint (100:220078) 2021.
Danijar Hafner, Timothy P. Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to trol: Learning behaviors by latent imagination. In 8th International Conference on Learning Representations. OpenReview.net, 2020. URL https: /openrevieww net. forum?id= S110TC4tDS.
Danijar Hafner, Timothy P. Lillicrap. Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In 9th International Conference on Learning Representations. OpenRe- view.net, 2021. URL httpt:tt//// /tppr////eeeeeeeieeeemmee net Fowrmmidddooaayy/eee
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy P. Lillicrap. Mastering diverse domains through world models. CoRR, abs/2301.04 2023, doi: 1044850/ARXX1.200000000 URL https: /tw.///o...///////////// org/10.../... 48550/ arxiv
Matteo Hessel, Manuel Kroiss, Aidan Clark, Iurii Kemaev, John Quan, Thomas Keck, Fabio Vi- ola, and Hado van Hasselt. Podracer architectures for scalable reinforcement learning. CoRR, 0ls/2104.06672, 2021. URL https: //apiia./a//////aa/eaaeaaeeiiaea org. abs 05222
Edward Hughes, Michael D Dennis, Jack Parker-Holder, Feryal Behbahani, Aditi Mavalankar, Yuge Shi. Tom Schaul, and Tim Rocktäschel. Position: Open-endedness is essential for artificial su- perhuman intelligence. Forty-first International Conference on Machine Learning, 2024. URL https: happs:///peerevee wpwpreveeeeeeeeeeeeeee net/ Dorum?1d=Bcc4222X77
Maximilian Igl. Gregory Farquhar, Jelena Luketina, Wendelin Boehmer, and Shimon Whiteson. Transient non-stationarity and generalisation deep reinforcement learning. arXiv preprint arXiv:2006.05826, 2020.
Matthew Jackson, Chris Lu, Louis Kirsch, Robert Lange, Shimon Whiteson, and Jakob Foerster. Discovering temporally-aware reinforcement learning algorithms. In International Conference on Learning Representations, 2024.
Matthew T Jackson, Minqi Jiang, Jack Parker-Holder, Risto Vuorio, Chris La, Greg Farquhar, Shimon Whiteson, and Jakob Foerster. Discovering general reinforcement learning algorithms with adversarial environment design. In Advances in Neural In- formation Processing Systems, volume 36. pp. 79980-79998, Curran Associates, Inc., wtw.::eeeddign.......eeeee..eeee neurips cetes.c///e httpeecccppppppppeeeeeeeeee (taprre__iieesppppe223333
2023. URL https:: https::/pprceeeengg... file fcs2d8a485746 orr.aaaae......aaareeoe000000rrr pdf.
Nick Jakobi. Evolutionary robotics and the radical envelope-of hypothesis. hypoteeess Adaptive behavior, 6(2):325-368, 1997.
Minqi Jiang, Michael Dennis, Jack Parker-Holder Jakob N. Foerster, Edward Grefenstette, and Tim Rocktäschel. Replay-guided adversarial environment design. In Advances in Neural Information Processing Systems, pp. 1884-1897, 2021a. URL https: / http::/:////////poooee///e hww.//aahh///ca 0eth http:////.......................
neurips cc/ htmll.
Minqi Jiang, Edward Grefenstette, and Tim Rocktäschel. Prioritized level replay. In Proceedings of the 38th International Conference on Machine Learning, volume 139, pp. 4940-4950. PMLR, 2021b. URL http::ttt/ccc
h http::///cocceeen.....o http:/g........rre////////////// htmml.
Minqi Jiang, Michael Dennis, Edward Grefenstette, and Tim Rocktäschel. minimax: Efficient base- lines for autocurricula in jax. In Agent Learning in Open-Endeding Workshop NeurIPS, 2023.
14

---

Published as a conference paper at ICLR 2025
Saman Kazemkhani, Aarav Pandya, Daphne Comelisse, Brennan Shacklett, and Eugene Vinitsky. Gpudrive: Data-driven, multi-agent driving simulation at million FPS. CoRR, 0000//24088081444, 2024. doi: 1044850/ARXXV.2200000555. URL https: /ttp://////oo..rr10.4485555///// 2408. 01584.
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of zero-shot generalisation in deep reinforcement learning. J. Artif. Intell. Res., 76:201-264, 2023. doi: 10.1613/jair...1e 1.t.///arr....174...... 14144. URL https: //tp:::////..../ ort.orgg///.... 11..1........................... 111.jaar....117..............
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska- Barwinska, al. Overcom- ing catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 114(13):3521-355 2017.
Sotetsu Koyamada, Shinri Okano, Soichiro Nishimorii Yu Murata, Keigo Habara, Haruka Kita, and Shin Ishii. Pgx: corddware-accelraate parallel game simulators for reinforcement learning. In Advances in Neural Information Processing Systems, volume 36, pp. 45716-45743, 2023.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1 179-1191, 2020.
Brenden M. Lake, Tomer D. Ullman, Joshua B. Tenenbaum, and Samuel J. Gershman. Building machines that learn and think like people. Behuvioral and Brain Sciences, 40:e253, 2017. doi: 10.10177/00145555160001111
Robert Tjarko Lange. gymnax: A JAX-based reinforcement learning environment library, 2022.
URL http: thubb com/ htt///bbbbrrttee gymnax.
Joel Lehman and Kenneth Stanley. Abandoning objectives: Evolution through the search for novelty alone. Evolutionary computation, 19:189-223. 06 2011. doi: 10.1 10.1162/EVCO 00005..
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. ArXiv e-prints, pp. arXiv-1607, 2016.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. CoRR, abs/2005.01 2020. URL
https: /ttp:///axvvv../
org/ abs 22000.00..02 01643.
Shiqian Li, Kewen Wu, Chi Zhang, and Yixin Zhu. I-PHRR: Interactive physical reasoning. In
The Twelfth International Conference on Learning Representations, 2024a. URL https://
openreview net/ wttp:/ooommmmmooooeeee
Zaijing Li, Yuquan Xie, Rui Shao, Gongwei Chen, Dongmei Jiang, and Liqiang Nie. Optimus- I: Hybrid multimodal memory empowered agents excel in long-horizon tasks. 2024b. URL https:/ //wpi///iii/////oo wwo//0.../......o. 15550/a0iiv....////aaaa 2408 03615.
Gunar E Liepins and Michael D Vose. Deceptiveness and genetic algorithm dynamics. In Founda- tions of genetic algorithmss, volume 1, pp. 36-50. Elsevier, 1991.
Chris Lu, Jakub Kuba, Alistair Letcher, Luke Metz, Christian Schroeder de Witt, and Jakob Foerster. Discovered policy optimisation. Advances in Neural Information Processing Systems, 35:16455- 16468, 2022.
Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob N. Foerster, Satinder Singh, and Feryal M. P. Behbahani. Structured state space models for in-context reinforcement learn- ing. In Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NewrIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URL /ptp:::p/pppp....ppeppp nips ct.ppper.cpppp...pccmmm 92d3d2a ce http://ppppppaappppppoopeapeeaep h h html.
15

---

Published as a conference paper at ICLR 2025
Viktor Makoviychuk, Lukasz Wawrzyniak, Yunrong Guo, Michelle Lu, Kier Storey, Miles Macklin, David Hoeller, Nikita Rudin, Arthur Allshire, Ankur Handa, and Gavriel State. Isaac gym: High performance GPU based physics simulation for robot learning. In Proceedings of the Neural Information Processing Systems Track on Datasets and Bench- marks I, NeurlPS Datasets and Benchmarks 2021, December 2021, virtual, 2021. URL attpo:///attaettteeeeccaaaaaaae proceeedig wtwpeeeeiiinnn....iiieeee neurips cc/paper/2021/ wwwchmarks Abstrratt coretttroundd..... html.
Abstract
Tambet Matiisnn, Avital Oliver. Taco Cohen, and John Schulman. Teacher-student curriculum learning. volume 31, pp. 3732-3740, 2020. doi: 10.1 10.1109/NNNS https:://oi..... orgg 1109 TNNLS. 2019 9.09..22 09119.2999066
URL
Michael Matthews, Michael Beukman, Benjamin Ellis, Mikayel Samvelyan, Matthew Jackson, Samuel Coward, and Jakob Foerster. Craftax: A lightning- benchmark for open-ended inforcement learning. In JCML, 2024.
Marvin Minsky. Steps toward artificial intelligence. Proceedings of the IRE, (9(1):9-128) 30, 1961.
Suvir Mirchandani, Fei Xia, Pete Florence, Brian Ichter, Danny Driess, Montserrat Gonzalez Are- nas, Kanishka Rao, Dorsa Sadigh, and Andy Zeng. Large language models general pattern machines. In Conference on Robot Learning, CoRL 2023, 6-9 November 2023, Atlanta, GA, USA, volume 229 of Proceedings of Machine Leurning Research, pp. 2498-2518. PMLR, 2023. / 9/m9//mmm https//////mmmmmaaaaaaaa. www.ncaaannnnnn... html.
URL https:tp/////c httpp//poeeeng...e/e/ mlr prwsspsssppp...... / v229
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning, nature, 51 (187540):529-5333 2015.
Sanmit Narvekar, Bei Peng, Matteo Leonetti, Jivko Sinapov, Matthew E. Taylor, and Peter Stone. Curriculum learning for reinforcement learning domains: A framework and survey. J. Mach. Learn. Res., 21:181:1-181:500 2020. URL /ttt:://fol http://jmlr.. http://mm......pppeeeeeeeee wttp://pppes//1122--11111 html.
Soroush Nasiriany, Abhiram Maddukuri, Lance Zhang, Adeet Parikh, Aaron Lo, Abhishek Joshi, Ajay Mandlekar, and Yuke Zhu. Robocasa: Large-scale simulation of everyday tasks for generalist robots. CoRR, abs/2406.00523,, 2024. doi: 1044550ARRXI.2406002522, URL
https://doi. http:://dooo.og/0..// org/10....... 48550/ aww.vvrriivv..../ 2406. 02523.
Allen Newell, John Shaw, and Herbert Simon. Report on general problem solving program. In IFIP congress, volume 256, pp. 64. Pittsburgh, PA, 1959.
Alexander Nikulin, Vladislav Kurenkov, Ilya Zisman, Viacheslav Sinii, Artem Agarkov, and Sergey Kolesnikov. Xland-minigrid: Scalable meta-reinforcemeet learning environments in JAX. In Amrrinsically-Motivvaet and Open-Ended Learning Workshop, NeurIPS2023, 2023. URL
https: wwwp//eieeeeeeeeeeeeeeeeeee.ewwa net forumm id wtt.fommmdxxAiaaaaaaa.... xAHDC4aHG..
Alexander Nikulin, Ilya Zisman, Alexey Zemtsov, Viacheslav Sinii, Vladislav Kurenkov, and Sergey Kolesnikov. Xland- I 100: A large-scale multi-task dataset for in-context reinforcement learning. CORR, abs/2406.08873, 2024. doi: 10. 00-45550/AAXXX1.240..09999 URL https: /do. org/ 10. 485408855553 246 08973.
Johan Obando-Ceron, Ghada Sokar, Timon Willi, Clare Lyle, Jesse Farebrother, Jakob Foerster, Gintare Karolina Dziugaite, Doina Precup, and Pablo Samuel Castro. Mixtures of experts unlock parameter scaling for deep rl. arXiv preprint arXiv:2402 08669, 2024.
Junhyuk Oh, Matteo Hessel, Wojciech M Czarnecki, Zhongwen Xu, Hado van Hasselt, Satinder Singh, and David Silver. Discovering reinforcement learning algorithms. Advances in Neural Information Processing Systems, 13:1060-1070, 2020.
Abby O'Neill, Abdul Rehman, Abhiram Maddukuri, Abhishek Gupta, Abhishek Padalkar, Abra- ham Lee, Acom Pooley, Agrim Gupta, Ajay Mandlekar, et al. Open x-embodiment: Robotic learning datasets and RT-X models : Open x-embodiment collaboration. In IEEE Interna- tional Conference Robotics and Automation, ICRA 2024, Yokohama, Japan, May 13-17,
16

---

Published as a conference paper at ICLR 2025
2024, pp. 6892 6903. IEEE, 2024. doi: 10.1 10.1099/CKA57147.002.100011 URL https: / 1t0.::/CC.2TTT11..222....1111...
Pierre-Yve Oudeyer, Frdric Kaplan, and Verena V Hafner. Intrinsic motivation systems for au-
tonomous mental development. IEEE transactions on evolutionary computation, (1(2):265-286 2007.
Emilio Parisotto, H. Francis Song, Jack W. Rae, Razvan Pascanu, Çaglar Gülççhrre Siddhant M. Jayakumar, Max Jaderberg, Raphaël Lopez Kaufman, Aidan Clark, Seb Noury, Matthew M. Botvinick, Nicolas Heess, and Raia Hadsell. Stabilizing transformers for reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 7487- 7498. PMLR, 2020. URL http: http:::pppccceeeeeeieeppppp ptw.eediccceeeesseeeeeee/.... www.gg.mmm.....m... mww.ppresss........m.. wwtps/vll///////// http:///111/paaireeeee http///paaiiittaaaeeeeee html.
Jack Parker-Holdee Mingi Jiang, Michael Dennis, Mikayel Samvelyan, Jakob Foerster, Edward Grefenstette, and Tim Rocktäschel. Evolving curricula with regret-based environment design. In Proceedings of the International Conference on Machine Learning, pp. 17473-17498, PMLR, 2022. URL https::ppp....m........tt http://poocedenngeeee mlr wtw.://ssss.p.......eeemmm sevess/vlee 62 wwwpaaaaaooeeeeeee 2ttper222a..mlll...... html.
Eduardo Pignatelli, Jarek Liesen, Robert Tjarko Lange, Chris Lu, Pablo Samuel Castro, and Laura Toni. Navix: Scaling minigrid environments with jax. arXiv preprint arXiv:2407.193% 2024.
Lerrel Pinto, James Davidson, Rahul Sukthankar, and Abhinav Gupta. Robust adversarial rein- forcement learning. In Proceedings of the 34th International Conference on Machine Learning. volume 70 of Proceedings of Machine Learning Research, pp. 2817-2826. PMLR, 06-11 Aug 2017. URL https: : http///oceedigngee mlr wttp:/eeessss///tttoooooo
Antonin Rafin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann.
Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268):1-8, 2021. URL httt:://imlll. wttpllrrrrggpappps/SSee 100-364410 html.
Aravind Rajeswaran, Vikash Kumar, Abhishek Gupta, Giulia Vezzani, John Schulman, Emanuel Todorov, and Sergey Levine. Learning complex dexterous manipulation with deep reinforcement learning and demonstrations. arXiv preprint arXiv: 1709. 1700: 2017.
Sharath Chandra Raparthy, Eric Hambro, Robert Kirk, Mikael Henaff, and Roberta Raileanu. Gen- eralization new sequential decision making tasks with in-context learning. In Forty-first In- ternational Conference on Machine Learning, ICML, 2024, Vienna, Austria, July 21-27, 2024. net
OpenReview.net, 2024. URL https: wttpeerrrrrw..ettttor.eeeeeeeeee wtt.m/id--////00
Scott Reed, Konrad ZoIna, Emilio Parisotto, Sergio Gómez Colmenarejo. Alexander Novikov, Gabriel Sarth-maron, Mai Giménez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, Tom Ec- cles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, and Nando de Freitas. A generalist agent. Transactions on Machine Learning Research, 2022. ISSN 2835-8856. URL https: http:////ppeeeeeeeeeeeee...o..e. net. forum? id-likKOkHjvj.
Alexander Rutherford, Benjamin Ellis, Matteo Gallici, Jonathan Cook, Andrei Lupu, Gardar Ing- varsson, Timon Willi, Akbir Khan, Christian Schroeder de Witt, Alexandra Souly, et al. Jaxmarl: Multi-agent environments in jax. arXiv preprint arXiv:2311.10090, 2023.
Alexander Rutherford, Michael Beukman, Timon Willi, Bruno Lacerda, Nick Hawes, and Jakob Foerster. No regrets: Investigating and improving regret approximations for curriculum discovery. arXiv, 2024. URL https: / /aaxiviivvv.////////aa/////aatva
wwwpwwwwwww.aaa/a/e/r/reaaaaaraa wwtp:b/////aaeeee.e..oo.e 2408 15099.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. CoRR, abs/1707.0634$1 2017. URL 1707.01 06347.
org/ /bs/
Carmelo Sferrazza, Dun-Ming Huang, Fangchen Liu, Jongmin Lee, and Pieter Abbeel. Body trans- former: Leveraging robot embodiment for policy learning. 2024, URL https: ///oi.oorrtt 10. 48550/8555 aariii... 2408. 06316
17

---

Published as a conference paper at ICLR 2025
Olivier Sigaud, Gianluca Baldassarre, Cédric Colas, Stéphane Doncieux, Richard J. Duro, Nicolas Perrin-Gilbee and Vieri Giuliano Santucci. definition of open-ended learning problems for goal-conditioned agents. CoRR, abs/231 2023, doi: 10.48550/ARXIV2233 ht.tps: //oii .rrg /t/.///./// 48500 rxi 2311. 00.34..
URL
Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with context- based representations. In International Conference on Machine Learning, pp. 9767-9779. PMLR, 2021.
Ghada Sokar, Rishabh Agarwal, Pablo Samuel Castro, and Utku Evei. The dormant neuron phe- nomenon in deep reinforcement learning. International Conference on Machine Learning, pp- 22145-32166 PMLR, 2023.
Lisa Soros and Kenneth Stanley. Identifying necessary conditions for open-ended evolution through the artificial life world of chromaria. In Artificial Life Conference Proceedings, pp. 793-800 MIT Press One Rogers Street, Cambridge, MA 22142-1209, USA journals-ifo 2014.
Kenneth O Stanley. Why open-endedness matters. Artificial life, 25(3):232 235, 2019.
Fan-Yun Sun, S. I. Harini, Angela Yi, Yihan Zhou, Alex Zook, Jonathan Tremblay, Logan Cross, Jiajun Wu, and Nick Haber. Factorsim: Generative simulation via factorized representation. 2024.
URL https: /twi. org/10 48550/ www....ww..............
Adaptive Agent Team, Jakob Bauer, Kate Baumli, Satinder Baveja, Feryal M. P. Behbahani, Avishkar Bhoopehand, Nathalie Bradley-Schmiee Michael Chang, Natalie Clay, Adrian Col- lister, Vibhavari Dasagi, Lucy Gonzalez, Karol Gregor, Edward Hughes, Sheleem Kashem, Maria Loks Moks-Thomppon, Hannah Openshaw, Jack Parker-Holder, Shreya Pathak, Nicolas Perez Nieves, Nemanja Rakicevic, Tim Rocktäschel, Yannick Schroecker, Jakub Sygnowski, Karl Tuyls, Sarah York, Alexander Zacherl, and Lei Zhang. Human- adaptation in an open-ended task space. CoRR, abs/2301.07608, 2023. doi: 10448550VaXXx:2311.77633 URL 48550/ arxiv 2301 07608.
http:://doi.rr///
Octo Model Team, Dibya Ghosh, Homer Walke, Karl Pertsch, Kevin Black, Oier Mees, Sudeep Dasari, Joey Hejna, Tobias Kreiman, Charles Xu, Jianlan Luo, You Liang Tan, Lawrence Yunliang Chen, Pannag Sanketi, Quan Vuong, Ted Xiao, Dorsa Sadigh, Chelsea Finn, and Sergey Levine. Octo: An open-source generalist robot policy. CoRR, abs/2405, 12213, 2024, doi: 10.48550/ ARXIV.2405.12213 URL https: / http:///.oo.......///./o 4854445554404 arxiv. 12201..221..
Open Ended Learning Team, Adam Stooke, Anuj Mahajan, Catarina Barros, Charlie Deck, Jakob Bauer, Jakub Sygnowski, Maja Trebacz, Max Jaderberg. Michaël Mathieu, Nat McAleese, Nathalie Bradley-
Nathaniel Wong, Nicolas Porcel, Roberta Raileanu, Steph Hughes- Fitt, Valentin Dalibard, and Wojciech Marian Czarnecki. Open-ended learning leads generally capable agents. CoRR, abs/2107.12808 2021. URL https: /tatta://arxiv.... arxiv. org/ abs/2107. 12808.
Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojeiech Zaremba, and Pieter Abbeel. Do- main randomization for transferring deep neural networks from simulation to the real world. In International Conference on Intelligent Robots and Systems, pp. 23-30. IEEE, 2017. doi: 10. 1109/RRO.20077.0 82021133. URL https: /ww. org /000000.0 1109/111 IROS 2017 8202133.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based con- trol. In International Conference on Intelligent Robots and Systems, pp. 5026-5033. IEEE, 2012. doi: 10. 00.1109/RROS.2012.3366 2012. 6386109.
URL https: //tp///ooio/ org. 1009//1011I.....
Georgios Tzannetos. Bárbara Gomes Ribeiro, Parameswaran Kamalaruban, and Adish Singla. Prox- imal curriculum for reinforcement learning agents. Trans. Mach. Learn. Res., 2023, 2023, URL .et
Dani Valevski, Yaniv Leviathan, Moab Arar, and Shlomi Fruchter. Diffusion models are real-time game engines, 2024. https: peattts://arx// //axxivv www.iiiooorraaaaoooo www./oog////o.....ooooo..o..s. hbcr//bs//444/ /2008 14837.
18

---

Published as a conference paper at ICLR 2025
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances Neural Information Processing Systems, pp. 5998 6008, 2017. URL https: http://poceeeegg..ppeccepoo neurips ec/ dee 91 fbd053cl c4844 Abst ratstttll..... htmll.
waaaassaffeee3555// htteer//00//aash/fffae4433333 paper/2017. hta////ahh//////
Oriol Vinyals, Igor Babuschkin, Wojeiech M Czamecki, Michaël Mathieu, Andrew Dudzik, Juny- oung Chung, David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. nature, 575(7782):350-354 2019.
Guanzhi Wang, Yuqi Xie, Yunfan Jiang, Ajay Mandlekar, Chaowei Xiao, Yuke Zhu, Linxi Fan, and Anima Anandkumar. Voyager: An open-ended embodied agent with large language mod- cls. Transactions Machine Learning Research, 2024. ISSN 2835-8856. URL https:
Rui Wang, Joel Lehman, Jeff Clune, and Kenneth O. Stanley. Paired Open-Endee Trailblazer (POET): Endlessly generating increasingly complex and diverse learning environments and their solutions. CoRR, abs/1901.011555 2019. URL http: http///axvvv... http://iii....aaa//a/a/aaaaaaaaa /100//190020 1111199111100 01753.
Tingwu Wang, Renjie Liao, Jimmy Ba, and Sanja Fidler. Nervenet:
Learning structured policy with graph neural networks. 6th International Conference on Learning Representations. OpenRe- view.net, 2018 URL https: / wpwp//pppee.ee..ee//////. net
Zihao Wang. Shaofei Cai, Anji Liu, Xiaojian Ma, and Yitao Liang. JARVIS-1: Open-world multi- task agents with ing in Open-Endedness Workshop, 2023. URL https: xzPkZyHlon
multimodal language models. Second Agent Learn-
https::/opeereriieeeeeeee /wpp///eepeeeeeeeeeee/ee net. net/forummid-
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances Neural Information Processing Systems, 33:14129-14142. 2020.
19

---

Published as a conference paper at ICLR 2025
APPENDIX
We structure the appendix as follows:
Appendix A describes the mathematical and computational logic behind Jax2D and Ap- pendix B performs speed tests on it.
Appendix C provides further details of the Kinet ix RL environment, while Appendix D shows examples of randomly generated levels.
Appendix E lists the hand-designed holdout levels and Appendix F shows example mor- phologies used in Figure 4.
Appendix describes the different network architectures 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 I. compares UED methods.
Appendix M performs small ablations study where we try removing aspects of our general agent training pipeline.
Appendix N compares the learnability of chosen vs randomly sampled environments over the course of training.
Appendix O Ablates the observation and action spaces of K intiix
Finally, Appendix Q briefly investigates lifelong learning aspects of the general agent.
A
JAX2D
This section provides an in-depth look into the logic behind Jax2D. Jax2D largely owes its heritage Box2D (Catto, 2007) and ImpulseEngine (Gaul, 2013), with most of the underlying framework being lifted from these engines and adapted for JAX. For more thorough account of some of the concepts behind rigid-body physics, recommend Erin Catto's talks.
A.1
CORE ENGINE
The main loop of Jax2D is summarised in Algorithm 1. Each part of the engine is subsequently explained as referenced.
Algorithm Jax2D main engine loop.
1: while true do
2: Apply gravity
3: Calculate collision manifolds (Appendices A.3.1, A.3.2, A.3.3 and A.3.4)
4:
Apply motors (Appendix A.5)
5:
Apply thrusters (Appendix A.6)
6:
if warm starting then
7:
Apply warm starting collision impulses (Appendix A.7)
8:
Apply warm starting joint impulses (Appendix A.7)
9:
end if
for 1=1 to num num_solver_stepy do
11:
Apply joint constraints (Appendices A.2 and A.4) Apply collision constraints (Appendices A.2 and A.3.5)
12:
13:
end for
14:
Euler step position and rotation
15: end while
htttps:/bbo2d..rgg/ i coss/
20

---

Published as a conference paper at ICLR 2025
A.2 IMPULSE RESOLUTION AND CONSTRAINT SOLVING
The core of Jax2D is impulse resolution, in which equal and opposite impulse is applied to pair of shapes order satisfy some constraint, For a given impulse j, the positional and angular velocities of shape affected follows.
m
1
where u is positional velocity, m is mass, w is angular velocity, r is the displacement from the centre of mass of the shape to the position the impulse being applied at and / is the rotational inertia.
We use x 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 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, con- taining the information shown in Table 1.
Table 1: Collision Manifold Specification
Attribute
Symbol
Data Type
Description
Position
p [float, float] [float, float]
Global position of the collision. Normalised vector along which the collision occurs. Positive penetration indicates an active collision.
Normal
Penetration
P
float
The resolution of collision takes place in two steps. First collision manifold is generated. This dependent on the exact shapes that are colliding (e.g. the logic for deriving a collision manifold between 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.) 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:
p Pa ft.
(3)
pb Pa
(4)
Pb
Pal
p To + Tb 1pb
Pal
A.3.3 POLYGON-CIRCLE COLLISION MANIFOLDS
The collision between polygon and circle b is calculated by first determining the closest point on any edge to the circle. For each edge, the centre of the circle 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 will instead be clipped to corner. Once this closest point p has been found, the collision manifold can
21

---

Published as a conference paper at ICLR 2025
be calculated.
ât + pb p
(6)
pb p
p t Tb pl
(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 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, then turned into an impulse that affects the two shapes. When two objects are deemed to have collided (i.e. 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 -. where is the restitution of the collision and v,, 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 would see perfectly elastic collision and the conservation of kinetic energy.
We first note that the velocity of point on object can be calculated by
e + X
where e is the velocity of the objects centre of mass, the angular velocity and is the point on the object relative to the centre of mass. Given this, we can derive the required impulse to resolve collision between objects and is
in - -11 ( (Va + (wa X Ta) (ra xn)
(wb rb))) f
Ma + m2 -1
[n)
In
where e is the restitution, î is the collision normal, va and vn are the respective positional velocities, wa and W/ are the respective angular velocities, Ta and Tb are the respective relative positions of the collision from the centre of masses, m. and my the respective masses and Ia and In are the respective rotational inertias.
Intuitively, the numerator represents the change in speed 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 min(ea,es) 6b).
22

---

Published as a conference paper at ICLR 2025
A.3.6 FRICTION COLLISIONS
As well as the collision impulse which along the collision normal, we calculate friction im- pulse which acts perpendicular to it against the relative movement. This follows Couloumb' 's Law:
ji - jlin!
(10)
where is is the friction impulse, In is the normal impulse and μ is the coefficient friction. js is therefore defined, similarly to Equation (12),
-(E (va + (wa x ra) Ub
(ra xx)² (rax (wb x rb)))
is clip
(11)
ma.
m
where f the normalised vector perpendicular to the normal of the collision.
Similar to restitution, every shape has its coefficient of friction, with the coefficient for a colli- sion defined as f v 12 p
A.3.7 POSITIONAL AND VELOCITY CORRECTIONS
In simulation of infinite temporal granularity, impulses would be enough guarantee reliable behaviour. However, since in practice we must quantise our simulation into discrete timesteps, only using impulses to solve constraints causes compounding errors 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 velocity correction. Decomposing Equation (12) that the numerator defines the change in speed that will occur along the collision normal between the 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 ap where the penetration and a 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 3p, where 3 is 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 directly at odds with the collision constraint. Therefore, when we connect two shapes with 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 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 constant collision with restitution at the point of joining, with the collision normal pushing the joined positions back towards each other. As with collisions, also apply velocity and positional corrections.
A.4.2 FIXED JOINTS
Jax2D also faciliates 'fixed' joint, in which additional rotational constraint enforces that the relative angle between two shapes remains constant, fixing them together effectively into a single rigid body.
23

---

Published as a conference paper at ICLR 2025
The rotational constraint applies angular impulse around the fixed joint, defined as
Wa W/
Jr
(12)
La
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 7(0, - 8, 0y), where 8, and O, are the respective rotations of the two shapes, Of is the target rotation at which they have been fixed at and y is a coefficient 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 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 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 the relative angular velocity of the two shapes is already bringing them back into within their limits. This is to allow motors push joints back within limits potentially faster than the angular velocity correction would do.
A.5 MOTORS
A revolute joint can have 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, 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
jr p tanh ((wa w/ A)
(13)
where is the target speed of the motor, A is the action being applied the motor (by human or artificial agent), is the motor power and p is a coefficient to control to what degree the power wanes it approaches the target angular velocity.
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 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- 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 (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, can effectively solve constraints only over multiple
24

---

Published as a conference paper at ICLR 2025
solver steps but also over multiple timesteps, with little extra cost. This technique is referred as "warm starting'.
record accumulated impulses throughout the solver steps and also to match collision manifolds across timesteps. Jax2D takes the simple approach of naively matching collision manifolds across adjacent timesteps if collision does not oceur 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.
Warm starting requires
A.8 PARALLELISED COMPUTATION AND BATCHED IMPULSE RESOLUTION
As well as being able to easily parallelise multiple Jax2D environments with the Jax vmap oper- ation, we also parallelise many of the calculations within single environment, providing further speed increases. The calculation collision manifolds is easily parallelised, as they have no side effects. The application of motors and thrusters is also parallelised. A more nuanced parallelisation the constraint solving.
As discussed in Appendix A.7, solving constraint can affect (and even unsolve) other constraints. For this reason, solving constraints sequentially provides a greater efficiency terms of solver steps, 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 be- tween every pair of shapes must be solved every solver step, as we can't '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 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 there- fore also arises as a tuneable parameter that trades off between simulation speed and accuracy. We 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), the potential for speed improvements significantly less.
25

---

Published as a conference paper at ICLR 2025
B JAX2D SPEED RESULTS
Here we investigate the runtime speed both Jax2D and Ki netix. For all comparisons we use single NVIDIA L40S GPU, on a server with AMD EPYC 9554 64-Core CPUs. We first compare Jax2D against Box2D (Catto, 2007). 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 different approaches of comparing speed: The first is by simply running the engines, and applying fixed actions, giving us raw speed measure of each engine. In the second approach, we compare speed when running the RI. training loop, to have a more realistic estimate for speed during training. We use PureJaxRL-style training for Jax2D (Lu al., 2022) and Stable Baselines 3 (Raffin et al., 2021) for Box2D. We use the flattened symbolic representation for Jax2D and use networks for both Box2D and Jax2D.
comparably-
The results are presented in Figure 7 and Table 2. First, inside an RL loop, Jax2D always outper- forms Box2D, and shows improved scaling once the number of parallel processes greatly exceeds the number physical CPU cores. When comparing just the engine, Box2D outperforms Jax2D when using fewer than 1024 environments, which point Jax2D overtakes Box2D.
RL
Engine Only
Steps per second
107
10a
Box2D
103
Jax2D
10
10.
102
103
104
10s
10
10-
102
10a
104
105
Number parallel environments
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.
Table The best-case steps per second for both Jax2D and Box2D, an RL loop and outside. In raw performance, Jax2D's best case is approximately 4.5x faster than Box2D, and this increases more than x0x 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 Kinet ix (s, M and L), as well as the dif- ferent observation spaces. Speed predictably decreases as we increase the environment size. Using the pixel-based observation requires more memory. cannot run many parallel environments as with the other observation spaces. Symbolic-Enin -Entity does not scale as well as Symbolic-Flait likely due to saturating memory bandwidth.
For actual runtimes, training the generalist agent for billion timesteps on a single L40S took around 7 hours for S, hours for M and 14 hours for L. Training such large number of timesteps is indeed nontrivial, but JAX and our Jax2D engine makes feasible. This could further be sped up by using multiple GPUs in parallel.
26

---

Published as a conference paper at ICLR 2025
10
Ridses
10
RL Symostt Flus RL-Symazii--mee Eniitt
brgans
10
Figure 8: The number of steps per second (SPS) in Kinetix for a variety of observation spaces. Symbolic-Fntity what we in experiments, while Symbolic-Flat is flattened (and therefore not permutation invariant) representation.
C KINETIX: FURTHER DETAILS
C.I ENVIRONMENT CLASS SIZES
The environment sizes we use are detailed in Table 3. Note that every level in Kinet x contains 4 large fixated polygons (floor, ceiling, left wall, right wall).
Table 3: The size of each environment class.
Entity
Small
Medium
Large
Polygons
6
12
Circles
3
Joints
2
Thrusters
2
Thruster Joint
4
4
4
Thruster Bindings
2
C.2 OBSERVATION SPACES
Kinetix allows for three observation spaces: Symboli -ntt Symbol ic-Flat. and Pixels. Both the symbolic observations common representation for shapes Table 4, joints Table 6 and thrusters Table
For use in Symbolic CumboliccEntity. Entity, construct 2 entities per joint: 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.
27

---

Published as a conference paper at ICLR 2025
Table 4: Information provided for shapes
Name
Dimensions
Position
Velocity
2
Inverse Mass
Inverse Inertia
1
Density
tanh (Angulr Velocity/100
OneHot (Role)
Ryyles
sin(Rotation)
cos(Rotation)
Friction
Restitution
OneHot (Shape'Type)
Thypes
Radius (only for circle)
Vertices (only for polygons)
8
TriangleOrRectangle (only for polygons)
2
Table 5: Information provided for thrusters
Table 6: Information provided for joints
Name
Dimensions
Name
Dimensions
Active
Active
Relative Position
IsFixed
Power
Relative Position w.r.t. from
2
sin sin(Rotation)
Relative Position w.r.t. to
2
costatttioon)
Motor Power
Motor Speed
Motor Permanently On
One Colour)
Ticolours
sin (Rotation)
cos (eoootttion)
28

---

Published as a conference paper at ICLR 2025
D RANDOMLY GENERATED LEVELS
We show 24 example random levels for size S (Figure 9), M (Figure 10) and L (Figure 11).
Figure 9: Randomly generated filtered levels from the DR distribution (S).
29

---

Published as a conference paper at ICLR 2025
Figure 10: Randomly generated filtered levels from the DR distribution (M..
30

---

Published as a conference paper at ICLR 2025
Figure 11: Randomly generated filtered levels from the DR distribution (L).
31

---

Published as a conference paper at ICLR 2025
E HAND-DESIGNED LEVELS LISTING
In this section we provide plots of the handmade levels. Figures 12 to 14 contain the full holdout sets for each environment size, respectively. We note that darker colour indicates that shape is fixated, i.e. that it has infinite mass and cannot move.
3-Car-One-kheal
S-Spin-Fall
S-Thruster-Aim
8-Thruster-Ovon
S-Thruster-Point
6-Tbruster-Up
8-Tbruster-Kait
S- monicycle-Balancee
s-Unicycle-Laat
s-Unicycle-Kight
Figure 12: Handmade levels (S).
32

---

Published as a conference paper at ICLR 2025
X- A-nn-laadd
N-Arn-Right
H- Ar-r--UUp
X-Car-Flip
X-Car-Left-Ban
H-Car-Right-2il
N-Car-Thrust
X-Catapult
X-Contrrapionn-Ke
X-Flinger
x-Tbrust-Cirele-Ovee
#-Turust-Control
N- H-Tbruster-Blec
Trrrutter-Coorrll coate--ooorrrttttteeo -aight
M-Thruater-Pal
M-Thruster-Pi Paitterr--all-Rhht 1-8ight
M-Thruster- Tof---at
T-Thrustor-Top-latte t- t-Baay
--Thruster-Top-Rkpe
M-Unicycle
Figure 13: Handmade levels (M).
33

---

Published as a conference paper at ICLR 2025
Letathweurd
Lntapteria
T
Urticcttenden Aaulle
Figure 14: Handmade levels (I.).
34

---

Published as a conference paper at ICLR 2025
F RANDOMLY GENERATED 3-SHAPE MORPHOLOGIES
Figure 15 shows a sample of the randomly-generate morphologies used for the analysis in Sec- tion 5.1.
Figure 15: Randomly generated 3-shape morphologies.
35

---

Published as a conference paper at ICLR 2025
G FURTHER NETWORK ARCHITECTURE DETAILS
We use the same actor-critic architecture for each observation space, consisting of five fully con- nected layers, of width 128, and tanh activation. However, how the input this network is obtained differs for each observation space. Since the environment is fully observable (except in the case of Pixels), we do use a recurrent network.
Pixels Inspired by the IMPALA architecture (Espeholt et al., 2018). we use two convolutional layers process the 125 125 observation. The first has 16 channels, a size of x 8 and stride of x while the second has 32 channels, size of x 4 and a stride of 2 2. The result of these layers is flattened before being passed to the main actor-critic network.
Symbolic-Flas -lat The Symbolic- Flat encoder is simply feed forward network with width of 512.
36

---

Published as a conference paper at ICLR 2025
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
7
0.995
AGAE
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
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
Transformer layers
Transformer Encoder Size
128
Transformer Size
16
Number of heads
8
SFT.
Batch Size N
12288
Rollout Length L
512
Update Period T
128
Buffer Size K
1024
Sample Ratio
0.5
37

---

Published as a conference paper at ICLR 2025
1 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 and In Figures 16 18, we plot the performance of the agents trained for Figure on each individual holdout level. We note that the fine-tuning base model is one trained on for 5B timesteps.
1.0
Inquaenall
0.3
0.2
1.0
0.9
0.2
1359
13TM
130
15CH
Environment Interactions
Figure 16: Specialist Agents on S.
38

---

Published as a conference paper at ICLR 2025
0.0
0.0
0.2
Talai Raod
1.0
0.4
0.0
1.0
Conteptiis
Tteat
0.0
30.0
0.2
0.0
0.0
0.2
1.0
E-trane Jol--att
Tiraaer Fall Hagh:
Deester Shoet
Desster Cop Lett
0.0
1.0
0.
0.0
02
0.
Environment Interactions
Figure 17: Specialist Agents on M.
39

---

Published as a conference paper at ICLR 2025
0.2
Tabus Fare
Rie-lined
1.0
1-0e hipe
Loter being
Eourreate Toree
Cerecle Tear tard
Cureace Reate Hard
0.0
1.0
I-trappe
0.0
0.0
0.4
0.2
0.0
Sabrry
0.2
1.0
1-0
Loch Moreeiee
Leffal
0.5
1.0
0.8
0.0
0.4
0.2
0.0
0.0
Environment Interactions
Figure 18: Specialist Agents on L.
40

---

Published as a conference paper at ICLR 2025
J GENERAL AGENT RESULTS BY HOLDOUT LEVEL
Next, we plot the performance of SFL and DR on individual levels, with results Figures 19 to 21. We that, generally, there is an upwards trend in the performance on most levels, but this is monotonic. Additionally, on some levels (e.g. M M-Thrust CerrustCConte Conttoo), performance decreases over training, potentially indicating bias in the levels trained on.
S:Candne--thhe
S-TrasdenOurr
STreucte-Pert
1.0
0.0
0.0
SHL
Nantim
02
5:Truake1.p
S-rruatt-WWa
S-Uricyse-tuanc
SUNcyceLott
S-incyclePigtt
1.0
0.0
0.0
15
45
15
26
45
15 28 33 45 15
15
Environment Interactions
Figure 19: DR SFL the full set of levels.
41

---

Published as a conference paper at ICLR 2025
5F1
01 Santom
0.8
0.0
3000
0.0
N-CaefIp
M-CaninB-Cany
1.0
0.2
0.0
N-Coffausien
Ktinggt
1.0
0.8
0.0
MrThrud-Centnl
ArThestte-oommtttn
0.0
0.4
0.2
M-Thruater Nel Right
MiThruster Shrett
N-Thraisrtt-eer Left
1.0
0.2
0.0
1.0
0.0
Brvironminnt Interactions
Figure DR vs SFL on the full set of levels.
42

---

Published as a conference paper at ICLR 2025
5PL
D1
0.0
Santom
0.2
0.00
Lo Hight
LGa Swiny doourd
Catpplo
Laraal Whoch Hard
1.0
0.0
0.2
0.0
L.Cpppt
Lande
LihaniTeus
L-GIaser
0.4
L-GraspeHaad
Utony
Mo
0.0
0.0
0.0
L-MEsCCoHoqppr
L-H LnCo
1.0
0.0
0.2
0.0
-.kgscc-Wbbe-liig
1.2
iPtrbll
L-Pattsmee)
1.0
0.4
L-Fartsmmer3
Lhat
L-2nho-4pp
0.4
0.3
0.0
02
0.0
Bivironment Intercctinss
Figure DR vs SFL on the full set of levels.
43

---

Published as a conference paper at ICLR 2025
K FURTHER GENERAL AGENT RESULTS
Figure 22 contains the performance of DR and SFL on each environment size. We can see that, every case, the agent's performance increases throughout training, indicating that learning 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 M and L, the agent is still able to zero-shot number of the levels.
Next, in Figure 24, we evaluate on a fixed set of randomly-generraed levels of the appropriate size. This 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 could expect further improvements by training this agent for longer.
M
0.8
Solve Rate
0.6
0.4
SRL Random
0.2
0
18
28
3B
4B
5B
1B
28
3B
4B
58
0
1B 28
3B
4B
5B
Environment Interactions
Figure 22: Results for DR and SFL for S, M and L,, respectively, In each pane, the training levels sampled from the DR distribution 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.
44

---

Published as a conference paper at ICLR 2025
0.8
0.6
Solve Rate
Solve Rate
Solve Rate
0.4
1B
2B
3B
4B
5B
18
2B
3B
4B
5B
1B
2B
3B
4B
5B
Train(M)[Eva((S)
Train(M)[Eval(M)
(M)Ea(())Evai(L)
0.8
0.6
0.4
Random
0.2
2B
3B
4B
Evaa(S))
Train(L)JEva((M)
Eraan(L)
0.8
0.6
0.4
0.2
Environment Interactions
Figure 23: Each row corresponds to the same agents, evaluated on each holdout set.
M
Solve Rate
0.4
DR
0.2
28
1B
Environment Interactions
Figure 24: Performance of the agents trained for Figure fixed of 1000 randomly-generated levels for each size.
45

---

Published as a conference paper at ICLR 2025
$
M
0.325
0.14
Solve Rate
0.20
0.300
0.12
0.275
0.15
0.10
0.250
0
9B
6B
12B
9B
12B
Environment Interactions
Figure 25: Performance of single seed, trained L, random levels.
3
Solve Rate
0.4
0.3
0.125
0.3
0.100
0.2
0.2
0.075
0.1
0
3B
6B
9B
12B
0
3B
6B
9B 12B
0
3B
6B
9B 12B
Environment Interactions
Figure 26: Performance a single seed, trained on L, on the holdout set of levels.
46

---

Published as a conference paper at ICLR 2025
L UED RESULTS
In Figure 27, we present results for two popular UED methods, PLR (Jiang et al., 2021b;a) and ACCEL (Poderr-Holder et al., 2022), with the hyperparamters listed in Table These results show neither method significantly outperformed DR, leading us to focus solely on SFL in the main text.
0.12
Solve Rate
0.3
0.10
ACCEL
0.2
0.06
18
28
38
48
5B
Environment Interactions
Figure 27: Solve rate on (left) M and (right) evaluation sets for PLR, ACCEL and DR.
Table 8: UED Hyperparameters.
Parameter
Value
PLR
Replay rate, Buffer size, K
0.5
8000
Scoring function Prioritisation
MaxMC
Rank
Temperature,
1.0
staleness coefficient
0.3
Duplicate check
ACCEL
Number of Edits
3
47

---

Published as a conference paper at ICLR 2025
M ABLATIONS
SFL
0.35
DR
SFL (No Filter)
Solve Rate
0.30
SFL (Symbolic-Flat)
0.25
0.20
0.15
0
1B
2B
3B
4B
5B
Environment Interactions
Figure 28: The solve rate on for different ablations. SFL denotes the training regime we used in the main text.
We perform ablations 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 re- moving the filtering, and find that there 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-Fles cantly worse than Symbolic-Entite likely due to the large number of symmetries inher- ent in the environment. Despite this, SFL with Symbolic-Flat does outperform DR with Symbolic iimbolicEEnttty.
large difference in performance for SFL. Secondly,
representation, and find that the performance of this method is signifi-
N LEARNABILITY OVER TRAINING
In Figure 29, we plot the learnaillity of the training levels, and the larger set of random levels these are sampled from. Overall, we find that random levels tend to have low Icarnability, whereas the top 1024 levels consistently has high learnability throughout training.
5
M
Solve Leamabllity
1.00
0.75
Tra ining Levels Random Levels
0.50
0.25
Environment Interactions
Figure 29: The learnability (scaled to between O and 1) over training for each of the environment sizes. We compute Icarnability for 12288 randomly-generann levels (shown in orange), and select the top 1024 of these (the learnability of this top subset is shown in blue).
O ALTERNATE OBSERVATION AND ACTION SPACES
In Figure 30, we consider multi-task setting, where we train agents on the holdout tasks with each combination of MuuliiiscreteeConnimooo action space and Pixels/Symbolic-Enie Entitty observation space. We find that MultiDiscrete actions outperform Continuous actions, and that Entity slightly outperforms Pixels (in addition to being significantly faster to run). This experiment validates our decision to Entity observations and MulriDiscrete actions for our main experiments.
48

---

Published as a conference paper at ICLR 2025
0.7
Solve Rate
0.6
Entity-MultiDiscreete
0.5
Entity-Continuous Pixels-MultiDiscrete Pixels-Continoous
0.4
0
100M
200M
Figure 30: We compare the different observation and action spaces on a multi-task setting, Here trained agents on all 74 holdout tasks, and show aggregate performance on this set of tasks. We plot mean and standard deviation over five seeds.
P CROSS-EMBODIMENT LEARNING
In Figure 31, we consider cross-cmbodiment learning, where we train single agent to control all of the Mujoco recreations. We compare this against agents trained individually for each task. For fairness, we allocate more samples to the single agent (500M vs 100M). We find that the single agent able to competently control morphologies, although less sample efficient when considering only single task. On some tasks (e.g. MuJoCo-Wand -aaoeerr we see improved learning from co-training with other morphologies.
39.75
0.23
Indivisual
Combined
Figure 31: Comparing the performance of agents individually trained against one jointly trained on the recreations of Mujoco tasks. Combined indicates the agent trained jointly, and all plots show mean and shade standard deviation over seeds. We note that the x-axis measures the total number of timesteps, i.c., for the Combined line, this includes all morphologies.
Q LIFELONG LEARNING
In Figure 32, we plot single training run where we first train agent random levels from the distribution for 5B timesteps. We then change this and train the agent on random M levels for timesteps and finally train it again random S levels for IB timesteps. We plot the performance of the agent on the heldout set of levels for both the and M size separately. As expected, training initially slightly improves performance timesteps on M improves performance by a larger margin. Going back to training on random S levels reduces the performance the M holdout set. This indicates level of forgetting plasticity loss in the agent.
the set of holdout levels. Then, training for 1B
49

---

Published as a conference paper at ICLR 2025
1.0
S
M
S
0.5
0.0
0
1B
2B
3B
4B
5B
6B
7B
Environment Interactions
Figure 32: A single run's training, where we first train the agent on S for 5B timesteps, then transition for IB and finally train on again for IB. We plot the aggregate evaluation performance on the S set of holdout levels on the left and the M holdout levels on the right.
50


Reviews:
[{"review_id": "FOpRu2yA7I", "reviewer": "Reviewer_pXuh", "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.", "questions": "See Weaknesses section. Additionally,\n- 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?\n- 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.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 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. \n- 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. \n- 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.", "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?\n- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xIQNBiOQSI", "reviewer": "Reviewer_VVCK", "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.", "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?\n\n2. Does Kinetix support parallel pixel rendering, such as for vectorized image-based experiments?\n\n3. Is SFL only choosing between the S, M, and L levels?\n\n4. Are the inputs into the policy purely one-hots? (One-hot encoding for each of the polygons, thrusters, and shapes.)\n\n5. 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.\n\n6. 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?\n\n7. How many environments were used during training, and how many environments were held-out for zero-shot evaluation?\n\n8. What was the generalist policy in Figure 5 trained on? A distribution of environments over all 4 tasks at their hard level?\n\n9. 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)?\n\n10. 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?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 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.\n\n2. Paper is clearly written.", "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?\n\n2. 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?\n\n3. 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.\n\n4. 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.\n\n5. There were no supplementary materials submitted, which would have been a good opportunity to show video rollouts of the trained agent in action.\n\n6. 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.\n\n7. 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.”\n\n8. Minor: I think the wrong citation was provided for MJX in Section 7 (that work seems to be Mujoco).\n\n9. Minor: Experiments would benefit from some comparison to prior approaches/architectures, though this is less important given this is mainly a systems/framework paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SJmSV8iRqK", "reviewer": "Reviewer_NaNR", "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.", "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?\n- 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?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 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 \n- 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\n- Have interpretable/handmade levels to understand the performance on different sizes/difficulties of tasks.", "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.  \n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "risOryPW97", "reviewer": "Reviewer_ppmD", "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.", "questions": "They are mentioned in the Weaknesses.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- The paper is well-written, organized, and straightforward.\n- Extensive testing across various task complexities validates its robustness in diverse 2D environments.\n- This paper has strong potential to serve as a valuable benchmark for future research.", "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.\n\n- **Filtering out:** The authors mention that trivial and unsolvable levels are filtered out. What quantitative metrics were used to determine this filtering.\n\n- **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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)