[Feature Request] Make PPOTrainer GAE recomputation frequency explicit and configurable

[JagerHoHo]

Motivation

PPOTrainer / OnPolicyTrainer currently appears to recompute GAE on every optimization epoch when add_gae=True.

The relevant call sequence is:

1. OnPolicyTrainer creates or accepts a GAE module.

2. It registers the module through:
   

   rl/torchrl/trainers/algorithms/on_policy.py

   Line 210 in 2a77e48

   ValueEstimatorHook(gae).register(self)

3. ValueEstimatorHook registers itself at the "pre_epoch" hook stage.
   

   rl/torchrl/trainers/trainers.py

   Line 2911 in 2a77e48

   trainer.register_op("pre_epoch", self)

4. Trainer.optim_steps() invokes "pre_epoch" hooks inside the optimization-epoch loop:
   

   rl/torchrl/trainers/trainers.py

   Lines 1607 to 1616 in 2a77e48

   	def optim_steps(self, batch: TensorDictBase) -> None:
   	average_losses = None
   	self._pre_optim_hook()
   	optim_steps_per_batch = self.optim_steps_per_batch
   	j = -1
   	for _ in range(self.num_epochs):
   	# LOGGING POINT 3: Pre-epoch logging (e.g., epoch-specific metrics)
   	self._pre_epoch_log_hook(batch)

Therefore, when num_epochs=N, GAE is computed N times for the same collected rollout. Because the critic is updated between epochs, subsequent calls use the updated critic and overwrite both advantage and value_target. When average_gae=True, advantage normalization is also repeated each epoch.

This behavior is surprising when migrating a conventional manual PPO loop to PPOTrainer. Many PPO implementations follow:

collect rollout → compute GAE once → optimize the fixed advantages and value targets for N epochs

The current trainer instead follows:

collect rollout
→ compute GAE
→ optimize epoch 1
→ recompute GAE with the updated critic
→ optimize epoch 2
→ ...

Per-epoch advantage recomputation is a valid PPO variant and has been studied in What Matters in On-Policy Reinforcement Learning?.
It has also been discussed in Stable-Baselines3 issue #445.

The problem is not necessarily that the current behavior is incorrect. It is a meaningful algorithmic choice that is currently selected implicitly by the hook placement and is not obvious from the add_gae argument.

Because GAE is recomputed using the updated critic before every epoch, the actor is effectively chasing a moving target. The advantages, value targets, and normalized advantage distribution can change between successive passes over the same rollout. As a result, the actor does not optimize one fixed empirical surrogate throughout the PPO update. This may be an intentional algorithmic variant, but it should be explicit and configurable because migrating to PPOTrainer can otherwise silently change the training semantics.

This is especially relevant for custom PPO-compatible objectives in which the advantage determines an explicit actor target rather than serving only as a linear policy-gradient coefficient.

Solution

Add an explicit option controlling when GAE is computed. For example:

PPOTrainer(
    ...,
    advantage_recomputation="per_epoch",
)

Supported values could be:

"per_batch"
    Compute GAE once when a rollout batch is received. Keep the resulting
    advantages and value targets fixed across all optimization epochs.

"per_epoch"
    Recompute GAE with the latest critic before every optimization epoch.
    This preserves the current behavior.

Alternatives

A simpler Boolean option could be introduced:

recompute_advantages_each_epoch: bool = True

Checklist

• I have checked that there is no similar issue in the repo (required)

[vmoens]

I think this conflates three distinct ways of using GAE:

1. Compute it once per collected rollout, before the optimization epochs. Advantages and value targets remain fixed.
2. [preferred] Compute it once at the beginning of every optimization epoch. Advantages are refreshed using the updated critic. This is the current trainer behavior.
3. Compute it on every sampled minibatch inside the loss.

The first two are training-loop decisions and belong outside the loss, as demonstrated by the TorchRL examples and implementations.

The third is only valid when sampled minibatches preserve the temporal ordering and trajectory context required by GAE, including correct episode boundaries and bootstrapping. Ordinary flattened and shuffled PPO minibatches do not satisfy those requirements. Supporting this requires a specialized replay-buffer sampler and compatible settings, and should remain a last-resort usage rather than a generally advertised trainer mode.

Consequently, I don’t think an advantage_recomputation option with values such as "per_batch" and "per_epoch" is the right abstraction. In particular, "per_batch" is ambiguous between a collected rollout and an optimization minibatch, although those have very different correctness requirements.

We can improve the documentation to make the existing semantics explicit:

• add_gae=True computes GAE at the pre_epoch stage.
• Users wanting fixed advantages should set add_gae=False and compute GAE once on the collected rollout before optimization.

Both supported approaches remain explicit in the training loop without adding another trainer configuration that obscures where and under which sampling assumptions GAE is valid.

[JagerHoHo]

I see. Thank you for the response.