This one is not a bug report but rather, an alternative suggestion for your consideration. learning_agents/ppo.py trains recurrent policies on fixed-length windows (PolicyWindowSize) taken from the flattened experience buffer. Windows start at every step, so many of them cross from one episode into the next. Today the trainer handles this with a loss mask: every slot after the first done in a window is dropped from the policy loss, so the previous episode’s memory never trains the next episode’s steps.
That works, but it leaves the loss with uneven sample counts. A step k steps into its episode is kept only in windows that start inside the same episode, so it lands in the loss min(k+1, W) times. With W = 16, the first step of every episode counts once, while mid-episode steps count 16times. The opening steps of every episode are therefore under-weighted by up to [1,2,...16] times. Those are the spawn and reset states the agent sees every episode. Nothing in the logged metrics shows this.
There’s also a smaller effect: the masked terms are still averaged over all B * W slots, so the logged loss/policy_* values scale with episode length.
Suggestion: reset and continue
Instead of masking the slots after a boundary, you could reset the propagated memory to the buffer’s recorded memory state at the boundary and keep training. For the first step of an episode, that’s the episode’s initial memory, which is what the agent actually had at that point in play. Every slot becomes a real training step:
- Steps get equal weight, except at the two ends of the rollout.
- The
B * Waverage counts only real slots, so the logged losses no longer depend on episode length. - Backpropagation through time still stops at the boundary, because the reset value comes from the buffer not the computation graph.
Here is a diagram I put together to try and help explain it.
Top: every window in a small buffer (window 4). Hatched slots are masked out, and R marks a memory reset. With the mask, B0 and C0 count once. With reset, they count 4 times. Bottom: a whole rollout with window 16. The mask ramps from 1 to 16 again at the start of every episode. Reset only ramps at the two ends of the rollout, where both versions run out of windows.
The changes required in ppo.py
The change to enable this is fairly small. In window construction add one reset flag per slot which is true when the previous slot ended an episode. These lines replace the mask construction.
done_flat = (terminated | truncated)
window_resets = np.zeros_like(window_indices, dtype=bool)
window_resets[:, 1:] = done_flat[window_indices[:, :-1]]
In train_policy, the flags are passed in place of the mask. The memory is reset where a flag is set, and mask[:,i] * is removed from the three loss terms:
policy_mem = mem[:,0]
for i in range(0, policy_window):
if i > 0:
policy_mem = torch.where(reset[:,i].unsqueeze(-1), mem[:,i], policy_mem)
act_enc_policy_mem = self.policy(torch.cat([self.encoder(obs[:,i]), policy_mem], dim=-1))
act_enc, policy_mem = act_enc_policy_mem[:,:self.act_enc_num], act_enc_policy_mem[:,self.act_enc_num:]
act_dist = self.decoder(act_enc)
ratio = schema_log_prob(self.action_schema, act_dist, act[:,i], mod[:,i]) - logp[:,i]
surr = torch.min(ratio, ratio.clamp(self.eps_clip_log_min, self.eps_clip_log_max)) # Note: this is the old surr without the neg adv fix
loss_surr += self.action_surr_weight * (adv[:,i] * -(surr / self.act_num)).mean() / policy_window
loss_ent += self.action_ent_weight * (-(schema_entropy(self.action_schema, act_dist, mod[:,i]) / self.act_num)).mean() / policy_window
loss_reg += self.action_reg_weight * (schema_regularization(self.action_schema, act_dist) / self.act_num).mean() / policy_window
A few miscellaneous points
- This relies on the last stored step of every episode carrying a
terminatedortruncatedflag. The buffer already guarantees that, and the current mask relies on it too. - Episode ends are not under-weighted by the mask, because it keeps the slots up to and including a
done. Only the start of the next episode loses weight. - The ramps at the two ends of the rollout remain with either version. Removing them would need something extra, such as per-step importance weights or windows anchored at episode starts.
- Whether equal weighting is the right objective is open to debate. A step seen with only a few steps of context can be argued to be a different prediction problem, not an under-sampled one. I mention it just to raise awareness and show an option that I chose to take. You may have users who have finetuned their hyper-parameters with the current behavior and this may affect those selections.
- There is an alternative formulation where you count how many slots each samples occurs in and then use that as an importance weight. In that version the loss for each slot would be divided by the total windows it occurs in.