the last question
How to Land on the Moon
June 30, 2025
LunarLanderContinuous-v3 is a small simulated control problem. At each step, the lander receives eight numbers describing its position, velocity, angle, angular velocity, and leg contact. It controls a main engine and a side engine. Both controls are continuous values.
The simulator advances one step and returns a reward. A safe landing gives a high reward. Crashing, using fuel, and moving away from the landing pad reduce it. The policy is the function that converts the eight observed values into the two engine controls.
I trained that policy with proximal policy optimization, or PPO. The first stable policy learned to hover until the episode timed out. It used all its fuel, never landed, and still scored above 100.
The hover showed that the policy had learned basic control. It could stop the fall and remain stable. It had not learned the final descent.
Code.
Rollouts
PPO alternates between collecting experience and updating the policy. During collection, the policy does not change.
env_name = 'LunarLanderContinuous-v3'
state_dim, action_dim = 8, 2
num_envs = 8
The vector environment returns observations with shape (8, 8): eight landers with eight state values each. The policy returns actions with shape (8, 2). One call to env.step(actions) advances every simulation and returns the next observations, eight rewards, and eight terminal flags.
Before the policy sees an observation, I multiply its components by OBS_SCALE:
OBS_SCALE = np.array([10, 6.666, 5, 7.5, 1, 2.5, 1, 1], dtype=np.float32)
This changes the numerical scale of each input but adds no information. These constants are part of this implementation; Gymnasium only defines the observation components.
Actor and critic
The model contains two separate multilayer perceptrons. Both receive the same scaled observation \(s \in \mathbb{R}^8\), but they produce different quantities.
The actor maps the observation to a two-element mean vector \(\mu(s)\). This vector parameterizes the Gaussian distribution used to sample the two engine controls. The actor has eight hidden layers with ReLU activations.
The critic maps the observation to one scalar \(V(s)\). This value estimates the discounted future reward expected from that state under the current policy. The critic does not select actions. During training, generalized advantage estimation combines rewards with successive value estimates to determine whether each sampled action produced a better or worse outcome than expected.
Both networks are stored inside the ActorCritic module, but they do not share hidden layers. The critic has four hidden layers with ReLU activations.
class ActorCritic(nn.Module):
def __init__(self, state_dim, action_dim, hidden_dim, actor_layers, critic_layers):
super(ActorCritic, self).__init__()
actor = [nn.Linear(state_dim, hidden_dim), nn.ReLU()]
for _ in range(actor_layers - 1):
actor.extend([nn.Linear(hidden_dim, hidden_dim), nn.ReLU()])
actor.append(nn.Linear(hidden_dim, action_dim))
self.actor = nn.Sequential(*actor)
self.log_std = nn.Parameter(torch.zeros(action_dim))
critic = [nn.Linear(state_dim, hidden_dim), nn.ReLU()]
for _ in range(critic_layers - 1):
critic.extend([nn.Linear(hidden_dim, hidden_dim), nn.ReLU()])
critic.append(nn.Linear(hidden_dim, 1))
self.critic = nn.Sequential(*critic)
def forward(self, state):
action_mean = self.actor(state)
action_std = self.log_std.exp()
value = self.critic(state)
return action_mean, action_std, value
Action distribution
The actor produces action_mean, the mean engine command \(\mu(s)\) for observation \(s\). log_std is a learned two-element parameter. The forward method converts it to a positive standard deviation:
The policy constructs a Gaussian distribution and samples a raw action:
dist = torch.distributions.Normal(action_mean, action_std)
raw_action = dist.sample()
action = torch.tanh(raw_action)
In mathematical notation:
\[u \sim \mathcal{N}(\mu(s), \operatorname{diag}(\sigma^2)), \qquad a = \tanh(u)\]Here, raw_action is \(u\), and action is \(a\). The Gaussian can produce any real number. Tanh limits each engine command to \((-1,1)\).
The transformation changes the probability density. The implementation calculates the corrected log probability directly:
def tanh_log_prob(raw_action, dist):
action = torch.tanh(raw_action)
logp_gaussian = dist.log_prob(raw_action).sum(-1)
correction = torch.log(1 - action**2 + 1e-6).sum(-1)
return logp_gaussian - correction
This code implements:
\[\log \pi(a \mid s) = \log \mathcal{N}(u;\mu(s),\operatorname{diag}(\sigma^2)) - \sum_i \log(1-\tanh^2(u_i))\]logp_gaussian is the first term. correction is the second. The small value 1e-6 prevents a logarithm of zero when tanh is close to -1 or 1.
The rollout stores raw_action. During the PPO update, both policies evaluate that same sample. The correction cancels in their probability ratio, but it keeps the stored log probability equal to the density after tanh.
Policy update
class PPO:
def __init__(self, actor_critic, pi_lr, vf_lr, gamma, lamda, K_epochs, eps_clip,
batch_size, vf_coef, entropy_coef):
self.actor_critic = actor_critic
self.states, self.actions = [], []
self.pi_optimizer = optim.Adam(
list(actor_critic.actor.parameters()) + [actor_critic.log_std],
lr=pi_lr
)
self.vf_optimizer = optim.Adam(actor_critic.critic.parameters(), lr=vf_lr)
self.gamma, self.lamda, self.K_epochs = gamma, lamda, K_epochs
self.eps_clip, self.batch_size = eps_clip, batch_size
self.vf_coef, self.entropy_coef = vf_coef, entropy_coef
The actor and critic use separate Adam optimizers. Their learning rates are 3e-4 and 1e-3, respectively.
Each rollout records the log probability assigned by the policy that collected it. During an update, the current policy evaluates the same state and raw action:
action_logprobs = tanh_log_prob(batch_actions, dist)
ratios = torch.exp(action_logprobs - batch_logprobs)
The exponential converts the difference between two log probabilities into their probability ratio:
\[r_t(\theta) = \frac{\pi_\theta(a_t \mid s_t)}{\pi_{\theta_{\mathrm{old}}}(a_t \mid s_t)}\]A ratio of one means that both policies assign the same density to the sampled action. A larger ratio means that the current policy favors it more. The advantage determines the desired direction. A positive advantage encourages a larger ratio. A negative advantage encourages a smaller ratio.
PPO uses the clipped objective (Schulman et al.):
\[L^{CLIP}(\theta) = \mathbb{E}_t \left[ \min \left( r_t(\theta) \hat{A}_t, \text{clip}(r_t(\theta), 1-\varepsilon, 1+\varepsilon) \hat{A}_t \right) \right]\]The clipped term removes further objective improvement after the ratio passes the permitted interval in the favorable direction. It does not impose a strict bound on the full policy change. The following code implements this objective directly:
def compute_loss(self, batch_states, batch_actions, batch_logprobs,
batch_advantages, batch_returns):
action_means, action_stds, state_values = self.actor_critic(batch_states)
dist = torch.distributions.Normal(action_means, action_stds)
action_logprobs = tanh_log_prob(batch_actions, dist)
ratios = torch.exp(action_logprobs - batch_logprobs)
actor_loss = -torch.min(
ratios * batch_advantages,
torch.clamp(ratios, 1-self.eps_clip, 1+self.eps_clip) * batch_advantages
).mean()
critic_loss = F.mse_loss(state_values.squeeze(-1), batch_returns)
entropy = dist.entropy().sum(-1).mean()
return actor_loss + self.vf_coef * critic_loss - self.entropy_coef * entropy
The loss is the clipped objective, plus value MSE, minus an entropy bonus. One backward pass feeds both optimizers.
Consider the action sampled from states[t]. The environment returns rewards[t] and the next observation. Before optimization starts, the critic assigns state_values[t] to the stored state. The update must measure whether that action led to more reward than the critic expected.
The code makes this comparison for every transition, then works backward through the rollout:
def compute_advantages(self, rewards, state_values, is_terminals):
T, N = rewards.shape
advantages, gae = torch.zeros_like(rewards), torch.zeros(N, device=rewards.device)
state_values_pad = torch.cat([state_values, state_values[-1:]], dim=0)
for t in reversed(range(T)):
delta = rewards[t] + self.gamma * state_values_pad[t + 1] * (1 - is_terminals[t]) - state_values_pad[t]
gae = delta + self.gamma * self.lamda * (1 - is_terminals[t]) * gae
advantages[t] = gae
returns = advantages + state_values_pad[:-1]
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
return advantages.reshape(-1), returns.reshape(-1)
The first line inside the loop calculates the temporal-difference residual:
\[\delta_t = r_t + \gamma(1-d_t)V(s_{t+1}) - V(s_t)\]The residual is positive when the reward and next-state value exceed the previous estimate. It is negative when they fall short. Here, \(d_t\) is one when the episode ends, so no future value crosses the episode boundary.
One transition can look better because of rewards that arrive several steps later. The second line carries the next advantage backward, discounted by gamma * lamda. This is the recursive form of generalized advantage estimation (Schulman et al.):
The implementation uses gamma = 0.99 and lamda = 0.95. It first adds each unstandardized advantage to the old value estimate. This produces returns[t], the critic target. It then standardizes only the advantages used by the actor.
The final transition needs the critic value of the observation after the rollout. This code does not store that observation. state_values[-1:] therefore repeats the final stored value as the bootstrap value. A terminal mask removes it when the final transition ends an episode. For a nonterminal final transition, it is an approximation rather than \(V(s_T)\).
Training loop
The rollout stores the raw actions, before the tanh:
def __call__(self, state):
action_np, state_tensor, raw_action = self.actor_critic.act(
state, deterministic=False, return_internals=True
)
self.states.append(state_tensor)
self.actions.append(raw_action)
return action_np
All eight environments step together:
def rollout(env, policy, num_steps=None, num_episodes=None):
states, _ = env.reset()
traj_rewards, traj_dones = [], []
ep_returns, ep_rets, step_count = [], np.zeros(env.num_envs), 0
while True:
states, rewards, terminated, truncated, _ = env.step(policy(states))
traj_rewards.append(rewards)
traj_dones.append(np.logical_or(terminated, truncated))
ep_rets += rewards
step_count += env.num_envs
if np.any(traj_dones[-1]):
for idx in np.where(traj_dones[-1])[0]:
ep_returns.append(ep_rets[idx])
ep_rets[idx] = 0.0
if (num_steps and step_count >= num_steps) or
(num_episodes and len(ep_returns) >= num_episodes):
break
return traj_rewards, traj_dones, ep_returns
100k steps per epoch, then 20 update passes over them at batch size 5000:
def update(self, rewards, dones):
with torch.no_grad():
rewards = torch.as_tensor(np.stack(rewards), dtype=torch.float32).to(device)
is_terms = torch.as_tensor(np.stack(dones), dtype=torch.float32).to(device)
old_states, old_actions = torch.cat(self.states), torch.cat(self.actions)
action_means, action_stds, old_state_values = self.actor_critic(old_states)
old_logprobs = tanh_log_prob(old_actions,
torch.distributions.Normal(action_means, action_stds))
old_state_values = old_state_values.squeeze(-1).view(-1, rewards.size(1))
advantages, returns = self.compute_advantages(rewards, old_state_values, is_terms)
dataset = TensorDataset(old_states, old_actions, old_logprobs, advantages, returns)
for _ in range(self.K_epochs):
for batch in DataLoader(dataset, batch_size=self.batch_size, shuffle=True):
batch_states, batch_actions, batch_logprobs, batch_advantages, batch_returns = batch
self.pi_optimizer.zero_grad()
self.vf_optimizer.zero_grad()
loss = self.compute_loss(batch_states, batch_actions, batch_logprobs,
batch_advantages, batch_returns)
loss.backward()
torch.nn.utils.clip_grad_norm_(
list(self.actor_critic.actor.parameters()) + [self.actor_critic.log_std],
max_norm=0.5
)
torch.nn.utils.clip_grad_norm_(self.actor_critic.critic.parameters(), max_norm=0.5)
self.pi_optimizer.step()
self.vf_optimizer.step()
self.states, self.actions = [], []
Gradients clipped at 0.5.
I evaluate every 10 epochs:
def evaluate_policy(actor_critic, n=16, render=False, num_episodes=None):
env = make_env(1 if render else n, render)
def policy(s): return actor_critic.act(s, deterministic=True)
if render and num_episodes:
_, _, ep_rets = rollout(env, policy, num_episodes=num_episodes)
else:
_, _, ep_rets = rollout(env, policy, num_steps=max_timesteps * (1 if render else n))
env.close()
return float(np.mean(ep_rets)) if ep_rets else 0.0
It stops when the moving average hits 250, which takes about 100 epochs.
One small step for the optimizer, one giant leap for the GPU bill.