Training AI Rocket Landing from Scratch: A Practical Analysis of PPO Reinforcement Learning

This project is a complete reinforcement learning practical case that uses the Proximal Policy Optimization (PPO) algorithm to train a neural network in a Unity 3D environment to control a rocket for soft landing. Key highlights include two-stage training (behavioral cloning pre-training + PPO fine-tuning), reward engineering optimization, and realistic physics simulation. It adopts a Python+Unity hybrid architecture to solve cross-language communication and physics simulation issues, providing reinforcement learning practitioners with a complete reference from environment design to algorithm implementation.

Original Author and Source

• Original Author/Maintainer: MichaelLam71
• Source Platform: GitHub
• Original Title: rocket-landing-ppo
• Original Link: https://github.com/MichaelLam71/rocket-landing-ppo
• Release Time: 2025

Project Overview

This project uses the PPO algorithm in a Unity 3D environment to train a neural network to control a rocket for suicidal deceleration descent and soft landing. It demonstrates a two-stage process from behavioral cloning pre-training to PPO fine-tuning, solving core challenges such as sparse rewards, physics simulation, and cross-language communication.

Hybrid Architecture Design

Adopts a Python+Unity hybrid architecture: The Python side (PyTorch) handles machine learning logic, while the Unity side processes physics simulation and rendering. Real-time communication is via TCP socket (port 5005). Python sends 3 action values (thrust, RCS X/Z axis torque), Unity returns 17 observation values (15 states + reward + termination flag), and the reset signal is triggered by a special thrust value (-999).

Two-Stage Training

1. Behavioral Cloning Pre-training: Use a PID controller to generate successful landing data, and train the neural network to imitate expert behavior with MSE loss supervision, solving the sparse reward exploration problem.
2. PPO Fine-tuning: Load the pre-trained model and use a terminal reward mechanism (success +100~300, crash -100) to avoid hovering/crashing issues caused by dense rewards.

Observation and Action Space

• Observation Space: 15 normalized values (position, velocity, upward vector, angular velocity, vector to landing pad), clipped to [-5,5] to ensure input stability.
• Action Space: 3-dimensional continuous vector (main engine thrust [0,1], RCS X/Z axis torque), simulating real rocket attitude control.

Physics Simulation Details

Includes details such as fuel consumption (calculated based on specific impulse), air resistance (standard formula), and symmetric inertia tensor settings. Rocket parameters: dry mass 22000kg, fuel 2000kg, thrust-to-weight ratio 2.0, maximum thrust 470880N.

Implements the standard PPO algorithm with key components:

• Dual Network Architecture: Actor outputs a Gaussian distribution policy, Critic estimates the state value function.
• GAE Advantage Estimation: Calculate action advantage with lambda=0.95.
• Clipped Objective Function: Clip the probability ratio to [0.95,1.05] to prevent large updates.
• Entropy Reward: Avoid premature policy convergence.
• Gradient Clipping: Limit the magnitude of abnormal batch updates.
• Linear Learning Rate Decay: Gradually reduce the learning rate during training.

Network structure: 15 inputs → 256 hidden layers →256 hidden layers → outputs (Actor:3-dimensional action distribution, Critic:1-dimensional state value).

The project provides complete training visualization (behavioral cloning loss curve, PPO training curve, JSON logs). Results show: Direct PPO training fails due to sparse rewards; behavioral cloning pre-training provides an effective starting point, allowing PPO to optimize landing quality (softer, more upright) and generalize to more difficult conditions (tilt, position offset, initial velocity).

This project provides the following experiences for reinforcement learning practitioners:

1. Pre-training Value: In sparse reward tasks, behavioral cloning pre-training can significantly accelerate learning.
2. Reward Engineering: Dense rewards easily lead to unexpected behaviors; terminal rewards are more stable.
3. Physics Simulation: Appropriate realistic details provide richer learning signals.
4. Hybrid Architecture: Combine the advantages of Python ML ecosystem and Unity physics rendering.

For beginners, this is an excellent reference project with clear code, detailed documentation, and coverage of the complete process.