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

Authors: Michael Matthews, Michael Beukman, Chris Lu, Jakob Nicolaus Foerster
Primary area: reinforcement learning
Keywords: reinforcement learning, open-endedness, unsupervised environment design, automatic curriculum learning, benchmark

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.

Decision: accept (oral)
Reviewer scores: 8, 8, 8, 8

Reviews and discussion:
<thread>
creation date:
1730661382895

thread_type:
official_review

<thread_object>
creator: reviewer_ppmd

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

strengths:
- The paper is well-written, organized, and straightforward.
- Extensive testing across various task complexities validates its robustness in diverse 2D environments.
- This paper has strong potential to serve as a valuable benchmark for future research.

weaknesses:
- **Real World Tasks:** While this paper provides a strong foundation in 2D simulations, expanding its scope to assess the agent’s adaptability to real-world tasks, such as 3D simulations or complex dynamics as seen in [1,2], would enhance its practical relevance. Bridging this gap could amplify the study’s contributions, offering broader insights into real-world generalization and scalability.

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

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

questions:
They are mentioned in the Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730696524255

thread_type:
official_review

<thread_object>
creator: reviewer_nanr

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

strengths:
- Provide a highly efficient 2D rigid-body physics engine, leveraging JAX for scalable computation, with speedups of up to 30x when training an RL agent, allowing for 
- The learnt agent is highly effective at Zero-Shot transfer in the S and M levels that are held out, indicating the efficacy of pre-training on a wide set of procedural generation tasks. Additionally, show faster convergence/higher performance with this initialization
- Have interpretable/handmade levels to understand the performance on different sizes/difficulties of tasks.

weaknesses:
- The JAX2D environment seems to be somewhat limited in its expressivities, modeling only 4 unique entities, which may not transfer to a wide set of domains/tasks outside of the ones studied.  
- The task/reward function seems to be fixed across all environments to collide the green and blue shaped objects, while avoiding red shapes. Additional reward shaping seems to be needed for effective training, leading to some limited applicability of generating this data at scale for any set of tasks.

questions:
- For the environment generator, it is mentioned that there may exist unsolvable levels which automatic curriculum methods can filter out. Could you clarify what was done here?
- How does the choice of algorithm affect the performance in your benchmark. Do you anticipate releasing a dataset of transitions from Kinetics which can be used for offline 2 online RL?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730716355177

thread_type:
official_review

<thread_object>
creator: reviewer_vvck

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

strengths:
1. Introduces a physics engine that provides “almost entirely dynamically specified” scenes, where environments with different robot morphologies can be vmap-ed and run in parallel, which is not doable with prior Jax-based sim frameworks like Brax.

2. Paper is clearly written.

weaknesses:
1. All environments in benchmark must fall under the goal of making green shape touch blue shape without touching red shape. This seems to mainly constrain the problem to single-step tasks, where the reward of minimizing the distance from green to blue always incentivizes progress. Was this unified goal constraint purposefully imposed by design, or was it a constraint of Jax implementation, where the reward function for all environments must be the same to be parallelizable?

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

10. Heatmaps were an interesting way to convey the agent’s performance. Why not simply graph the agent’s x distance from the goal, with x-axis being the training iterations?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730720497538

thread_type:
official_review

<thread_object>
creator: reviewer_pxuh

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

strengths:
- Kinetix provides 66 hand-designed levels while having the option to edit tasks with a graphical editor or to randomly generate more levels with rejection sampling. 
- The unified goal and dynamics within all environments encourage policies to have physical reasoning capabilities instead of merely memorizing the solution for some particular task, which is a valuable objective for researchers to pursue. 
- Kinetix provides a way to generate unlimited environments and tasks with a unified goal, objects, and dynamics, which could be of interest to multiple research communities like generalist RL policy learning, meta-learning, world modeling, spatial understanding and physics reasoning, and so on.

weaknesses:
- The paper notes that as the generated environments increase in complexity, they may become unsolvable, which could contribute to the lower performance observed in the Large-level environments. If so, how does this impact the usability and interpretability of the benchmark results? To what extent does this affect the performance results reported in Figure 3?
- It is unclear whether the proposed benchmark supports visual observations, which are essential for training generalist policies and building agents that can operate in real-world settings.
- Although Kinetix can generate a vast range of environments, it is unclear how this benchmark would generalize to tasks or environments outside of its defined task distribution.

questions:
See Weaknesses section. Additionally,
- The range of the y axis for the four plots on the right in Figure 5 are missing. Are these also from 0 to 1?
- How long does training take for training on 1B Kinetix environments? Would be good to see if training on such a large number of environments itself would be a bottleneck for learning generalist agents.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731689814862

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response

comment:
We are grateful for the reviewers' thorough evaluations of our work, in particular for agreeing with us on the value of the Kinetix benchmark and its utility for investigating generalist RL training.

$\color{red}{pXuh}$: "contributes to the development of generalist RL agents".

$\color{green}{VVCK}$: "different robot morphologies can be vmap-ed and run in parallel, which is not doable with prior Jax-based sim frameworks".

$\color{cyan}{NaNR}$: "The learnt agent is highly effective at Zero-Shot transfer ... indicating the efficacy of pre-training on a wide set of procedural generation tasks".

$\color{magenta}{ppmD}$: "well-written, organized, and straightforward" and "strong potential to serve as a valuable benchmark for future research".

We are glad that reviewers appreciated the "novel" ($\color{red}{pXuh}$) and "highly efficient" ($\color{cyan}{NaNR}$) Jax2D engine we introduced that can `vmap` over different tasks, something which is "not doable with prior Jax-based sim" ($\color{green}{VVCK}$).

We appreciate the recognition of the results of our "highly effective" ($\color{cyan}{NaNR}$) generalist agent that underwent "extensive testing" ($\color{magenta}{ppmD}$) and has learned general strategies "instead of merely memorizing the solution for some particular task" ($\color{red}{pXuh}$).

Finally, every reviewer noted future work that the Kinetix environment could facilitate, including "meta-learning, world modeling" ($\color{red}{pXuh}$), "lifelong learning, transfer learning, and cross-embodiment learning" ($\color{green}{VVCK}$), "offline 2 online RL" ($\color{cyan}{NaNR}$) and that "this paper has strong potential to serve as a valuable benchmark for future research" ($\color{magenta}{ppmD}$).


# Additional Experiments

Reviewer $\color{green}{VVCK}$ recommended running additional experiments in cross-embodiment learning and lifelong learning. While these don't compose the main thrust of the paper, we thought they were useful in providing examples of the diverse range of research that can be done with Kinetix. We have run initial experiments in these domains and have added them to the updated manuscript (Appendices P and Q).

Reviewer $\color{green}{VVCK}$ also advised we investigate experiments with a continuous action space and pixel-based observations, as these results would "strengthen the argument of the paper". We found this argument compelling and have provided results in the new Appendix O. Reviewer $\color{red}{pXuh}$ also found it "unclear whether the proposed benchmark supports visual observations" - we hope that this is now clearer in the updated manuscript that Kinetix does natively support visual observations (rendered entirely inside JAX).

# Offline Dataset
As suggested by reviewer $\color{cyan}{NaNR}$, we added a dataset of transitions to the anonymous GitHub repository that can be used for behaviour cloning, or offline to online RL research. This consists of full trajectories from over 40000 distinct levels, totalling more than 10M transitions (and even larger datasets can be generated). We hope that this will enable future work in these areas using Kinetix.

# Reward Formulation

$\color{green}{VVCK}$ and $\color{cyan}{NaNR}$ highlighted our simple reward formulation as a potential downside of Kinetix. We summarise our general response to this (adapted from our individual response) as follows:

- We would argue that the green-blue-red goal formulation is actually very expressive, as evidenced by the diversity of tasks in our 66 handmade holdout environments. This formulation is generally referred to as a Reach-Avoid problem and has a rich literature behind it [1,2,3].
- Conditioning the general agent on a goal/reward function for each task would have massively complicated the learning process.
- Even with our current goal formulation, the generalist agent fails to make any progress on the more complex holdout tasks. We would argue that this implies the current formulation is complex enough to be useful to study.

We have added some additional text to the manuscript to clarify this.

# Videos of Trained Agents

Reviewer $\color{green}{VVCK}$ stated that they would like to see some videos of the trained agent in action. We completely agree that this is a useful resource and we encourage all the reviewers to take another look at the linked anonymous GitHub repository, which now contains more of these videos.

# References

[1] - Jaime F. Fisac, Mo Chen, Claire J. Tomlin, S. Shankar Sastry. Reach-Avoid Problems with Time-Varying Dynamics, Targets and Constraints.

[2] - Benoit Landry, Mo Chen, Scott Hemley, Marco Pavone. Reach-Avoid Problems via Sum-of-Squares Optimization and Dynamic Programming.

[3] - Kostas Margellos, John Lygeros. Hamilton-Jacobi formulation for reach-avoid differential games.
</content>
</thread_object>

</thread>


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

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

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