Two RL Agents Team Up to 'Chopstick-Catch a Rocket': Mechazilla Captures Super Heavy Booster in Practice

From Falcon 9 to Super Heavy: A More Complex RL Challenge

A few days ago, a developer shared on Reddit a project using a PPO agent to land a Falcon 9 on a droneship. Now he's escalated the problem: getting two interdependent reinforcement learning agents to collaborate on the task of a Super Heavy booster being captured by the launch tower's mechanical arms (i.e., "Mechazilla"). One agent flies the booster, the other controls the tower arms and carriage. The entire simulation runs in the browser, letting you watch both agents perform in real time.

According to the author, the training environment uses PPO implemented in JAX, running on a single GPU, with heavy domain randomization: wind, gusts, engine failure, reignition failure, sensor noise and latency, mass and thrust variations, and limited propellant. Each attempt starts from several kilometers up with engines off.

The Core Difficulty of Dual-Agent Training: You Can't Start Both From Scratch

The author quickly found that the two agents cannot be trained from scratch simultaneously. The tower arms get no learnable signal when the booster never appears, and the booster never receives a capture reward when the arms never close. This is a classic sparse-reward multi-agent coordination problem.

His solution is staged training: first train the booster with a simple rule-based tower (closing the arms when the booster is nearby); then freeze the booster and train the tower agent from scratch; finally fine-tune the booster to work with the learned tower. This curriculum learning strategy is the key that made the project work.

Stage 1: Booster + Rule-Based Tower

Even with the rule-based tower's help, the booster took a long time to learn the full approach. The author found that short rollouts (1.6 seconds) caused the agent to "forget" how to capture, because a capture requires several seconds of slow, stabilized descent while the tower arms close. Extending the rollout to 6.4 seconds with λ = 0.99 solved this problem. This suggests that in long-horizon tasks, the rollout length must cover the full successful behavior cycle.

PPO (Proximal Policy Optimization) is a policy-gradient-based reinforcement learning algorithm that ensures training stability by limiting the magnitude of policy updates. Its core is a clipped surrogate objective that prevents the new policy from diverging too far from the old one. PPO is an on-policy algorithm, requiring trajectories from the current policy for updates. In this article, the rollout length (the length of each sampled trajectory) directly affects PPO's ability to model long-horizon tasks: if the rollout is too short, advantage estimation cannot cover the full successful behavior cycle, causing the policy to "forget" long-range dependencies. The author extended the rollout from 1.6 to 6.4 seconds precisely so PPO can observe the complete capture process and correctly assign rewards.

Stage 2: Training the Tower Agent From Scratch

After freezing the booster, the tower agent is trained from scratch. Initially it just flails its arms and keeps hitting the booster. The author also tried slowly fading out the rule-based tower, but training collapsed immediately once the assistance disappeared. After several fixes, the from-scratch tower ultimately won out, achieving even higher capture reliability than the rule-based tower.

Curriculum Learning is a training strategy that guides agents to learn by gradually increasing task difficulty or changing the environment distribution. In reinforcement learning, curriculum learning is often used to address sparse rewards and exploration difficulties. The staged training in this article is an example of curriculum learning: first let the booster learn basic approach under a simple rule-based tower, then freeze the booster to train the tower agent, and finally fine-tune jointly. This "easy-to-hard" curriculum design avoids the chicken-and-egg dilemma of starting both agents from scratch simultaneously. Notably, the key to curriculum learning lies in smooth transitions between stages; the author's failed rule-based tower fading strategy shows that directly changing the environment distribution can undermine the stability of learned policies, whereas staged replacement is more robust.

Stage 3: Booster Adapts to the Learned Tower

When retraining the booster to work with the learned tower, something unexpected happened: performance quickly deteriorated. The booster approached slightly too fast and slipped off the arms. Drastically lowering the learning rate solved the problem, and the pair now clearly outperforms the rule-based tower setup. This shows that during joint fine-tuning, a lower learning rate helps stabilize policy updates and avoid destroying already-learned delicate behaviors.

Sparse Reward refers to the situation where the agent only receives a reward signal upon completing the entire task (e.g., a successful capture), with no feedback on intermediate steps. In the Super Heavy capture task, success is only determined when the tower arms close and the booster hovers stably, so the agent is extremely unlikely to trigger this reward by chance during random exploration, making learning highly inefficient. Sparse reward problems are typically mitigated through reward shaping, curriculum learning, or hierarchical reinforcement learning. In this article, the author uses a rule-based tower to provide dense "approach" signals early on, essentially transforming sparse rewards into learnable dense feedback through environment design—a common approach to solving sparse-reward multi-agent coordination.

Technical Details and Engineering Insights

The project uses PPO + JAX single-GPU training, with domain randomization covering the main sources of uncertainty in real rocket recovery. The real-time browser demo lowers the barrier for the public to understand RL. The author admits he is not an ML expert and especially welcomes feedback on the multi-agent part.

From an engineering perspective, this project validates several practical lessons: curriculum learning is necessary for sparse-reward cooperative tasks; rollout length must match the task's time scale; directly fading out an auxiliary rule-based policy is less robust than staged replacement; and using a smaller learning rate during joint training prevents policy collapse.

Domain Randomization is a technique for improving the robustness and generalization of reinforcement learning policies by randomizing environment parameters (such as physical properties, noise, latency, etc.) during training, so the agent learns to adapt to various uncertainties, thereby narrowing the sim-to-real gap. In this article, the author heavily randomized wind, gusts, engine failure, sensor noise and latency, mass and thrust variations, and limited propellant, covering the main disturbance sources in real rocket recovery. The cost of domain randomization is increased training difficulty, but it can significantly improve policy performance in unseen scenarios. For high-risk tasks like Super Heavy capture, domain randomization is an indispensable engineering tool.

Limitations and Open Questions

Currently, the two agents' collaboration still relies on staged training rather than end-to-end joint optimization. The rule-based tower's assistance is fully replaced in stage two, but the reason the fading strategy failed has not been analyzed in depth. The specific parameter distributions for domain randomization and details of the reward function design are also not elaborated in the original text. For multi-agent reinforcement learning researchers, these are directions worth further exploration.

The author's online demo is available at https://qforge-dev.github.io/hoverslam , along with a playground for the earlier Falcon simulation.

One-line Summary

A developer uses dual-agent reinforcement learning to simulate a Super Heavy booster being caught by mechanical arms in the browser, employing staged curriculum learning to solve the coordination challenge.

Paragraph Summary

This article presents a project that uses reinforcement learning to simulate a Super Heavy booster being captured by the launch tower's mechanical arms. Two PPO agents control the booster and the tower arms respectively, trained on a single JAX GPU with heavy domain randomization. The author found that the two agents cannot be trained from scratch simultaneously, so he adopted staged curriculum learning: first train the booster with a rule-based tower, then freeze the booster to train the tower, and finally fine-tune jointly. Along the way, he solved the forgetting problem caused by short rollouts and found that the learning rate must be drastically lowered during joint fine-tuning. The project validates practical lessons such as curriculum learning, rollout length matching, auxiliary policy replacement, and low-learning-rate fine-tuning, but collaboration still relies on staged training rather than end-to-end optimization, and some details remain undisclosed.