[Bug]: n_steps > 1 is silently ignored when replay_buffer_class is passed explicitly

[Koustav-github]

🐛 Bug

When replay_buffer_class is passed explicitly to an off-policy algorithm, n_steps > 1
is silently ignored and the algorithm trains with 1-step returns.

In OffPolicyAlgorithm._setup_model(), both the n-step buffer selection and its validation
live inside the if self.replay_buffer_class is None: branch. When the user supplies the
class, neither runs: self.n_steps keeps its value, but the buffer is not an
NStepReplayBuffer, so sample() returns discounts=None and train() falls back to
discounts = self.gamma. No exception and no warning are raised.

To Reproduce

1. n_steps is accepted but never applied when the buffer class is given explicitly

from stable_baselines3 import SAC
from stable_baselines3.common.buffers import ReplayBuffer

model = SAC("MlpPolicy", "Pendulum-v1", n_steps=3, buffer_size=200,
            replay_buffer_class=ReplayBuffer)
print(model.n_steps, type(model.replay_buffer).__name__)

3 ReplayBuffer

n_steps=3 is stored on the model, but the buffer built cannot compute n-step returns.
sample() therefore returns discounts=None, and train() falls back to
discounts = self.gamma, i.e. 1-step returns. Nothing is raised or warned.

2. The same configuration is rejected or accepted depending only on whether the class is named

from stable_baselines3 import SAC
from stable_baselines3.common.envs import BitFlippingEnv
from stable_baselines3.her import HerReplayBuffer

env_kwargs = dict(n_bits=4, continuous=True)

try:
    SAC("MultiInputPolicy", BitFlippingEnv(**env_kwargs), n_steps=3, buffer_size=200)
except AssertionError as e:
    print("no explicit class ->", type(e).__name__ + ":", e)

m = SAC("MultiInputPolicy", BitFlippingEnv(**env_kwargs), n_steps=3, buffer_size=200,
        replay_buffer_class=HerReplayBuffer, learning_starts=10)
print("explicit HerReplayBuffer -> no error;", type(m.replay_buffer).__name__, "n_steps =", m.n_steps)

no explicit class -> AssertionError: N-step returns are not supported for Dict observation spaces yet.
explicit HerReplayBuffer -> no error; HerReplayBuffer n_steps = 3

Relevant log output / Error message

No error and no warning are raised in the silent cases -- that is the bug.

The only exception produced anywhere in the repro is:

    AssertionError: N-step returns are not supported for Dict observation spaces yet.

It comes from the case that already behaves CORRECTLY: a Dict observation space with
n_steps=3 and no explicit replay_buffer_class. It is shown to highlight the inconsistency,
not because it is the problem being reported.

System Info

No response

Checklist

• My issue does not relate to a custom gym environment. (Use the custom gym env template instead)
• I have checked that there is no similar issue in the repo
• I have read the documentation
• I have provided a minimal and working example to reproduce the bug
• I've used the markdown code blocks for both code and stack traces.

[Koustav-github]

I'd be happy to work on this if you agree with the direction — could you assign it to me?

Proposed fix: move the n-step check out of the if self.replay_buffer_class is None: branch
in _setup_model(), so it applies to the resolved buffer class regardless of where it came from:

• n_steps > 1 + Dict observation space → raise (currently only raised on the default path)
• n_steps > 1 + a class that is not an NStepReplayBuffer subclass → raise, naming the class
• n_steps > 1 + an NStepReplayBuffer subclass → forward n_steps/gamma, as the default path already does

The last case uses issubclass rather than an identity check on purpose, so a user subclass of
NStepReplayBuffer keeps working — today such a subclass silently receives the buffer's own
defaults (n_steps=3, gamma=0.99) instead of the values passed to the algorithm.

One open question I'd like your call on: the existing Dict-obs check is an assert. I'd lean
towards ValueError so both rejections share a type and survive python -O, but that changes
the exception type on an error path — happy to keep assert if you prefer the smaller diff.

This has been the behaviour since 2.7.0, when n_steps was introduced; nothing about the
interaction with replay_buffer_class has changed since.

[Koustav-github]

Draft PR opened at #2275 with tests and a changelog entry — no hurry at all, just flagging it. Happy to adjust the assert vs ValueError choice (and anything else) whenever you're back.

[araffin]

Hello,

n_steps parameter is only meant to be used with the default replay buffer settings.
If you pass replay_buffer_class=MyCustomBuffer, you must pass replay_buffer_kwargs={"n_steps": 3} and then it should work.

Instead of changing the way n_steps behave, I would prefer to emit a warning when n_steps > 1 and a custom class is passed (replay_buffer_class not None).

[Koustav-github]

Verified — replay_buffer_class=NStepReplayBuffer with replay_buffer_kwargs={"n_steps": 3}
does compute n-step returns (discounts is not None), so the supported path works as you
describe. Dropping the behaviour change; I'll rework the PR as documentation plus the warning
you suggested.

Noting the documentation label — the current docstring is:

:param n_steps: When n_step > 1, uses n-step return (with the NStepReplayBuffer) when updating the Q-value network.

It names NStepReplayBuffer, but doesn't say that n_steps is ignored when replay_buffer_class
is passed, or that replay_buffer_kwargs is the way to configure a custom buffer. I'd propose
expanding it across off_policy_algorithm.py, dqn.py, sac.py, td3.py and ddpg.py to make
both explicit. Happy to adjust the wording or the set of files.

One question that affects both the docs and the warning text: gamma isn't forwarded either. With

SAC("MlpPolicy", "Pendulum-v1", gamma=0.95,
    replay_buffer_class=NStepReplayBuffer, replay_buffer_kwargs={"n_steps": 3})

the buffer keeps its own default gamma=0.99 while the algorithm uses 0.95, so n-step returns
get discounted differently from the critic target. Should the docs and warning tell users to pass
both n_steps and gamma in replay_buffer_kwargs?

[araffin]

Dropping the behaviour change; I'll rework the PR as documentation plus the warning
you suggested.

Please don't use a LLM when replying to issues.

Should the docs and warning tell users to pass
both n_steps and gamma in replay_buffer_kwargs?

probably, but I would only trigger it when n_steps > 1

[Koustav-github]

probably, but I would only trigger it when n_steps > 1

Alright, on it.

[Koustav-github]

hey @araffin,

kindly check the commits now, warning is now triggered when n_steps > 1 and a custom class is passed, instead of changing the way n_steps behave.

[Koustav-github]

hey @araffin,

I have addressed the issues and resolved it the one that u asked for in the last conversation on the pr branch. #2275

[Koustav-github]

hey @araffin,

i saw the "llm generated" tag on pr #2275 given by u. So, its better to say upfront, yes, I have used llm to root-cause the issue and draft them, but regarding "tell your human", I was always here and I do understand that

model = SAC(..., replay_buffer_class=HerReplayBuffer, n_steps=4)

this trained in 1 step (silently) because n_steps/gamma is only forwarded in the auto-selection path and nothing warns.
and Thanks for approving the changes. :)