Fix PPO importance-sampling ratio bias with squashed-Gaussian policies

[jmtoepperwien]

Fix PPO importance-sampling ratio bias with squashed-Gaussian policies

Problem

When PPO uses a squashed-Gaussian policy (4-tuple output: action, z, mean, log_std), the importance-sampling ratio becomes biased:

• During rollout, old_log_probs stored in the buffer are computed without the tanh change-of-variables correction
• During the PPO update, new_log_probs include the correction
• The correction does not cancel in the ratio:

  ratio = exp(log_π_new(a) - log_π_old(a))
        = exp([log N(z|new) - log_correction] - log N(z|old))
        = [N(z|new)/N(z|old)] × 1/(1 - tanh(z)²)
  

• The spurious 1/(1 - tanh(z)²) factor explodes as actions approach ±1, causing gradient instability

Solution

Two commits:

1. Always apply tanh correction at rollout time (97e8391)

   • Remove the sac= gate from sample_nondeterministic_logprobs
   • Apply the correction whenever the model outputs a 4-tuple (which is the canonical signal that squashing is active)
   • This ensures old_log_probs stored in the buffer are already corrected, so the correction cancels correctly in the PPO update

2. Check tanh_squash attribute explicitly (5e59c37)

   • Add tanh_squash=True to SACModel as a class attribute (matching PPOModel)
   • Pass tanh_squash=getattr(self.model, "tanh_squash", False) to sample_nondeterministic_logprobs at each call site
   • Makes the intent explicit and allows future models to opt in/out cleanly