Score squashed gSDE actions on the pre-tanh sample

[shaneraphel]

Description

With use_sde=True and squash_output=True, StateDependentNoiseDistribution.log_prob recovered the Gaussian sample by inverting tanh of the stored action. In float32, tanh of a large sample rounds to exactly 1, and the clipped inverse (atanh(1 - eps)) is a different point. The PPO surrogate then differentiates that other point. At a unit-variance mean of 12 the reconstructed latent is about 8.32 instead of 12.78, and the mean gradient changes sign. The initial probability ratio stays 1 either way, because each route uses its own old likelihood.

sample and mode now keep the pre-tanh sample when squashing. The rollout buffer stores it, and PPO and A2C pass it back into evaluate_actions. Likelihoods that only have the squashed action still invert it, as before.

Motivation and Context

Fixes the likelihood path reported in #2285. The unsaturated control is unchanged: tests/test_sde.py passes. A new test checks that a saturated float32 action inverts to a different point, and that log_prob given the drawn sample matches x - mean for a unit-variance Normal, while inverting the stored action gives the opposite gradient sign.

• I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)

Checklist

• I've read the CONTRIBUTION guide (required)
• I have updated the changelog accordingly (docs/misc/changelog.md) (required).
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.
• I have opened an associated PR on the SB3-Contrib repository (if necessary)
• I have opened an associated PR on the RL-Zoo3 repository (if necessary)
• I have reformatted the code using make format (required)
• I have checked the codestyle using make check-codestyle and make lint (required)
• I have ensured make pytest and make type both pass. (required)
• I have checked that the documentation builds using make doc (required)

tests/test_gsde_tanh_latent.py and tests/test_sde.py: 15 passed, 2 skipped. I did not run make format, make pytest, or make doc. ruff check on the edited files reported no errors. ruff format --check also wants to reflow lines in policies.py and ppo.py that this change does not touch, so I left those as they are.

[Ibrahimkhan4real]

Hi @shaneraphel, thanks for taking this on! I originally reported #2285.

I gave this branch a try (ea03e76) and found a weird bug where the stored pre-tanh sample doesn’t look right. Right now it’s read from policy.action_dist when the step gets added to the buffer, but callback.on_step() runs between forward() and rollout_buffer.add(). If a callback calls predict(), for example EvalCallback, it reuses the same distribution object and ends up overwriting the sample.

With a callback that calls predict() every step, I saw 64/64 stored samples differ from the ones actually used for acting. Without that callback, it was 0/64. So those steps later get replayed with a sample that came from a different observation. Easiest fix is probably just caching the sample inside forward() directly instead of trying to read it back from the shared dist object later.

import numpy as np, torch as th
from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import BaseCallback

class PredictInStep(BaseCallback):
    def __init__(self, other_obs):
        super().__init__(); self.other_obs = other_obs
    def _on_step(self):
        self.model.predict(self.other_obs, deterministic=True)
        return True

def run(with_predict):
    model = PPO("MlpPolicy", "Pendulum-v1", use_sde=True, n_steps=64, batch_size=64, n_epochs=1,
                policy_kwargs=dict(squash_output=True), seed=0, device="cpu")
    fwd, stored = [], []
    orig_forward = model.policy.forward
    def forward(obs, deterministic=False):
        out = orig_forward(obs, deterministic)
        fwd.append(model.policy.action_dist.gaussian_actions.detach().clone())
        return out
    model.policy.forward = forward
    orig_add = model.rollout_buffer.add
    def add(*a, gaussian_actions=None, **k):
        stored.append(gaussian_actions.detach().clone())
        return orig_add(*a, gaussian_actions=gaussian_actions, **k)
    model.rollout_buffer.add = add
    other = np.array([1.0, 0.0, 0.0], dtype=np.float32)
    model.learn(64, callback=PredictInStep(other) if with_predict else None)
    mism = sum(not th.allclose(f, s) for f, s in zip(fwd, stored))
    print(f"predict-in-callback={with_predict}: {mism}/{len(stored)} stored samples differ")

run(False); run(True)

[Ibrahimkhan4real]

I’d be happy to continue from your branch, as in fix this and add the missing tests. Would that be OK with you?