Code reference

Auto-generated from source docstrings. For the conceptual contract see Core concepts — the environment contract, the backups, and the MAP naming law.

The module layout follows MAP: buffers carry the Mode, the *_1p / *_2p modules carry the Players, and the *_base modules hold what both player counts of an Algorithm share.

Backups — the value operators

safety_sb3.backups

The Bellman backups — defined ONCE, used by every learner.

This library solves two different safety problems, and they take different value operators. Before this module existed the operators were re-implemented at four call sites (the numpy rollout buffer, the torch rollout buffer, and two SAC train() bodies); they silently diverged, and the PPO family spent the whole 2026 campaign optimizing a fixed point that was neither problem's value. Every backup now routes through here so that cannot recur.

A learner is then written without reference to a particular operator — it takes a mode and calls :func:target. Ordinary sum-of-rewards RL is one of the modes (:data:CUMULATIVE), so it is a special case of the same code, not a separate library.

Convention (throughout the library)::

g(s) >= 0   <=>  s is OUTSIDE the failure set   ("safe")
l(s) >= 0   <=>  s is INSIDE the target set     ("reached")
V is MAXIMIZED;  V(s) >= 0  <=>  s is in the solution set

The operators

Avoid (stay safe forever) — Fisac et al. 2019; ISAACS eq. 6/7::

V(s) = (1 - gamma) * g(s)  +  gamma * min( g(s), V(s') )

Reach-avoid (reach the target while staying safe throughout) — Hsu et al. RSS'21 eq. 15; Gameplay Filters eq. 6a::

V(s) = (1 - gamma) * min(l(s), g(s))
     + gamma * min( g(s), max( l(s), V(s') ) )

Cumulative (the STANDARD RL operator — maximize the discounted sum of rewards; Bellman 1957, and every textbook since)::

V(s) = r(s)  +  gamma * V(s')

This one is not a safety operator and does not follow the sign convention above: its first argument is a reward, not a margin, V is an expected discounted return with no zero-level-set meaning, and there is no g, no l and no failure set. It is here because the learners in this library are written against :func:target rather than against a fixed backup, so plain sum-of-rewards RL costs one branch and comes out as the degenerate member of the family. Use it to run a nominal (reward-maximizing) baseline through exactly the same actor/critic/entropy code as the safety learners — the standard control for "is the safety operator doing the work, or is it just SAC?". Do not read a cumulative V as a safety certificate; V >= 0 means nothing here.

The two safety operators share the shape (1 - gamma) * anchor + gamma * backup. The anchor is the "episode terminates now" payoff — 1 - gamma is the per-step termination probability of the discounted formulation, so the anchor is what the trajectory scores if it stops here:

• avoid: stopping now is a win iff you are safe -> g
• reach-avoid: iff you are in the target AND safe -> min(l, g)

The reach-avoid anchor is the same expression as the finite-horizon terminal condition V_H = min(l, g) (Gameplay Filters eq. 5b); that identity is the structural tell that it is not a stylistic choice.

Why the anchors are not interchangeable

Anchoring reach-avoid on g makes "stay safe forever, never reach" a fixed point at V = g > 0 — a win — when the reach-avoid value of such a trajectory is max_t min(l_t, min_{s<=t} g_s) = max_t l_t < 0, a loss. The resulting fixed point is neither problem's value; RSS'21's under-approximation theorem (RA_gamma nested inside RA) stops applying, so the critic can wrongly certify reachability and is unsound to filter with. RSS'21 says of the g-anchored form (its eq. 13) that it approximates "safety or liveness problems, but not both".

Note that ISAACS anchors on g and is correct to: it is a pure avoid game with no target set and no l anywhere in the paper. Gameplay Filters extends ISAACS to reach-avoid and changes the anchor when it does. The mixture — a g anchor with a max(l, V') recursion — appears in none of the papers.

Avoid is not expressible as a reach-avoid instance

Tempting, and wrong: to make the reach-avoid operator compute the avoid value you would need the anchor to reduce (min(l, g) = g, i.e. l >= g) and the recursion to reduce (max(l, V') = V', i.e. l <= V'). Since the avoid recursion caps V' <= g, that needs l >= g >= V' >= l — satisfiable only in the degenerate case. A large negative l buys the recursion and destroys the anchor (V == l everywhere, empty safe set); a large positive l buys the anchor and destroys the recursion (V == g everywhere, no lookahead — coming failures never propagate). Use the avoid operator for avoid problems.

All functions here are elementwise and shape-agnostic, and accept either numpy arrays or torch tensors (not_done may be float 0/1 or bool-as-float).

avoid_target

avoid_target(g, v_next, not_done, gamma)

Avoid (safety) target: (1-g)*g + g*min(g, V'); terminal -> g.

Parameters:

Name	Type	Description	Default
g	Array	safety margin g(s), >= 0 iff safe.	required
v_next	Array	bootstrap value V(s').	required
not_done	Array	1.0 on non-terminal steps, 0.0 on terminal steps.	required
gamma	float	discount.	required

Source code in safety_sb3/backups.py

def avoid_target(g: Array, v_next: Array, not_done: Array,
                 gamma: float) -> Array:
  """Avoid (safety) target: ``(1-g)*g + g*min(g, V')``; terminal -> ``g``.

  :param g: safety margin ``g(s)``, ``>= 0`` iff safe.
  :param v_next: bootstrap value ``V(s')``.
  :param not_done: 1.0 on non-terminal steps, 0.0 on terminal steps.
  :param gamma: discount.
  """
  xp = _xp(g)
  v_to_go = xp.minimum(g, v_next)
  return (not_done * ((1.0 - gamma) * g + gamma * v_to_go)
          + (1.0 - not_done) * g)

reach_avoid_target

reach_avoid_target(
    g, l, v_next, not_done, gamma, terminal_type="all"
)

Reach-avoid target: (1-g)*min(l,g) + g*min(g, max(l, V')).

Parameters:

Name	Type	Description	Default
g	Array	safety margin g(s), >= 0 iff safe.	required
l	Array	target margin l(s), >= 0 iff in the target set.	required
v_next	Array	bootstrap value V(s').	required
not_done	Array	1.0 on non-terminal steps, 0.0 on terminal steps.	required
gamma	float	discount.	required
terminal_type	str	target at terminal steps. "all" (default) -> min(l, g), the finite-horizon terminal condition (Gameplay Filters eq. 5b); "g" -> g alone. Both are offered by the reference (safe_adaptation_dev/utils/train.py). With timeout bootstrapping disabled, timeouts arrive here as terminal steps and "all" is correct for them: a timeout IS the horizon cutoff of eq. 5b.	'all'

Source code in safety_sb3/backups.py

def reach_avoid_target(g: Array, l: Array, v_next: Array, not_done: Array,
                       gamma: float, terminal_type: str = "all") -> Array:
  """Reach-avoid target: ``(1-g)*min(l,g) + g*min(g, max(l, V'))``.

  :param g: safety margin ``g(s)``, ``>= 0`` iff safe.
  :param l: target margin ``l(s)``, ``>= 0`` iff in the target set.
  :param v_next: bootstrap value ``V(s')``.
  :param not_done: 1.0 on non-terminal steps, 0.0 on terminal steps.
  :param gamma: discount.
  :param terminal_type: target at terminal steps. ``"all"`` (default) ->
      ``min(l, g)``, the finite-horizon terminal condition (Gameplay Filters
      eq. 5b); ``"g"`` -> ``g`` alone. Both are offered by the reference
      (``safe_adaptation_dev/utils/train.py``). With timeout bootstrapping
      disabled, timeouts arrive here as terminal steps and ``"all"`` is
      correct for them: a timeout IS the horizon cutoff of eq. 5b.
  """
  check_terminal_type(terminal_type)
  xp = _xp(g)
  anchor = xp.minimum(l, g)
  v_to_go = xp.minimum(g, xp.maximum(l, v_next))
  terminal = g if terminal_type == "g" else anchor
  return (not_done * ((1.0 - gamma) * anchor + gamma * v_to_go)
          + (1.0 - not_done) * terminal)

cumulative_target

cumulative_target(reward, v_next, not_done, gamma)

Standard RL target: r + gamma * V'; terminal -> r.

The ordinary discounted-cumulative-reward backup, so that reward-maximizing RL is a mode of this library rather than a separate implementation. Unlike the two safety operators this one carries no margin semantics: the first argument is a reward, and the resulting V is an expected return whose sign means nothing.

Parameters:

Name	Type	Description	Default
reward	Array	immediate reward r(s, a) (NOT a margin).	required
v_next	Array	bootstrap value V(s').	required
not_done	Array	1.0 on non-terminal steps, 0.0 on terminal steps.	required
gamma	float	discount.	required

Source code in safety_sb3/backups.py

def cumulative_target(reward: Array, v_next: Array, not_done: Array,
                      gamma: float) -> Array:
  """Standard RL target: ``r + gamma * V'``; terminal -> ``r``.

  The ordinary discounted-cumulative-reward backup, so that reward-maximizing
  RL is a mode of this library rather than a separate implementation. Unlike
  the two safety operators this one carries no margin semantics: the first
  argument is a **reward**, and the resulting ``V`` is an expected return whose
  sign means nothing.

  :param reward: immediate reward ``r(s, a)`` (NOT a margin).
  :param v_next: bootstrap value ``V(s')``.
  :param not_done: 1.0 on non-terminal steps, 0.0 on terminal steps.
  :param gamma: discount.
  """
  return reward + gamma * not_done * v_next

target

target(
    mode,
    g,
    v_next,
    not_done,
    gamma,
    l=None,
    terminal_type="all",
)

Dispatch to the operator for mode. See the module docstring.

Parameters:

Name	Type	Description	Default
mode	str	backups.AVOID, backups.REACH_AVOID or backups.CUMULATIVE.	required
g	Array	the safety margin g(s) for the safety modes; for CUMULATIVE this slot carries the reward instead.	required
l	Optional[Array]	required for REACH_AVOID; ignored for the other modes.	None

Source code in safety_sb3/backups.py

def target(mode: str, g: Array, v_next: Array, not_done: Array, gamma: float,
           l: Optional[Array] = None, terminal_type: str = "all") -> Array:
  """Dispatch to the operator for ``mode``. See the module docstring.

  :param mode: ``backups.AVOID``, ``backups.REACH_AVOID`` or
      ``backups.CUMULATIVE``.
  :param g: the safety margin ``g(s)`` for the safety modes; for
      ``CUMULATIVE`` this slot carries the **reward** instead.
  :param l: required for ``REACH_AVOID``; ignored for the other modes.
  """
  check_mode(mode)
  if mode == AVOID:
    return avoid_target(g, v_next, not_done, gamma)
  if mode == CUMULATIVE:
    return cumulative_target(g, v_next, not_done, gamma)
  if l is None:
    raise ValueError("mode='reach-avoid' requires the target margin l")
  return reach_avoid_target(g, l, v_next, not_done, gamma, terminal_type)

On-policy learners (PPO family)

safety_sb3.ppo_base

PPO, written without a fixed Bellman backup — the shared half of the family.

:class:AbstractPPO is everything the PPO-family learners share regardless of how many players are in the game: the env guards, the rsl_rl-parity training recipe, the KL-adaptive learning rate, gamma annealing, and the choice of rollout buffer from _MODE. It deliberately owns no rollout loop — that is the axis its subclasses split on:

:class:`~safety_sb3.ppo_1p.AbstractPPO1P` — one actor, one buffer.
:class:`~safety_sb3.ppo_2p.AbstractPPO2P` — two actors, two value nets,
    two buffers, and a phase machine.

In the PPO family the backup is carried by the rollout buffer, not by train(): compute_returns_and_advantage substitutes the mode's operator for the one-step TD target inside GAE. So a learner selects its problem simply by declaring _MODE, and :func:safety_sb3.buffers_rollout.rollout_buffer_classes hands it the matching buffers.

The rsl_rl-parity pieces (validated: they are what lets this learn hard robot tasks where stock-SB3-PPO-on-g stalls):

• No timeout bootstrapping (bootstrap_on_timeout=False, the default). Stock PPO adds gamma * V(terminal) to the reward at truncations (on_policy_algorithm.collect_rollouts). In the two safety modes the reward IS the physical safety margin g(s) — an absolute quantity — so adding a value estimate destroys its semantics and silently corrupts the backup the moment episodes start reaching the time limit. rsl_rl's SafetyPPO overrides process_env_step for exactly this reason; we gate it in the collect loops.

• KL-adaptive learning rate (adaptive_lr=True with desired_kl). rsl_rl raises the LR when the update KL sits well below target and lowers it when above. SB3's target_kl only early-stops; it never adjusts the LR, so in the low-KL / weak-gradient regime typical of the sparse safety signal the policy barely moves. Enabling this recovers rsl_rl's behavior.

Observation normalization is orthogonal and handled the SB3 way — wrap the env in VecNormalize — which mirrors rsl_rl's built-in running obs normalizer.

All knobs are constructor arguments, so downstream projects configure the algorithm without subclassing.

AbstractPPO

Bases: GammaAnnealMixin, PPO

PPO whose Bellman backup is chosen by _MODE (see the module docstring).

Parameters:

Name	Type	Description	Default
bootstrap_on_timeout	bool	if False (default) skip PPO's timeout value bootstrapping — correct whenever the reward is a margin g(s).	False
normalize_obs	bool	wrap the env in VecNormalize(norm_obs=True, norm_reward=False) — obs normalization is needed to match rsl_rl on hard robot tasks; reward normalization is refused (it corrupts g).	False
adaptive_lr	bool	enable rsl_rl-style KL-adaptive learning rate.	False
desired_kl	float | None	target KL for the adaptive LR controller.	0.01
lr_bounds	tuple[float, float]	(min, max) bounds for the adaptive LR.	(1e-05, 0.01)
adaptive_lr_factor	float	multiplicative step for the adaptive LR.	1.5
gamma_anneal		discount-factor annealing (ON by default). True anneals gamma 0.99 -> 0.9999 over the first 50% of training then holds (the reach-avoid boundary only sharpens as gamma -> 1; see gamma_anneal.py). False keeps gamma constant; a callable frac -> gamma supplies a custom schedule. GPU-resident path: pass a :class:~safety_sb3.tensor_env.TensorVecEnv and everything (rollout, buffer, backup, minibatching) stays on device — no numpy bounce. Detected automatically; env metrics() (curriculum levels etc.) are forwarded to the logger every rollout.	True

Source code in safety_sb3/ppo_base.py

class AbstractPPO(GammaAnnealMixin, PPO):
    """PPO whose Bellman backup is chosen by ``_MODE`` (see the module docstring).

    :param bootstrap_on_timeout: if False (default) skip PPO's timeout value
        bootstrapping — correct whenever the reward is a margin ``g(s)``.
    :param normalize_obs: wrap the env in VecNormalize(norm_obs=True,
        norm_reward=False) — obs normalization is needed to match rsl_rl on hard
        robot tasks; reward normalization is refused (it corrupts ``g``).
    :param adaptive_lr: enable rsl_rl-style KL-adaptive learning rate.
    :param desired_kl: target KL for the adaptive LR controller.
    :param lr_bounds: (min, max) bounds for the adaptive LR.
    :param adaptive_lr_factor: multiplicative step for the adaptive LR.
    :param gamma_anneal: discount-factor annealing (ON by default). ``True``
        anneals gamma 0.99 -> 0.9999 over the first 50% of training then holds
        (the reach-avoid boundary only sharpens as gamma -> 1; see
        ``gamma_anneal.py``). ``False`` keeps gamma constant; a callable
        ``frac -> gamma`` supplies a custom schedule.

    GPU-resident path: pass a :class:`~safety_sb3.tensor_env.TensorVecEnv` and
    everything (rollout, buffer, backup, minibatching) stays on device — no numpy
    bounce. Detected automatically; env ``metrics()`` (curriculum levels etc.)
    are forwarded to the logger every rollout.
    """

    #: the Bellman operator this learner converges to; concretes set it
    _MODE = backups.AVOID

    def __init__(
        self,
        *args,
        learning_rate=3e-4,
        rollout_buffer_class=None,
        rollout_buffer_kwargs=None,
        bootstrap_on_timeout: bool = False,
        normalize_obs: bool = False,
        adaptive_lr: bool = False,
        desired_kl: float | None = 0.01,
        lr_bounds: tuple[float, float] = (1e-5, 1e-2),
        adaptive_lr_factor: float = 1.5,
        gamma_anneal=True,
        **kwargs,
    ):
        backups.check_mode(self._MODE)
        # Guard the reward-normalization footgun + optionally add obs norm.
        # env is the 2nd positional PPO arg (policy, env, ...) or a kwarg.
        args = list(args)
        if "env" in kwargs:
            kwargs["env"] = guard_and_normalize_env(kwargs["env"], normalize_obs)
        elif len(args) >= 2:
            args[1] = guard_and_normalize_env(args[1], normalize_obs)
        args = tuple(args)

        # GPU-resident path? (detected from the env; see tensor_env.py)
        _env = kwargs.get("env", args[1] if len(args) >= 2 else None)
        self._tensor_path = bool(getattr(_env, "is_tensor_env", False))
        if self._tensor_path and bootstrap_on_timeout:
            raise ValueError(
                "bootstrap_on_timeout=True is not supported on the tensor path "
                "(it is wrong for safety margins in any case)."
            )

        # Default buffer: the pair implementing THIS learner's mode -- numpy or
        # torch depending on the path.
        if rollout_buffer_class is None:
            numpy_cls, tensor_cls = rollout_buffer_classes(self._MODE)
            rollout_buffer_class = tensor_cls if self._tensor_path else numpy_cls

        self.bootstrap_on_timeout = bool(bootstrap_on_timeout)
        self.adaptive_lr = bool(adaptive_lr)
        self.desired_kl = desired_kl
        self.lr_min, self.lr_max = float(lr_bounds[0]), float(lr_bounds[1])
        self.adaptive_lr_factor = float(adaptive_lr_factor)
        if self.adaptive_lr:
            if callable(learning_rate):
                raise ValueError(
                    "adaptive_lr=True requires a float initial learning_rate, "
                    f"got a schedule: {learning_rate!r}"
                )
            # Mutable current LR, adjusted each update from the measured KL.
            self._adaptive_lr = float(learning_rate)

        super().__init__(
            *args,
            learning_rate=learning_rate,
            rollout_buffer_class=rollout_buffer_class,
            rollout_buffer_kwargs=rollout_buffer_kwargs,
            **kwargs,
        )
        # Resolve the gamma-anneal schedule now that super() has set self.gamma.
        self._setup_gamma_anneal(gamma_anneal)

    def _setup_model(self) -> None:
        # Builds policy, optimizer, and rollout buffer.
        super()._setup_model()
        self._check_rollout_buffer(self.rollout_buffer)

    def _check_rollout_buffer(self, buffer) -> None:
        check_rollout_buffer_mode(self, buffer)

    def collect_rollouts(self, *args, **kwargs):
        """Refuse to fall through to SB3's ``OnPolicyAlgorithm.collect_rollouts``.

        This class has no rollout loop by design — that is the Players axis.
        Without this guard, instantiating it directly would inherit the stock
        loop, which bootstraps values on timeout (corrupting a margin reward) and
        never offers the buffer its per-step extras (so a reach-avoid buffer
        would silently back up ``l = 0`` everywhere).
        """
        raise NotImplementedError(
            f"{type(self).__name__} has no rollout loop. Use a concrete learner "
            "(SafetyPPO1P, ReachAvoidPPO2P, ...), or AbstractPPO1P / "
            "AbstractPPO2P.")

    # ---------------------------------------------------------- adaptive LR
    def _update_learning_rate(self, optimizers) -> None:
        """Set the optimizer LR.  With ``adaptive_lr`` use the KL-controlled
        value instead of SB3's progress schedule (called at each ``train()``)."""
        if not self.adaptive_lr:
            super()._update_learning_rate(optimizers)
            return
        self.logger.record("train/learning_rate", self._adaptive_lr)
        update_learning_rate(optimizers, self._adaptive_lr)

    def train(self) -> None:
        super().train()
        # Adjust the LR for the NEXT update from the KL just measured, mirroring
        # rsl_rl: KL >> target -> shrink LR; KL << target -> grow LR.
        if self.adaptive_lr and self.desired_kl is not None:
            kl = self.logger.name_to_value.get("train/approx_kl")
            if kl is not None and kl > 0.0:
                if kl > self.desired_kl * 2.0:
                    self._adaptive_lr = max(
                        self.lr_min, self._adaptive_lr / self.adaptive_lr_factor
                    )
                elif kl < self.desired_kl / 2.0:
                    self._adaptive_lr = min(
                        self.lr_max, self._adaptive_lr * self.adaptive_lr_factor
                    )
            self.logger.record("train/adaptive_lr", self._adaptive_lr)

collect_rollouts

collect_rollouts(*args, **kwargs)

Refuse to fall through to SB3's OnPolicyAlgorithm.collect_rollouts.

This class has no rollout loop by design — that is the Players axis. Without this guard, instantiating it directly would inherit the stock loop, which bootstraps values on timeout (corrupting a margin reward) and never offers the buffer its per-step extras (so a reach-avoid buffer would silently back up l = 0 everywhere).

Source code in safety_sb3/ppo_base.py

def collect_rollouts(self, *args, **kwargs):
    """Refuse to fall through to SB3's ``OnPolicyAlgorithm.collect_rollouts``.

    This class has no rollout loop by design — that is the Players axis.
    Without this guard, instantiating it directly would inherit the stock
    loop, which bootstraps values on timeout (corrupting a margin reward) and
    never offers the buffer its per-step extras (so a reach-avoid buffer
    would silently back up ``l = 0`` everywhere).
    """
    raise NotImplementedError(
        f"{type(self).__name__} has no rollout loop. Use a concrete learner "
        "(SafetyPPO1P, ReachAvoidPPO2P, ...), or AbstractPPO1P / "
        "AbstractPPO2P.")

safety_sb3.ppo_1p

Single-player PPO — the 1P half of the PPO family, and its three modes.

AbstractPPO                 (ppo_base.py — recipe, buffers, adaptive LR)
└─ AbstractPPO1P            one actor, one value net, one rollout buffer
    ├─ SafetyPPO1P          _MODE = AVOID
    ├─ ReachAvoidPPO1P      _MODE = REACH_AVOID   (+ _ReachAvoidPlumbing)
    └─ CumulativePPO1P      _MODE = CUMULATIVE

Everything specific to one player lives here: the two rollout loops (numpy / GPU-resident), which are SB3's OnPolicyAlgorithm.collect_rollouts with the timeout value-bootstrap gated off and one added call — the buffer is handed each step's extras so it can keep whatever its own operator needs (see :mod:safety_sb3.buffers_rollout).

The three concrete learners below are one line each. That is the point of the split: the mode axis costs a line, the players axis costs a loop.

AbstractPPO1P

Bases: AbstractPPO

PPO with the ordinary single-actor rollout, backup chosen by _MODE.

Source code in safety_sb3/ppo_1p.py

class AbstractPPO1P(AbstractPPO):
    """PPO with the ordinary single-actor rollout, backup chosen by ``_MODE``."""

    def _collect_rollouts_tensor(
        self,
        env,
        callback: BaseCallback,
        rollout_buffer,
        n_rollout_steps: int,
    ) -> bool:
        """GPU-resident rollout: policy forward, env step, buffer add and the
        backup all stay on device. Episode stats and env ``metrics()``
        (curriculum levels etc.) are recorded to the logger directly."""
        assert self._last_obs is not None
        self.policy.set_training_mode(False)
        rollout_buffer.reset()
        callback.on_rollout_start()

        dev = env.device
        obs = self._last_obs
        if not th.is_tensor(obs):  # first call after _setup_learn
            obs = th.as_tensor(np.asarray(obs), dtype=th.float32, device=dev)
        episode_starts = th.as_tensor(
            np.asarray(self._last_episode_starts, dtype=np.float32), device=dev
        ) if not th.is_tensor(self._last_episode_starts) else self._last_episode_starts

        low = th.as_tensor(self.action_space.low, dtype=th.float32, device=dev)
        high = th.as_tensor(self.action_space.high, dtype=th.float32, device=dev)

        ep_ret = getattr(self, "_t_ep_ret", None)
        if ep_ret is None or ep_ret.shape[0] != env.num_envs:
            self._t_ep_ret = th.zeros(env.num_envs, device=dev)
            self._t_ep_len = th.zeros(env.num_envs, device=dev)
        fin_ret, fin_len = [], []

        n_steps = 0
        while n_steps < n_rollout_steps:
            with th.no_grad():
                actions, values, log_probs = self.policy(obs)
            clipped = th.clamp(actions, low, high)

            new_obs, rewards, dones, timeouts, l_x = env.step_tensor(clipped)
            self.num_timesteps += env.num_envs
            n_steps += 1

            callback.update_locals(locals())
            if not callback.on_step():
                return False

            rollout_buffer.record_extras(l_x)
            rollout_buffer.add(obs, actions, rewards, episode_starts,
                               values.flatten(), log_probs)

            self._t_ep_ret += rewards
            self._t_ep_len += 1.0
            if bool(dones.any()):
                d = dones.bool()
                fin_ret.append(self._t_ep_ret[d])
                fin_len.append(self._t_ep_len[d])
                self._t_ep_ret = th.where(d, th.zeros_like(self._t_ep_ret), self._t_ep_ret)
                self._t_ep_len = th.where(d, th.zeros_like(self._t_ep_len), self._t_ep_len)

            obs = new_obs
            episode_starts = dones.float()

        with th.no_grad():
            last_values = self.policy.predict_values(obs)
        rollout_buffer.compute_returns_and_advantage(
            last_values=last_values.flatten(), dones=dones.float())

        self._last_obs = obs
        self._last_episode_starts = episode_starts

        if fin_ret:
            self.logger.record("rollout/ep_rew_mean", float(th.cat(fin_ret).mean()))
            self.logger.record("rollout/ep_len_mean", float(th.cat(fin_len).mean()))
        for k, v in (env.metrics() or {}).items():
            self.logger.record(f"env/{k}", float(v))

        callback.update_locals(locals())
        callback.on_rollout_end()
        return True

    def collect_rollouts(
        self,
        env: VecEnv,
        callback: BaseCallback,
        rollout_buffer: RolloutBuffer,
        n_rollout_steps: int,
    ) -> bool:
        """SB3 ``OnPolicyAlgorithm.collect_rollouts`` with two safety changes:
        the timeout value-bootstrap is gated by ``self.bootstrap_on_timeout``
        (off by default — the reward is the physical margin g(s)), and the buffer
        is offered each step's ``infos`` before ``add()`` so it can capture the
        extras its operator needs. On the GPU-resident path this dispatches to
        ``_collect_rollouts_tensor``."""
        # Anneal gamma (buffer.gamma) for THIS rollout's GAE, then collect. The
        # tensor path bypasses SB3's _update_current_progress_remaining, so apply
        # it here explicitly (idempotent on the numpy path).
        self._apply_gamma_anneal()
        if self._tensor_path:
            return self._collect_rollouts_tensor(
                env, callback, rollout_buffer, n_rollout_steps)
        assert self._last_obs is not None, "No previous observation was provided"
        self.policy.set_training_mode(False)

        n_steps = 0
        rollout_buffer.reset()
        if self.use_sde:
            self.policy.reset_noise(env.num_envs)

        callback.on_rollout_start()

        while n_steps < n_rollout_steps:
            if (
                self.use_sde
                and self.sde_sample_freq > 0
                and n_steps % self.sde_sample_freq == 0
            ):
                self.policy.reset_noise(env.num_envs)

            with th.no_grad():
                obs_tensor = obs_as_tensor(self._last_obs, self.device)
                actions, values, log_probs = self.policy(obs_tensor)
            actions = actions.cpu().numpy()

            clipped_actions = actions
            if isinstance(self.action_space, spaces.Box):
                if self.policy.squash_output:
                    clipped_actions = self.policy.unscale_action(clipped_actions)
                else:
                    clipped_actions = np.clip(
                        actions, self.action_space.low, self.action_space.high
                    )

            new_obs, rewards, dones, infos = env.step(clipped_actions)
            self.num_timesteps += env.num_envs

            callback.update_locals(locals())
            if not callback.on_step():
                return False

            self._update_info_buffer(infos, dones)
            n_steps += 1

            if isinstance(self.action_space, spaces.Discrete):
                actions = actions.reshape(-1, 1)

            # Timeout bootstrapping — DISABLED by default for safety (see the
            # module docstring). Stock PPO would corrupt g(s) here.
            if self.bootstrap_on_timeout:
                for idx, done in enumerate(dones):
                    if (
                        done
                        and infos[idx].get("terminal_observation") is not None
                        and infos[idx].get("TimeLimit.truncated", False)
                    ):
                        terminal_obs = self.policy.obs_to_tensor(
                            infos[idx]["terminal_observation"]
                        )[0]
                        with th.no_grad():
                            terminal_value = self.policy.predict_values(terminal_obs)[0]
                        rewards[idx] += self.gamma * terminal_value

            rollout_buffer.record_extras(infos)

            rollout_buffer.add(
                self._last_obs,
                actions,
                rewards,
                self._last_episode_starts,
                values,
                log_probs,
            )
            self._last_obs = new_obs
            self._last_episode_starts = dones

        with th.no_grad():
            values = self.policy.predict_values(obs_as_tensor(new_obs, self.device))

        rollout_buffer.compute_returns_and_advantage(last_values=values, dones=dones)

        callback.update_locals(locals())
        callback.on_rollout_end()
        return True

collect_rollouts

collect_rollouts(
    env, callback, rollout_buffer, n_rollout_steps
)

SB3 OnPolicyAlgorithm.collect_rollouts with two safety changes: the timeout value-bootstrap is gated by self.bootstrap_on_timeout (off by default — the reward is the physical margin g(s)), and the buffer is offered each step's infos before add() so it can capture the extras its operator needs. On the GPU-resident path this dispatches to _collect_rollouts_tensor.

Source code in safety_sb3/ppo_1p.py

def collect_rollouts(
    self,
    env: VecEnv,
    callback: BaseCallback,
    rollout_buffer: RolloutBuffer,
    n_rollout_steps: int,
) -> bool:
    """SB3 ``OnPolicyAlgorithm.collect_rollouts`` with two safety changes:
    the timeout value-bootstrap is gated by ``self.bootstrap_on_timeout``
    (off by default — the reward is the physical margin g(s)), and the buffer
    is offered each step's ``infos`` before ``add()`` so it can capture the
    extras its operator needs. On the GPU-resident path this dispatches to
    ``_collect_rollouts_tensor``."""
    # Anneal gamma (buffer.gamma) for THIS rollout's GAE, then collect. The
    # tensor path bypasses SB3's _update_current_progress_remaining, so apply
    # it here explicitly (idempotent on the numpy path).
    self._apply_gamma_anneal()
    if self._tensor_path:
        return self._collect_rollouts_tensor(
            env, callback, rollout_buffer, n_rollout_steps)
    assert self._last_obs is not None, "No previous observation was provided"
    self.policy.set_training_mode(False)

    n_steps = 0
    rollout_buffer.reset()
    if self.use_sde:
        self.policy.reset_noise(env.num_envs)

    callback.on_rollout_start()

    while n_steps < n_rollout_steps:
        if (
            self.use_sde
            and self.sde_sample_freq > 0
            and n_steps % self.sde_sample_freq == 0
        ):
            self.policy.reset_noise(env.num_envs)

        with th.no_grad():
            obs_tensor = obs_as_tensor(self._last_obs, self.device)
            actions, values, log_probs = self.policy(obs_tensor)
        actions = actions.cpu().numpy()

        clipped_actions = actions
        if isinstance(self.action_space, spaces.Box):
            if self.policy.squash_output:
                clipped_actions = self.policy.unscale_action(clipped_actions)
            else:
                clipped_actions = np.clip(
                    actions, self.action_space.low, self.action_space.high
                )

        new_obs, rewards, dones, infos = env.step(clipped_actions)
        self.num_timesteps += env.num_envs

        callback.update_locals(locals())
        if not callback.on_step():
            return False

        self._update_info_buffer(infos, dones)
        n_steps += 1

        if isinstance(self.action_space, spaces.Discrete):
            actions = actions.reshape(-1, 1)

        # Timeout bootstrapping — DISABLED by default for safety (see the
        # module docstring). Stock PPO would corrupt g(s) here.
        if self.bootstrap_on_timeout:
            for idx, done in enumerate(dones):
                if (
                    done
                    and infos[idx].get("terminal_observation") is not None
                    and infos[idx].get("TimeLimit.truncated", False)
                ):
                    terminal_obs = self.policy.obs_to_tensor(
                        infos[idx]["terminal_observation"]
                    )[0]
                    with th.no_grad():
                        terminal_value = self.policy.predict_values(terminal_obs)[0]
                    rewards[idx] += self.gamma * terminal_value

        rollout_buffer.record_extras(infos)

        rollout_buffer.add(
            self._last_obs,
            actions,
            rewards,
            self._last_episode_starts,
            values,
            log_probs,
        )
        self._last_obs = new_obs
        self._last_episode_starts = dones

    with th.no_grad():
        values = self.policy.predict_values(obs_as_tensor(new_obs, self.device))

    rollout_buffer.compute_returns_and_advantage(last_values=values, dones=dones)

    callback.update_locals(locals())
    callback.on_rollout_end()
    return True

SafetyPPO1P

Bases: AbstractPPO1P

On-policy avoid RL — Fisac et al. 2019::

V(s) = (1-γ)·g + γ·min(g, V')          terminal: V(s) = g

g(s) is the safety margin and rides on the reward channel; V(s) >= 0 iff the state is in the safe set. There is no target set and no l: this learner never reads info["l_x"] and its buffer has nowhere to put it. For a task with a target to reach use :class:ReachAvoidPPO1P — the two operators are not interchangeable (see :mod:safety_sb3.backups).

Source code in safety_sb3/ppo_1p.py

class SafetyPPO1P(AbstractPPO1P):
    """On-policy **avoid** RL — Fisac et al. 2019::

        V(s) = (1-γ)·g + γ·min(g, V')          terminal: V(s) = g

    ``g(s)`` is the safety margin and rides on the reward channel; ``V(s) >= 0``
    iff the state is in the safe set. There is no target set and no ``l``: this
    learner never reads ``info["l_x"]`` and its buffer has nowhere to put it.
    For a task with a target to reach use :class:`ReachAvoidPPO1P` — the two
    operators are not interchangeable (see :mod:`safety_sb3.backups`).
    """

    _MODE = backups.AVOID

ReachAvoidPPO1P

Bases: _ReachAvoidPlumbing, AbstractPPO1P

On-policy reach-avoid RL — Hsu et al. RSS'21 eq. 15::

V(s) = (1-γ)·min(l, g) + γ·min(g, max(l, V'))

g (safety margin) rides on the reward channel; l (target margin) is read from info["l_x"] on the numpy path and returned directly by TensorVecEnv.step_tensor on the GPU-resident path — in both cases by the buffer, not by this class.

Unlike the SAC family this is fully on-policy: it pairs with vectorized envs at large n_envs (the GPU-parallel regime where off-policy reach-avoid saturates; validated on a 1280-env MuJoCo quadruped in unitree_rl_mjlab).

Source code in safety_sb3/ppo_1p.py

class ReachAvoidPPO1P(_ReachAvoidPlumbing, AbstractPPO1P):
    """On-policy **reach-avoid** RL — Hsu et al. RSS'21 eq. 15::

        V(s) = (1-γ)·min(l, g) + γ·min(g, max(l, V'))

    ``g`` (safety margin) rides on the reward channel; ``l`` (target margin) is
    read from ``info["l_x"]`` on the numpy path and returned directly by
    ``TensorVecEnv.step_tensor`` on the GPU-resident path — in both cases by the
    *buffer*, not by this class.

    Unlike the SAC family this is fully on-policy: it pairs with vectorized envs
    at large ``n_envs`` (the GPU-parallel regime where off-policy reach-avoid
    saturates; validated on a 1280-env MuJoCo quadruped in ``unitree_rl_mjlab``).
    """

    _MODE = backups.REACH_AVOID

CumulativePPO1P

Bases: AbstractPPO1P

Ordinary reward-maximizing PPO — V(s) = r + γ·V(s').

Not a safety learner: the reward channel carries a reward, not a margin, and V >= 0 means nothing. It is here so a nominal baseline runs through every line of the same code as the safety learners — the standard control for "is the safety operator doing the work, or is it just PPO?".

Source code in safety_sb3/ppo_1p.py

class CumulativePPO1P(AbstractPPO1P):
    """Ordinary reward-maximizing PPO — ``V(s) = r + γ·V(s')``.

    Not a safety learner: the reward channel carries a reward, not a margin, and
    ``V >= 0`` means nothing. It is here so a *nominal* baseline runs through
    every line of the same code as the safety learners — the standard control for
    "is the safety operator doing the work, or is it just PPO?".
    """

    _MODE = backups.CUMULATIVE

safety_sb3.ppo_2p

Two-player (adversarial) PPO — the 2P half of the PPO family.

AbstractPPO                 (ppo_base.py)
└─ AbstractPPO2P            two actors, two value nets, two buffers, phases
    ├─ SafetyPPO2P          _MODE = AVOID
    └─ ReachAvoidPPO2P      _MODE = REACH_AVOID   (+ _ReachAvoidPlumbing)

The control player MAXIMIZES the value, the disturbance player MINIMIZES it, and a league of archived opponents damps cycling. Both players are ordinary PPO policies over disjoint SUB-spaces of the env's concatenated action space (the env exposes one Box(ctrl_dim + dstb_dim) and splits it; g rides on the reward channel, l — reach-avoid only — on info["l_x"]).

• the targets are computed on-policy per rollout, inside the rollout buffer; the min player trains on the SAME targets with a NEGATED advantage (negation commutes with advantage normalization) — the zero-sum property;
• training alternates in phases (dstb pretrain, then K dstb / M ctrl rollout cycles); the frozen player acts stochastically;
• with n_envs > 1, league opponents are assigned to env SLICES so a single on-policy batch contains rollouts against the whole archived population (validated at 1280 GPU-parallel envs in unitree_rl_mjlab; off-policy two-player SAC saturates at ~32-64 envs).

self.policy is always the CONTROL policy — predict(), save(), and downstream filter wrappers see the deployable controller.

Why this is NOT the same object as :mod:safety_sb3.sac_2p

Both classes hold two actors, and it is tempting to read them as one algorithm with two optimizers. They are not, and the difference is forced by where each family's critic takes its action:

• SAC 2P has ONE twin critic over the JOINT action Q(s, [a_ctrl, a_dstb]), one replay buffer, and continuous updates. Because the critic takes the action, that single Q is the game value, and both actors differentiate through the same object — ctrl ascends it, dstb descends it. Nothing needs to be scheduled, and replay makes the learner indifferent to whose policy generated the data. This is the minimax game as ISAACS formulates it.
• PPO 2P (here) has TWO independent state-only value nets V(s), two rollout buffers, and a phase machine. A state-only critic cannot represent "value of this joint action", so each player needs its own advantage, its own importance ratio, and — because PPO is on-policy — its own freshly collected data. That is exactly what the phase machine exists for, and exactly why SAC needs no phases. This is an alternating best-response approximation, not the simultaneous saddle point.

They are therefore not interchangeable, and there is deliberately no cross-family AbstractTwoPlayer base: the two differ in their state — one critic vs two, one buffer vs two, phases vs none — not merely in their update loops. A shared parent could only hold the names.

AbstractPPO2P

Bases: AbstractPPO

Two-player on-policy game with league opponents, backup chosen by _MODE.

Parameters:

Name	Type	Description	Default
ctrl_action_dim	int	number of leading action dims belonging to the control player; the rest are the disturbance.	-1
dstb_pretrain_rollouts	int	rollouts spent training only the disturbance before the alternation starts.	20
ctrl_rollouts_per_cycle	int	control-player rollouts per alternation cycle.	4
dstb_rollouts_per_cycle	int	disturbance-player rollouts per alternation cycle.	1
use_leaderboard	bool	archive opponents and assign them to env slices.	False
dstb_learning_rate	float | None	disturbance-player learning rate; None inherits the ctrl value.	None
dstb_ent_coef	float | None	disturbance-player entropy coefficient; None inherits the ctrl value.	None

Source code in safety_sb3/ppo_2p.py

class AbstractPPO2P(AbstractPPO):
  """Two-player on-policy game with league opponents, backup chosen by ``_MODE``.

  :param ctrl_action_dim: number of leading action dims belonging to the control
      player; the rest are the disturbance.
  :param dstb_pretrain_rollouts: rollouts spent training only the disturbance
      before the alternation starts.
  :param ctrl_rollouts_per_cycle: control-player rollouts per alternation cycle.
  :param dstb_rollouts_per_cycle: disturbance-player rollouts per alternation cycle.
  :param use_leaderboard: archive opponents and assign them to env slices.
  :param dstb_learning_rate: disturbance-player learning rate; ``None`` inherits
      the ctrl value.
  :param dstb_ent_coef: disturbance-player entropy coefficient; ``None`` inherits
      the ctrl value.
  """

  def __init__(
    self,
    policy,
    env,
    # Sentinel default keeps SB3 ``load()`` working: load() constructs the
    # class BARE (cls(policy, env=None, _init_setup_model=False)) then restores
    # __dict__, so the saved ctrl_action_dim replaces the sentinel before
    # _setup_model() ever runs. Direct construction must pass a real value
    # (validated in _split_action_space).
    ctrl_action_dim: int = -1,
    dstb_pretrain_rollouts: int = 20,
    ctrl_rollouts_per_cycle: int = 4,
    dstb_rollouts_per_cycle: int = 1,
    use_leaderboard: bool = False,
    leaderboard_dir: str = "ppo_2p_leaderboard",
    save_top_k_ctrl: int = 5,
    save_top_k_dstb: int = 5,
    softmax_rationality: float = 3.0,
    leaderboard_freq_cycles: int = 1,
    num_slices: int = 8,
    dstb_learning_rate: float | None = None,
    dstb_ent_coef: float | None = None,
    **kwargs,
  ) -> None:
    # Per-player optimization: the two players sit in different optimization
    # regimes (ctrl is usually warm-started and precise; dstb is fresh and
    # exploratory). The reference setup used ctrl lr 5e-4 / dstb 3e-4 and
    # ctrl ent 5e-4 / dstb 2e-3. None -> inherit the ctrl value.
    self._dstb_lr = None if dstb_learning_rate is None else float(dstb_learning_rate)
    self._dstb_ent = None if dstb_ent_coef is None else float(dstb_ent_coef)
    self.ctrl_action_dim = int(ctrl_action_dim)
    self._dstb_pretrain = int(dstb_pretrain_rollouts)
    self._ctrl_per_cycle = int(ctrl_rollouts_per_cycle)
    self._dstb_per_cycle = int(dstb_rollouts_per_cycle)
    self._use_lb = bool(use_leaderboard)
    self._lb_dir = leaderboard_dir
    self._lb_kc = int(save_top_k_ctrl)
    self._lb_kd = int(save_top_k_dstb)
    self._lb_rationality = float(softmax_rationality)
    self._lb_freq_cycles = int(leaderboard_freq_cycles)
    self._n_slices = int(num_slices)
    self._rollouts_done = 0
    super().__init__(policy, env, **kwargs)

  # --- spaces / model setup --------------------------------------------------

  def _split_action_space(self):
    full = self.action_space
    assert isinstance(full, spaces.Box) and len(full.shape) == 1
    c = self.ctrl_action_dim
    assert c > 0, (
      f"ctrl_action_dim not set — pass it to {type(self).__name__}(...) (the -1 "
      "default exists only so SB3 load() can construct the class bare).")
    ctrl = spaces.Box(low=full.low[:c], high=full.high[:c], dtype=full.dtype)
    dstb = spaces.Box(low=full.low[c:], high=full.high[c:], dtype=full.dtype)
    return ctrl, dstb

  def _setup_model(self) -> None:
    # Build the CTRL policy/buffer over the ctrl SUB-space by temporarily
    # narrowing the action space seen by the base setup.
    full_space = self.action_space
    ctrl_space, dstb_space = self._split_action_space()
    self.action_space = ctrl_space
    super()._setup_model()
    self.action_space = full_space
    self._full_action_space = full_space
    self._ctrl_space, self._dstb_space = ctrl_space, dstb_space

    # Min player: same policy class/kwargs over the dstb sub-space.
    self.dstb_policy = self.policy_class(
      self.observation_space, dstb_space, self.lr_schedule,
      use_sde=self.use_sde, **self.policy_kwargs,
    ).to(self.device)
    # Both buffers must carry THIS learner's backup, so both are looked up from
    # _MODE (hardcoding the reach-avoid buffer here once gave the avoid learner's
    # min player the wrong game), and both get the same buffer kwargs (the ctrl
    # buffer used to receive terminal_type and the dstb buffer silently did not).
    numpy_cls, tensor_cls = rollout_buffer_classes(self._MODE)
    dstb_buf_cls = tensor_cls if self._tensor_path else numpy_cls
    self.dstb_rollout_buffer = dstb_buf_cls(
      self.n_steps, self.observation_space, dstb_space,
      device=self.device, gamma=self.gamma, gae_lambda=self.gae_lambda,
      n_envs=self.n_envs, **(self.rollout_buffer_kwargs or {}),
    )
    self._check_rollout_buffer(self.dstb_rollout_buffer)

    self._leaderboard: Leaderboard | None = None
    if self._use_lb:
      self._leaderboard = Leaderboard(
        self._lb_kc, self._lb_kd, self._lb_rationality, self._lb_dir
      )
      self._scratch_dstb = [
        self.policy_class(
          self.observation_space, dstb_space, self.lr_schedule,
          use_sde=self.use_sde, **self.policy_kwargs,
        ).to(self.device)
        for _ in range(self._lb_kd)
      ]
    self._slice_opps: list[int] = [_ZERO, _RANDOM] + [_CURRENT] * max(
      self._n_slices - 2, 0
    )
    n = self.n_envs
    self._slice_of_env = (np.arange(n) * self._n_slices // max(n, 1)).clip(
      max=self._n_slices - 1
    )
    # Per-episode league flags (numpy path).
    self._ever_gneg = np.zeros(n, dtype=bool)
    self._alloc_reach_flags(self._ever_gneg)

    # Per-player KL-adaptive LR state: a SHARED controller cross-contaminates
    # (a dstb phase's KL would move the LR the next ctrl phase trains with,
    # and vice versa). Each player keeps its own scalar; train() swaps the
    # active one into AbstractPPO's controller and stores it back after.
    if self.adaptive_lr:
      self._adaptive_lr_ctrl = float(self._adaptive_lr)
      self._adaptive_lr_dstb = float(self._dstb_lr if self._dstb_lr is not None
                                     else self._adaptive_lr)
    if self._dstb_lr is not None:
      for grp in self.dstb_policy.optimizer.param_groups:
        grp["lr"] = self._dstb_lr

  # --- phase machine -----------------------------------------------------------

  def _phase(self) -> str:
    if self._rollouts_done < self._dstb_pretrain:
      return "dstb"
    k = (self._rollouts_done - self._dstb_pretrain) % (
      self._dstb_per_cycle + self._ctrl_per_cycle
    )
    return "dstb" if k < self._dstb_per_cycle else "ctrl"

  def _cycle_len(self) -> int:
    return self._dstb_per_cycle + self._ctrl_per_cycle

  # --- league episode outcome --------------------------------------------------
  # The win condition is "ctrl survived to the time limit without ever violating
  # g", AND -- in reach-avoid only -- "reached the target at some point". The
  # reach conjunct is supplied by _ReachAvoidPlumbing; the no-ops below are the
  # avoid game's answer, which is that survival IS the win because there is no
  # target set. Nothing here has to be switched back off.

  def _alloc_reach_flags(self, like) -> None:
    """Allocate extra per-episode flags the mode's win condition needs."""

  def _note_reach(self, step_l) -> None:
    """Record per-step target-margin evidence (``infos`` or an ``l_x`` tensor)."""

  def _reach_flag(self):
    """Extra conjunct in the win condition; ``True`` == no extra requirement."""
    return True

  def _clear_reach(self, dones) -> None:
    """Clear the flags of episodes that just ended."""

  def _episode_success(self, dones, timeouts):
    """Training outcome scored on the league board (numpy path)."""
    return dones & timeouts & ~self._ever_gneg & self._reach_flag()

  def _episode_success_tensor(self, dones, timeouts):
    """Torch twin of :meth:`_episode_success`."""
    return dones & timeouts & ~self._ever_gneg_t & self._reach_flag()

  # --- opponents ---------------------------------------------------------------

  def _resample_slice_opponents(self) -> None:
    if self._leaderboard is None:
      self._slice_opps = [_ZERO, _RANDOM] + [_CURRENT] * max(self._n_slices - 2, 0)
      return
    self._slice_opps = self._leaderboard.sample_dstb_slices(self._n_slices)
    self._loaded_scratch: dict[int, th.nn.Module] = {}
    for opp in {o for o in self._slice_opps if o >= 0}:
      scratch = self._scratch_dstb[len(self._loaded_scratch) % len(self._scratch_dstb)]
      self._leaderboard.load_actor(
        scratch, "dstb", self._leaderboard.dstb_steps[opp]
      )
      scratch.to(self.device)
      self._loaded_scratch[opp] = scratch

  def _dstb_actions_for_ctrl_phase(self, obs_tensor) -> np.ndarray:
    """Per-slice opponent disturbance (stochastic), in dstb sub-space units."""
    n = self.n_envs
    out = np.zeros((n, self._dstb_space.shape[0]), dtype=np.float32)
    cur_cache = None
    for si, opp in enumerate(self._slice_opps):
      mask = self._slice_of_env == si
      if not mask.any() or opp == _ZERO:
        continue
      if opp == _RANDOM:
        out[mask] = np.random.uniform(
          self._dstb_space.low, self._dstb_space.high,
          size=(int(mask.sum()), self._dstb_space.shape[0]),
        )
      elif opp == _CURRENT:
        if cur_cache is None:
          with th.no_grad():
            a, _, _ = self.dstb_policy(obs_tensor)
          cur_cache = a.cpu().numpy()
        out[mask] = cur_cache[mask]
      else:
        with th.no_grad():
          a, _, _ = self._loaded_scratch[opp](obs_tensor)
        out[mask] = a.cpu().numpy()[mask]
    return out

  def _dstb_actions_tensor(self, obs: th.Tensor) -> th.Tensor:
    """Per-slice opponent disturbance, torch end-to-end (ctrl phase)."""
    n = self.n_envs
    dev = obs.device
    d = self._dstb_space.shape[0]
    lo = th.as_tensor(self._dstb_space.low, dtype=th.float32, device=dev)
    hi = th.as_tensor(self._dstb_space.high, dtype=th.float32, device=dev)
    out = th.zeros(n, d, device=dev)
    slice_of_env = getattr(self, "_slice_of_env_t", None)
    if slice_of_env is None or slice_of_env.shape[0] != n:
      self._slice_of_env_t = th.as_tensor(self._slice_of_env, device=dev)
      slice_of_env = self._slice_of_env_t
    cur = None
    for si, opp in enumerate(self._slice_opps):
      mask = slice_of_env == si
      if not bool(mask.any()) or opp == _ZERO:
        continue
      if opp == _RANDOM:
        out[mask] = lo + (hi - lo) * th.rand(int(mask.sum()), d, device=dev)
      elif opp == _CURRENT:
        if cur is None:
          with th.no_grad():
            cur, _, _ = self.dstb_policy(obs)
        out[mask] = cur[mask]
      else:
        with th.no_grad():
          a, _, _ = self._loaded_scratch[opp](obs)
        out[mask] = a[mask]
    return out

  def _collect_rollouts_tensor(self, env, callback, rollout_buffer,
                               n_rollout_steps: int) -> bool:
    """Torch twin of the two-player rollout: both actions each step, env gets
    the concatenation, only the ACTIVE player's data is stored."""
    del rollout_buffer
    phase = self._phase()
    active = self.dstb_policy if phase == "dstb" else self.policy
    passive = self.policy if phase == "dstb" else self.dstb_policy
    buf = self.dstb_rollout_buffer if phase == "dstb" else self.rollout_buffer

    assert self._last_obs is not None
    active.set_training_mode(False)
    passive.set_training_mode(False)
    if phase == "ctrl":
      self._resample_slice_opponents()

    dev = env.device
    obs = self._last_obs
    if not th.is_tensor(obs):
      obs = th.as_tensor(np.asarray(obs), dtype=th.float32, device=dev)
    episode_starts = (self._last_episode_starts
                      if th.is_tensor(self._last_episode_starts)
                      else th.as_tensor(np.asarray(
                        self._last_episode_starts, dtype=np.float32), device=dev))
    c_lo = th.as_tensor(self._ctrl_space.low, dtype=th.float32, device=dev)
    c_hi = th.as_tensor(self._ctrl_space.high, dtype=th.float32, device=dev)
    d_lo = th.as_tensor(self._dstb_space.low, dtype=th.float32, device=dev)
    d_hi = th.as_tensor(self._dstb_space.high, dtype=th.float32, device=dev)
    ever_gneg = getattr(self, "_ever_gneg_t", None)
    if ever_gneg is None or ever_gneg.shape[0] != env.num_envs:
      self._ever_gneg_t = th.zeros(env.num_envs, dtype=th.bool, device=dev)
      self._alloc_reach_flags(self._ever_gneg_t)
    slice_of_env = getattr(self, "_slice_of_env_t", None)
    if slice_of_env is None:
      self._slice_of_env_t = th.as_tensor(self._slice_of_env, device=dev)

    buf.reset()
    callback.on_rollout_start()
    n_steps = 0
    while n_steps < n_rollout_steps:
      with th.no_grad():
        actions, values, log_probs = active(obs)
      if phase == "dstb":
        with th.no_grad():
          ctrl_a, _, _ = passive(obs)
        dstb_a = actions
      else:
        ctrl_a = actions
        dstb_a = self._dstb_actions_tensor(obs)
      full = th.cat([th.clamp(ctrl_a, c_lo, c_hi),
                     th.clamp(dstb_a, d_lo, d_hi)], dim=1)

      new_obs, rewards, dones, timeouts, l_x = env.step_tensor(full)
      self.num_timesteps += env.num_envs
      n_steps += 1
      callback.update_locals(locals())
      if not callback.on_step():
        return False

      buf.record_extras(l_x)  # the buffer keeps what its operator needs
      buf.add(obs, actions, rewards, episode_starts,
              values.flatten(), log_probs)

      self._note_reach(l_x)
      self._ever_gneg_t |= rewards < 0.0
      if self._leaderboard is not None and bool(dones.any()):
        d_b = dones.bool()
        succ = self._episode_success_tensor(d_b, timeouts.bool())
        b = self._leaderboard.board
        if phase == "ctrl":
          for si, opp in enumerate(self._slice_opps):
            m = d_b & (self._slice_of_env_t == si)
            if not bool(m.any()):
              continue
            col = (b.shape[1] - 1 if opp == _ZERO
                   else b.shape[1] - 2 if opp in (_RANDOM, _CURRENT) else opp)
            self._leaderboard.ema_score(-1, col, float(succ[m].float().mean()))
        else:
          self._leaderboard.ema_score(
            -1, b.shape[1] - 2, float(succ[d_b].float().mean()))
        self._clear_reach(d_b)
        self._ever_gneg_t = th.where(d_b, th.zeros_like(self._ever_gneg_t),
                                     self._ever_gneg_t)

      obs = new_obs
      episode_starts = dones.float()

    with th.no_grad():
      last_values = active.predict_values(obs)
    buf.compute_returns_and_advantage(last_values=last_values.flatten(),
                                      dones=dones.float())
    self._last_obs = obs
    self._last_episode_starts = episode_starts
    for k, v in (env.metrics() or {}).items():
      self.logger.record(f"env/{k}", float(v))
    callback.update_locals(locals())
    callback.on_rollout_end()
    return True

  # --- rollout collection --------------------------------------------------------

  def collect_rollouts(
    self,
    env: VecEnv,
    callback: BaseCallback,
    rollout_buffer: RolloutBuffer,
    n_rollout_steps: int,
  ) -> bool:
    """Two-player rollout: both actions computed each step, env receives the
    concatenation, only the ACTIVE player's data is stored (in its buffer)."""
    # This override does not call the AbstractPPO1P loop, so anneal gamma here;
    # the helper updates both the ctrl and dstb rollout buffers (idempotent).
    self._apply_gamma_anneal()
    if self._tensor_path:
      return self._collect_rollouts_tensor(env, callback, rollout_buffer,
                                           n_rollout_steps)
    del rollout_buffer  # phase decides the buffer
    phase = self._phase()
    active_policy = self.dstb_policy if phase == "dstb" else self.policy
    passive_policy = self.policy if phase == "dstb" else self.dstb_policy
    buf = self.dstb_rollout_buffer if phase == "dstb" else self.rollout_buffer

    assert self._last_obs is not None
    active_policy.set_training_mode(False)
    passive_policy.set_training_mode(False)
    if phase == "ctrl":
      self._resample_slice_opponents()

    n_steps = 0
    buf.reset()
    callback.on_rollout_start()

    while n_steps < n_rollout_steps:
      with th.no_grad():
        obs_tensor = obs_as_tensor(self._last_obs, self.device)
        actions, values, log_probs = active_policy(obs_tensor)
      actions = actions.cpu().numpy()

      if phase == "dstb":
        with th.no_grad():
          ctrl_a, _, _ = passive_policy(obs_tensor)  # frozen ctrl, stochastic
        ctrl_np = ctrl_a.cpu().numpy()
        dstb_np = actions
      else:
        ctrl_np = actions
        dstb_np = self._dstb_actions_for_ctrl_phase(obs_tensor)

      ctrl_clipped = np.clip(ctrl_np, self._ctrl_space.low, self._ctrl_space.high)
      dstb_clipped = np.clip(dstb_np, self._dstb_space.low, self._dstb_space.high)
      full_actions = np.concatenate([ctrl_clipped, dstb_clipped], axis=1)

      new_obs, rewards, dones, infos = env.step(full_actions)
      self.num_timesteps += env.num_envs

      callback.update_locals(locals())
      if not callback.on_step():
        return False
      self._update_info_buffer(infos, dones)
      n_steps += 1

      # Timeout bootstrap (ACTIVE player's value fn) — DISABLED by default:
      # the reward is the physical margin g(s) (see the AbstractPPO docstring).
      if self.bootstrap_on_timeout:
        for idx, done in enumerate(dones):
          if (
            done
            and infos[idx].get("terminal_observation") is not None
            and infos[idx].get("TimeLimit.truncated", False)
          ):
            terminal_obs = active_policy.obs_to_tensor(
              infos[idx]["terminal_observation"]
            )[0]
            with th.no_grad():
              terminal_value = active_policy.predict_values(terminal_obs)[0]
            rewards[idx] += self.gamma * terminal_value

      buf.record_extras(infos)  # the buffer keeps what its operator needs
      buf.add(
        self._last_obs, actions, rewards,
        self._last_episode_starts, values, log_probs,
      )

      # episode flags -> league board (training outcomes)
      self._note_reach(infos)
      self._ever_gneg |= rewards < 0.0
      if self._leaderboard is not None and dones.any():
        timeouts = np.array(
          [bool(info.get("TimeLimit.truncated", False)) for info in infos]
        )
        succ = self._episode_success(dones, timeouts)
        b = self._leaderboard.board
        if phase == "ctrl":
          for si, opp in enumerate(self._slice_opps):
            m = dones & (self._slice_of_env == si)
            if not m.any():
              continue
            col = (b.shape[1] - 1 if opp == _ZERO
                   else b.shape[1] - 2 if opp in (_RANDOM, _CURRENT) else opp)
            self._leaderboard.ema_score(-1, col, float(succ[m].mean()))
        else:
          self._leaderboard.ema_score(
            -1, b.shape[1] - 2, float(succ[dones].mean())
          )
        self._clear_reach(dones)
        self._ever_gneg[dones] = False

      self._last_obs = new_obs
      self._last_episode_starts = dones

    with th.no_grad():
      values = active_policy.predict_values(obs_as_tensor(new_obs, self.device))
    buf.compute_returns_and_advantage(last_values=values, dones=dones)

    callback.update_locals(locals())
    callback.on_rollout_end()
    return True

  # --- updates -----------------------------------------------------------------

  def _update_learning_rate(self, optimizers) -> None:
    """Phase-aware: after the base update (adaptive scalar or ctrl schedule),
    re-apply the dstb player's own fixed LR in dstb phases (the base call
    would otherwise apply the ctrl schedule to the dstb optimizer)."""
    super()._update_learning_rate(optimizers)
    if (not self.adaptive_lr and self._dstb_lr is not None
        and self._phase() == "dstb"):
      from stable_baselines3.common.utils import update_learning_rate
      update_learning_rate(
        optimizers if not isinstance(optimizers, list) else optimizers[0],
        self._dstb_lr)
      self.logger.record("game/lr_dstb", self._dstb_lr)

  def train(self) -> None:
    phase = self._phase()
    if phase == "dstb":
      # Min player: SAME targets, NEGATED advantage; PPO.train()
      # runs against the swapped-in dstb policy/buffer with the DSTB player's
      # own LR (adaptive state) and entropy coefficient.
      self.dstb_rollout_buffer.advantages *= -1.0
      ctrl_policy, ctrl_buf = self.policy, self.rollout_buffer
      ctrl_ent = self.ent_coef
      self.policy, self.rollout_buffer = self.dstb_policy, self.dstb_rollout_buffer
      if self._dstb_ent is not None:
        self.ent_coef = self._dstb_ent
      if self.adaptive_lr:
        self._adaptive_lr = self._adaptive_lr_dstb
      try:
        super().train()
      finally:
        if self.adaptive_lr:
          self._adaptive_lr_dstb = float(self._adaptive_lr)
          self.logger.record("game/lr_dstb", self._adaptive_lr_dstb)
        self.policy, self.rollout_buffer = ctrl_policy, ctrl_buf
        self.ent_coef = ctrl_ent
    else:
      if self.adaptive_lr:
        self._adaptive_lr = self._adaptive_lr_ctrl
      super().train()
      if self.adaptive_lr:
        self._adaptive_lr_ctrl = float(self._adaptive_lr)
        self.logger.record("game/lr_ctrl", self._adaptive_lr_ctrl)

    self._rollouts_done += 1
    self.logger.record("game/phase_is_dstb", float(phase == "dstb"))
    self.logger.record("game/rollouts_done", self._rollouts_done)

    # snapshot + prune every N full cycles (after pretrain)
    if (
      self._leaderboard is not None
      and self._rollouts_done > self._dstb_pretrain
      and (self._rollouts_done - self._dstb_pretrain)
      % (self._cycle_len() * self._lb_freq_cycles) == 0
    ):
      self._leaderboard.prune(
        self.num_timesteps, self.policy, self.dstb_policy
      )
      self.logger.record(
        "game/archived_dstb", float(len(self._leaderboard.dstb_steps))
      )

  # --- persistence ---------------------------------------------------------------

  def _excluded_save_params(self):
    return super()._excluded_save_params() + [
      "_leaderboard", "_scratch_dstb", "_loaded_scratch",
    ]

  def _get_torch_save_params(self):
    state_dicts, tensors = super()._get_torch_save_params()
    state_dicts = state_dicts + ["dstb_policy", "dstb_policy.optimizer"]
    return state_dicts, tensors

collect_rollouts

collect_rollouts(
    env, callback, rollout_buffer, n_rollout_steps
)

Two-player rollout: both actions computed each step, env receives the concatenation, only the ACTIVE player's data is stored (in its buffer).

Source code in safety_sb3/ppo_2p.py

def collect_rollouts(
  self,
  env: VecEnv,
  callback: BaseCallback,
  rollout_buffer: RolloutBuffer,
  n_rollout_steps: int,
) -> bool:
  """Two-player rollout: both actions computed each step, env receives the
  concatenation, only the ACTIVE player's data is stored (in its buffer)."""
  # This override does not call the AbstractPPO1P loop, so anneal gamma here;
  # the helper updates both the ctrl and dstb rollout buffers (idempotent).
  self._apply_gamma_anneal()
  if self._tensor_path:
    return self._collect_rollouts_tensor(env, callback, rollout_buffer,
                                         n_rollout_steps)
  del rollout_buffer  # phase decides the buffer
  phase = self._phase()
  active_policy = self.dstb_policy if phase == "dstb" else self.policy
  passive_policy = self.policy if phase == "dstb" else self.dstb_policy
  buf = self.dstb_rollout_buffer if phase == "dstb" else self.rollout_buffer

  assert self._last_obs is not None
  active_policy.set_training_mode(False)
  passive_policy.set_training_mode(False)
  if phase == "ctrl":
    self._resample_slice_opponents()

  n_steps = 0
  buf.reset()
  callback.on_rollout_start()

  while n_steps < n_rollout_steps:
    with th.no_grad():
      obs_tensor = obs_as_tensor(self._last_obs, self.device)
      actions, values, log_probs = active_policy(obs_tensor)
    actions = actions.cpu().numpy()

    if phase == "dstb":
      with th.no_grad():
        ctrl_a, _, _ = passive_policy(obs_tensor)  # frozen ctrl, stochastic
      ctrl_np = ctrl_a.cpu().numpy()
      dstb_np = actions
    else:
      ctrl_np = actions
      dstb_np = self._dstb_actions_for_ctrl_phase(obs_tensor)

    ctrl_clipped = np.clip(ctrl_np, self._ctrl_space.low, self._ctrl_space.high)
    dstb_clipped = np.clip(dstb_np, self._dstb_space.low, self._dstb_space.high)
    full_actions = np.concatenate([ctrl_clipped, dstb_clipped], axis=1)

    new_obs, rewards, dones, infos = env.step(full_actions)
    self.num_timesteps += env.num_envs

    callback.update_locals(locals())
    if not callback.on_step():
      return False
    self._update_info_buffer(infos, dones)
    n_steps += 1

    # Timeout bootstrap (ACTIVE player's value fn) — DISABLED by default:
    # the reward is the physical margin g(s) (see the AbstractPPO docstring).
    if self.bootstrap_on_timeout:
      for idx, done in enumerate(dones):
        if (
          done
          and infos[idx].get("terminal_observation") is not None
          and infos[idx].get("TimeLimit.truncated", False)
        ):
          terminal_obs = active_policy.obs_to_tensor(
            infos[idx]["terminal_observation"]
          )[0]
          with th.no_grad():
            terminal_value = active_policy.predict_values(terminal_obs)[0]
          rewards[idx] += self.gamma * terminal_value

    buf.record_extras(infos)  # the buffer keeps what its operator needs
    buf.add(
      self._last_obs, actions, rewards,
      self._last_episode_starts, values, log_probs,
    )

    # episode flags -> league board (training outcomes)
    self._note_reach(infos)
    self._ever_gneg |= rewards < 0.0
    if self._leaderboard is not None and dones.any():
      timeouts = np.array(
        [bool(info.get("TimeLimit.truncated", False)) for info in infos]
      )
      succ = self._episode_success(dones, timeouts)
      b = self._leaderboard.board
      if phase == "ctrl":
        for si, opp in enumerate(self._slice_opps):
          m = dones & (self._slice_of_env == si)
          if not m.any():
            continue
          col = (b.shape[1] - 1 if opp == _ZERO
                 else b.shape[1] - 2 if opp in (_RANDOM, _CURRENT) else opp)
          self._leaderboard.ema_score(-1, col, float(succ[m].mean()))
      else:
        self._leaderboard.ema_score(
          -1, b.shape[1] - 2, float(succ[dones].mean())
        )
      self._clear_reach(dones)
      self._ever_gneg[dones] = False

    self._last_obs = new_obs
    self._last_episode_starts = dones

  with th.no_grad():
    values = active_policy.predict_values(obs_as_tensor(new_obs, self.device))
  buf.compute_returns_and_advantage(last_values=values, dones=dones)

  callback.update_locals(locals())
  callback.on_rollout_end()
  return True

SafetyPPO2P

Bases: AbstractPPO2P

Two-player on-policy avoid game — ISAACS proper (Hsu, Nguyen, Fisac 2022, eq. 7), with league opponents::

V(s) = (1-γ)·g + γ·max_ctrl min_dstb min(g, V')

the robust-invariance value under a worst-case disturbance. No target set, no l — the paper has neither, so info["l_x"] is never read and the league scores an episode a WIN when the ctrl player survives to the time limit without ever violating g.

This is the class for an adversarial avoid task ("stay standing against a worst-case force"). Do NOT emulate it by giving :class:ReachAvoidPPO2P a degenerate l: no l reduces the reach-avoid operator to avoid (:mod:safety_sb3.backups proves the two conditions are contradictory).

Source code in safety_sb3/ppo_2p.py

class SafetyPPO2P(AbstractPPO2P):
  """Two-player on-policy **avoid** game — ISAACS proper (Hsu, Nguyen, Fisac
  2022, eq. 7), with league opponents::

      V(s) = (1-γ)·g + γ·max_ctrl min_dstb min(g, V')

  the robust-invariance value under a worst-case disturbance. No target set, no
  ``l`` — the paper has neither, so ``info["l_x"]`` is never read and the league
  scores an episode a WIN when the ctrl player survives to the time limit without
  ever violating ``g``.

  This is the class for an adversarial *avoid* task ("stay standing against a
  worst-case force"). Do NOT emulate it by giving :class:`ReachAvoidPPO2P` a
  degenerate ``l``: no ``l`` reduces the reach-avoid operator to avoid
  (:mod:`safety_sb3.backups` proves the two conditions are contradictory).
  """

  _MODE = backups.AVOID

ReachAvoidPPO2P

Bases: _ReachAvoidPlumbing, AbstractPPO2P

Two-player on-policy reach-avoid game — Gameplay Filters (Hsu et al. 2024, eq. 6a), which extends ISAACS to reach-avoid: anchor min(l, g).

Needs a target margin l (info["l_x"] on the numpy path, step_tensor's l_x on the GPU-resident path). The league scores an episode a ctrl WIN only if the controller both survived to the time limit and reached the target. For a two-player avoid game (no target set) use :class:SafetyPPO2P.

Source code in safety_sb3/ppo_2p.py

class ReachAvoidPPO2P(_ReachAvoidPlumbing, AbstractPPO2P):
  """Two-player on-policy **reach-avoid** game — Gameplay Filters (Hsu et al.
  2024, eq. 6a), which extends ISAACS to reach-avoid: anchor ``min(l, g)``.

  Needs a target margin ``l`` (``info["l_x"]`` on the numpy path,
  ``step_tensor``'s ``l_x`` on the GPU-resident path). The league scores an
  episode a ctrl WIN only if the controller both survived to the time limit and
  reached the target. For a two-player *avoid* game (no target set) use
  :class:`SafetyPPO2P`.
  """

  _MODE = backups.REACH_AVOID

Off-policy learners (SAC family)

safety_sb3.sac_base

SAC, written without a fixed Bellman backup — the shared half of the family.

:class:AbstractSAC is everything the SAC-family learners share regardless of how many players are in the game: the constructor and mode check, replay-buffer selection and validation, the entropy-temperature bounds / reset, gamma annealing, the GPU-resident collection loop, and the TD target itself. It deliberately owns no train() — that is the axis its subclasses split on:

:class:`~safety_sb3.sac_1p.AbstractSAC1P` — one actor, one entropy temp.
:class:`~safety_sb3.sac_2p.AbstractSAC2P` — two actors over disjoint action
    sub-spaces, two entropy temps, ONE twin critic over the joint action.

Which value function a learner converges to is a parameter, not a class: the update loops form the soft next-state value V' and hand it to :func:safety_sb3.backups.target together with self._MODE.

``_MODE = backups.AVOID``        -> SafetySAC1P / SafetySAC2P
``_MODE = backups.REACH_AVOID``  -> ReachAvoidSAC1P / ReachAvoidSAC2P
``_MODE = backups.CUMULATIVE``   -> CumulativeSAC1P (ordinary SAC)

Each of those is a one-line specialization; the mode is also selectable per instance via the mode= constructor argument. This follows Bertsekas's abstract-DP framing (define the algorithm against an abstract operator, then choose the operator) and it is why the reach-avoid learner cannot drift from the avoid learner again: before this class they were two hand-copied train() bodies, and the copy had already lost the min_alpha/max_alpha clamp.

Note that SAC needs no reach-avoid mixin, unlike the on-policy family: its whole mode delta is the buffer choice and the l-carrying flag below, both of which already dispatch on _MODE right here.

AbstractSAC

Bases: GammaAnnealMixin, SAC

SAC whose Bellman backup is chosen by _MODE (see the module docstring).

Subclasses SB3's SAC to reuse the actor, entropy regularization, replay buffer, etc. Abstract: it has no train().

GPU-resident path: pass a TensorVecEnv (detected via is_tensor_env, optionally auto-wrapped in TensorVecNormalize with normalize_obs=True) and collection + replay stay on device end-to-end (:class:safety_sb3.tensor_replay.TensorReplayBuffer); train() is unchanged -- the tensor buffer returns the same named tuples. NOTE on the tensor path set gradient_steps explicitly (>=1): the SB3 convention gradient_steps=-1 ("as many as env steps") would mean num_envs updates per vector step.

Parameters:

Name	Type	Description	Default
mode	str | None	backup to converge to, one of :data:safety_sb3.backups.MODES. Defaults to the class's _MODE, so subclasses need only set that.	None
terminal_type	str	terminal target for the reach-avoid backup; ignored by the other modes. See :func:safety_sb3.backups.reach_avoid_target.	'all'

Source code in safety_sb3/sac_base.py

class AbstractSAC(GammaAnnealMixin, SAC):
  """SAC whose Bellman backup is chosen by ``_MODE`` (see the module docstring).

  Subclasses SB3's SAC to reuse the actor, entropy regularization, replay
  buffer, etc. Abstract: it has no ``train()``.

  GPU-resident path: pass a ``TensorVecEnv`` (detected via ``is_tensor_env``,
  optionally auto-wrapped in ``TensorVecNormalize`` with ``normalize_obs=True``)
  and collection + replay stay on device end-to-end
  (:class:`safety_sb3.tensor_replay.TensorReplayBuffer`); ``train()`` is
  unchanged -- the tensor buffer returns the same named tuples. NOTE on the
  tensor path set ``gradient_steps`` explicitly (>=1): the SB3 convention
  ``gradient_steps=-1`` ("as many as env steps") would mean num_envs updates
  per vector step.

  :param mode: backup to converge to, one of :data:`safety_sb3.backups.MODES`.
      Defaults to the class's ``_MODE``, so subclasses need only set that.
  :param terminal_type: terminal target for the reach-avoid backup; ignored by
      the other modes. See :func:`safety_sb3.backups.reach_avoid_target`.
  """

  #: default backup; subclasses set this, ``mode=`` overrides it per instance
  _MODE = backups.AVOID
  _tensor_store_l = False  # reach-avoid flips this (buffer stores l_x)

  def __init__(self, *args, mode: str | None = None,
               replay_buffer_class=None, terminal_type: str = "all",
               normalize_obs: bool = False, gamma_anneal=True,
               min_alpha: float | None = 1e-3, max_alpha: float | None = None,
               **kwargs):
    # The mode must be known before super().__init__ -> _setup_model, which
    # picks the buffer (an l-carrying one for reach-avoid).
    self._MODE = backups.check_mode(self._MODE if mode is None else mode)
    self.terminal_type = backups.check_terminal_type(terminal_type)
    if self._MODE == backups.REACH_AVOID:
      if replay_buffer_class is None:
        replay_buffer_class = ReachAvoidReplayBuffer
      self._tensor_store_l = True  # tensor buffer stores l(s)
    if "env" in kwargs:
      kwargs["env"] = guard_and_normalize_env(kwargs["env"], normalize_obs)
    elif len(args) >= 2:
      args = list(args)
      args[1] = guard_and_normalize_env(args[1], normalize_obs)
    _env = kwargs.get("env", args[1] if len(args) >= 2 else None)
    self._tensor_path = bool(getattr(_env, "is_tensor_env", False))
    # Entropy-temperature (alpha) FLOOR/ceiling (reference: min_alpha=1e-3).
    # SB3's auto-entropy is unbounded; a floor stops alpha collapsing to ~0
    # (deterministic, no exploration) esp. after a gamma-jump alpha reset.
    # Set before super().__init__ -> _setup_model reads them.
    self._min_alpha = None if min_alpha is None else float(min_alpha)
    self._max_alpha = None if max_alpha is None else float(max_alpha)
    super().__init__(*args, replay_buffer_class=replay_buffer_class, **kwargs)
    # Discount-factor annealing (ON by default): the REFERENCE-FAITHFUL
    # discrete-jump schedule (gamma 0.99 -> 0.999 @20% -> 0.9999 @40%, hold);
    # each jump resets alpha via _on_gamma_jump (Q-scale shift). self.gamma is
    # read in the TD target of train(). See gamma_anneal.py.
    self._setup_gamma_anneal(gamma_anneal)
    self._check_replay_buffer()

  def _check_replay_buffer(self) -> None:
    """Reach-avoid needs a buffer that stores ``l(s)``.

    On load (``_init_setup_model=False``) the buffer is built later; only
    validate when it already exists. The tensor path builds its own
    l-carrying buffer.
    """
    if self._MODE != backups.REACH_AVOID or self._tensor_path:
      return
    if self.replay_buffer is not None and not isinstance(
        self.replay_buffer, ReachAvoidReplayBuffer):
      raise TypeError(
        f"{type(self).__name__} is in mode={self._MODE!r} and needs a "
        "ReachAvoidReplayBuffer (it stores l(s)); got "
        f"{type(self.replay_buffer).__name__}."
      )

  def _setup_model(self) -> None:
    if not self._tensor_path:
      super()._setup_model()
      self._setup_entropy_bounds()
      return
    from .tensor_replay import TensorReplayBuffer
    # Keep the numpy buffer SB3 allocates in _setup_model negligible, then
    # replace it with the device-resident buffer at the real size.
    real_size = self.buffer_size
    self.buffer_size = self.env.num_envs
    super()._setup_model()
    self.buffer_size = real_size
    self.replay_buffer = TensorReplayBuffer(
      real_size,
      obs_dim=int(np.prod(self.observation_space.shape)),
      act_dim=int(np.prod(self.action_space.shape)),
      n_envs=self.env.num_envs,
      device=str(self.device),
      store_l=self._tensor_store_l,
    )
    self._setup_entropy_bounds()

  # --- entropy-temperature (alpha) bounds + reset (issues 6 & 7) -----------
  def _setup_entropy_bounds(self) -> None:
    """Compute the log-space alpha clamp bounds and snapshot the init
    log_ent_coef (so a gamma jump can reset alpha to it)."""
    import math
    # A NON-POSITIVE bound means "no bound": log() is undefined there, and a
    # floor of 0 is exactly the request to let alpha decay freely. Treating it
    # as None rather than raising matters because `--min-alpha 0` is the
    # natural way to ASK for an unfloored run, and a math domain error at
    # construction kills the run before it starts (E056 lost two cells this
    # way, leaving the min_alpha hypothesis untested).
    self._log_min_alpha = (None if not self._min_alpha or self._min_alpha <= 0
                           else math.log(self._min_alpha))
    self._log_max_alpha = (None if not self._max_alpha or self._max_alpha <= 0
                           else math.log(self._max_alpha))
    lec = getattr(self, "log_ent_coef", None)   # None for a FIXED ent_coef
    self._init_log_ent_coef = None if lec is None else lec.detach().clone()

  @staticmethod
  def _polyak(params, target_params, tau: float) -> None:
    """Soft target update as ONE fused foreach op instead of a Python loop.

    SB3's ``polyak_update`` iterates in Python and issues two kernels per tensor
    (``mul_`` then ``add``), so a 6-layer twin critic costs ~56 launches every
    target update. The foreach ops below do the identical arithmetic in two
    multi-tensor launches regardless of parameter count.

    BITWISE-IDENTICAL, deliberately. ``_foreach_lerp_`` would be one launch
    instead of two and is algebraically the same
    (``t + tau*(p-t)`` == ``t*(1-tau) + p*tau``), but it reorders the float32
    arithmetic -- measured drift 4.8e-07 absolute / 6.6e-05 relative per update
    at tau=0.01. Keeping SB3's mul-then-add order costs one extra launch and
    buys exact equality with every checkpoint we have already validated, which
    matters on the certificate path where the whole artifact is a learned
    level set.

    This is worth doing because the update is launch-bound, not compute-bound
    (measured: t_grad flat from batch 512 to 16384), and because a Python loop
    over parameters is hostile to CUDA-graph capture.
    """
    tl = list(target_params)
    if not tl:                     # e.g. batch_norm_stats on a net without BN
      return
    with th.no_grad():
      th._foreach_mul_(tl, 1.0 - tau)
      th._foreach_add_(tl, list(params), alpha=tau)

  # --- logging over the DUMP WINDOW, not the last call ----------------------
  # On the tensor path _dump_logs() fires every 50k env-steps, which at 1024
  # envs is ~49 train() calls. ``Logger.record`` OVERWRITES ("if called many
  # times, last value will be used"), so a per-train() record() reaches wandb as
  # a 1-in-49 sample and ~96% of the window is discarded unseen. That is how
  # E057 could go critic_loss 0.0014 -> 1.95e+07 inside a single window with
  # only the endpoints visible.
  #
  # ``record_mean`` fixes the average. It does NOT fix an onset spike, which a
  # 49-call mean dilutes -- hence the running MAX below, reset on dump.

  def _record_window_max(self, key: str, value: float) -> None:
    """Running max of ``value`` since the last dump, recorded under ``key``."""
    store = getattr(self, "_win_max_store", None)
    if store is None:
      store = self._win_max_store = {}
    cur = store.get(key)
    store[key] = value if cur is None else max(cur, value)
    self.logger.record(key, store[key])

  def _reset_window_max(self) -> None:
    """Start a fresh window. Called immediately after a dump."""
    self._win_max_store = {}

  def _clamp_entropy_temps(self) -> None:
    """Clamp the learned entropy temperature into [min_alpha, max_alpha].
    Called after each entropy-coefficient optimizer step in train()."""
    if self._log_min_alpha is None and self._log_max_alpha is None:
      return
    lec = getattr(self, "log_ent_coef", None)
    if lec is not None:
      with th.no_grad():
        lec.clamp_(min=self._log_min_alpha, max=self._log_max_alpha)

  def _entropy_optimizer_lr(self) -> float:
    """LR for the (ctrl) entropy-coefficient optimizer. Falls back to the
    shared ``learning_rate`` when no dedicated ``ent_coef_lr`` was set
    (the two-player learners set ``self._ent_coef_lr``); keeps a rebuilt
    optimizer on the configured entropy lr rather than the shared one."""
    lr = getattr(self, "_ent_coef_lr", None)
    return float(self.lr_schedule(1)) if lr is None else float(lr)

  def _reset_entropy_temp(self) -> None:
    """Reset the learned entropy temperature to its init value and rebuild
    its optimizer (clears Adam moments) -- the reference alpha reset on a
    gamma jump. No-op for a fixed ent_coef."""
    lec = getattr(self, "log_ent_coef", None)
    if lec is None or self._init_log_ent_coef is None:
      return
    with th.no_grad():
      lec.data.copy_(self._init_log_ent_coef)
    if getattr(self, "ent_coef_optimizer", None) is not None:
      self.ent_coef_optimizer = th.optim.Adam(
        [lec], lr=self._entropy_optimizer_lr())

  def _on_gamma_jump(self, old_gamma: float, new_gamma: float) -> None:
    """A discrete gamma jump shifts the Q-scale, so the tuned entropy
    temperature is stale -> reset it (reference: reset alpha on every jump)."""
    self._reset_entropy_temp()
    logger = getattr(self, "logger", None)
    if logger is not None:
      logger.record("train/alpha_reset_gamma", float(new_gamma))

  # --- tensor collection ---------------------------------------------------

  def collect_rollouts(self, env, callback, train_freq: TrainFreq,
                       replay_buffer, action_noise=None, learning_starts=0,
                       log_interval=None) -> RolloutReturn:
    # Anneal self.gamma before the train() that follows this collection reads
    # it in the TD target. The tensor path bypasses SB3's
    # _update_current_progress_remaining, so apply it here (idempotent).
    self._apply_gamma_anneal()
    if self._tensor_path:
      return self._collect_rollouts_tensor(
        env, callback, train_freq, replay_buffer,
        learning_starts=learning_starts, log_interval=log_interval)
    return super().collect_rollouts(
      env, callback, train_freq, replay_buffer, action_noise=action_noise,
      learning_starts=learning_starts, log_interval=log_interval)

  def _tensor_policy_actions(self, obs: "th.Tensor") -> "th.Tensor":
    """Post-warmup policy actions for the tensor collect: on device, in the
    env's action range, shape ``(num_envs, action_dim)``.

    Single-player: the ctrl actor. The two-player learners override this to
    sample BOTH players and concatenate ``[a_ctrl, a_dstb]`` so the composed
    action matches the env's ``ctrl_dim + dstb_dim`` action space (the
    numpy-path analog is :meth:`~safety_sb3.sac_2p.AbstractSAC2P._sample_action`).
    """
    with th.no_grad():
      return self.actor(obs, deterministic=False)

  def _collect_rollouts_tensor(self, env, callback, train_freq: TrainFreq,
                               replay_buffer, learning_starts=0,
                               log_interval=None) -> RolloutReturn:
    """Torch-native off-policy collection: actor forward, env.step_tensor,
    replay add — all on device. Mirrors OffPolicyAlgorithm.collect_rollouts
    step/episode accounting (one vector step counts 1 toward train_freq)."""
    self.policy.set_training_mode(False)
    callback.on_rollout_start()
    dev = env.device

    obs = self._last_obs
    if not th.is_tensor(obs):  # first call after _setup_learn
      obs = th.as_tensor(np.asarray(obs), dtype=th.float32, device=dev)

    low = th.as_tensor(self.action_space.low, dtype=th.float32, device=dev)
    high = th.as_tensor(self.action_space.high, dtype=th.float32, device=dev)

    if getattr(self, "_t_ep_ret", None) is None \
            or self._t_ep_ret.shape[0] != env.num_envs:
      self._t_ep_ret = th.zeros(env.num_envs, device=dev)
      self._t_ep_len = th.zeros(env.num_envs, device=dev)
      self._tensor_last_dump = 0
    fin_ret, fin_len = [], []

    num_collected_steps, num_collected_episodes = 0, 0
    continue_training = True
    while should_collect_more_steps(train_freq, num_collected_steps,
                                    num_collected_episodes):
      if self.num_timesteps < learning_starts:
        actions = low + (high - low) * th.rand(
          (env.num_envs, low.shape[0]), device=dev)
      else:
        actions = self._tensor_policy_actions(obs)
      actions = th.clamp(actions, low, high)

      new_obs, g, dones, timeouts, l_x = env.step_tensor(actions)
      # An env may EXECUTE something other than what it was handed: a safety
      # filter replacing an uncertified proposal with its fallback's action is
      # the case this exists for. The buffer must hold the transition that
      # ACTUALLY happened -- storing the proposal against the resulting
      # (reward, next state) would fit the critic to a transition that never
      # occurred. Off-policy learning makes the substitution exactly correct
      # with no importance correction, which is why filtered training is run
      # with SAC. An env that does not override leaves `actions` untouched.
      executed = getattr(env, "executed_action", None)
      if executed is not None:
        actions = executed
      self.num_timesteps += env.num_envs
      num_collected_steps += 1

      callback.update_locals(locals())
      if not callback.on_step():
        continue_training = False
        break

      replay_buffer.add_batch(
        obs, new_obs, actions, g, dones, timeouts,
        l_x=l_x if self._tensor_store_l else None)

      self._t_ep_ret += g
      self._t_ep_len += 1.0
      n_done = int(dones.sum())
      if n_done:
        d = dones.bool()
        fin_ret.append(self._t_ep_ret[d])
        fin_len.append(self._t_ep_len[d])
        self._t_ep_ret = th.where(d, th.zeros_like(self._t_ep_ret),
                                  self._t_ep_ret)
        self._t_ep_len = th.where(d, th.zeros_like(self._t_ep_len),
                                  self._t_ep_len)
        num_collected_episodes += n_done
        self._episode_num += n_done
      obs = new_obs

    self._last_obs = obs
    if fin_ret:
      self.logger.record("rollout/ep_rew_mean", float(th.cat(fin_ret).mean()))
      self.logger.record("rollout/ep_len_mean", float(th.cat(fin_len).mean()))
    # periodic dump (numpy path dumps on episode boundaries; with thousands
    # of parallel envs use an env-step cadence instead)
    if self.num_timesteps - self._tensor_last_dump >= 50_000:
      self._tensor_last_dump = self.num_timesteps
      for k, v in (env.metrics() or {}).items():
        self.logger.record(f"env/{k}", float(v))
      self._dump_logs()
      self._reset_window_max()   # the window just closed; start the next one

    callback.on_rollout_end()
    return RolloutReturn(num_collected_steps * env.num_envs,
                         num_collected_episodes, continue_training)

  def train(self, gradient_steps: int, batch_size: int) -> None:
    """Refuse to fall through to SB3's ``SAC.train``.

    This class has no update loop by design — that is the Players axis. Without
    this guard, instantiating it directly would inherit stock SAC's train() and
    silently converge to the CUMULATIVE fixed point while ``_MODE`` claimed
    otherwise: exactly the "different fixed points under one name" bug the
    library exists to prevent.
    """
    raise NotImplementedError(
      f"{type(self).__name__} has no update loop. Use a concrete learner "
      "(SafetySAC1P, ReachAvoidSAC2P, ...), or AbstractSAC1P / AbstractSAC2P "
      "with mode= if you want the loop without a named mode.")

  def _bellman_target(self, replay_data, v_next: "th.Tensor") -> "th.Tensor":
    """The backup for THIS instance's mode -- see :mod:`safety_sb3.backups`.

    ``replay_data.rewards`` carries the safety margin ``g(s)`` in the safety
    modes and the reward ``r`` in ``CUMULATIVE``; ``l_x`` is present only on an
    l-carrying buffer (reach-avoid) and is ignored by the other modes.
    """
    return backups.target(
      self._MODE, replay_data.rewards, v_next, 1.0 - replay_data.dones,
      self.gamma, l=getattr(replay_data, "l_x", None),
      terminal_type=self.terminal_type)

train

train(gradient_steps, batch_size)

Refuse to fall through to SB3's SAC.train.

This class has no update loop by design — that is the Players axis. Without this guard, instantiating it directly would inherit stock SAC's train() and silently converge to the CUMULATIVE fixed point while _MODE claimed otherwise: exactly the "different fixed points under one name" bug the library exists to prevent.

Source code in safety_sb3/sac_base.py

def train(self, gradient_steps: int, batch_size: int) -> None:
  """Refuse to fall through to SB3's ``SAC.train``.

  This class has no update loop by design — that is the Players axis. Without
  this guard, instantiating it directly would inherit stock SAC's train() and
  silently converge to the CUMULATIVE fixed point while ``_MODE`` claimed
  otherwise: exactly the "different fixed points under one name" bug the
  library exists to prevent.
  """
  raise NotImplementedError(
    f"{type(self).__name__} has no update loop. Use a concrete learner "
    "(SafetySAC1P, ReachAvoidSAC2P, ...), or AbstractSAC1P / AbstractSAC2P "
    "with mode= if you want the loop without a named mode.")

safety_sb3.sac_1p

Single-player SAC — the 1P half of the SAC family, and its three modes.

AbstractSAC                 (sac_base.py — ctor, buffers, entropy, backup)
└─ AbstractSAC1P            the single-actor update loop
    ├─ SafetySAC1P          _MODE = AVOID
    ├─ ReachAvoidSAC1P      _MODE = REACH_AVOID
    └─ CumulativeSAC1P      _MODE = CUMULATIVE

There is exactly ONE train() here and the three learners below are one line each. No reach-avoid mixin is needed on this side of the library: SAC's entire mode delta is the replay-buffer choice and the TD target, and both already dispatch on _MODE inside :class:~safety_sb3.sac_base.AbstractSAC.

That de-duplication is load-bearing, not cosmetic. These used to be two hand-copied train() bodies, and the copy had silently dropped _clamp_entropy_temps() — so min_alpha/max_alpha did nothing on the reach-avoid path. One shared body makes that unrepresentable.

AbstractSAC1P

Bases: AbstractSAC

SAC's single-actor update loop; the backup is chosen by _MODE.

Source code in safety_sb3/sac_1p.py

class AbstractSAC1P(AbstractSAC):
  """SAC's single-actor update loop; the backup is chosen by ``_MODE``."""

  def train(self, gradient_steps: int, batch_size: int) -> None:
    """Largely follows the original SAC train method from stable_baselines3.
    The TD target is the one for ``self._MODE`` (:meth:`_bellman_target`), so
    this one body serves every single-player mode in the library.
    """
    # Switch to train mode (this affects batch norm / dropout)
    self.policy.set_training_mode(True)
    # Update optimizers learning rate
    optimizers = [self.actor.optimizer, self.critic.optimizer]
    if self.ent_coef_optimizer is not None:
      optimizers += [self.ent_coef_optimizer]

    # Update learning rate according to lr schedule
    self._update_learning_rate(optimizers)

    # Logging scalars accumulate ON DEVICE; one sync per train() call, not four
    # per gradient step. Each `.item()` drains the CUDA queue and blocks the CPU
    # while Python re-queues the next update's ~300 kernel launches -- measured:
    # t_grad is FLAT from batch 512 to 16384 (11.47 -> 14.20 ms), the textbook
    # launch-bound signature, and at gradient_steps=64 the loop was paying 256
    # stalls per collect cycle. It is also a HARD PREREQUISITE for CUDA-graph
    # capture: `.item()` is a capture error.
    _z = th.zeros((), device=self.device)
    _ninf = th.full((), -float("inf"), device=self.device)
    ent_coef_loss_sum, ent_coef_sum = _z.clone(), _z.clone()
    actor_loss_sum, critic_loss_sum = _z.clone(), _z.clone()
    critic_loss_max, ent_coef_max = _ninf.clone(), _ninf.clone()
    n_ent_coef_loss = 0

    for gradient_step in range(gradient_steps):
      # Sample replay buffer
      replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env)

      # We need to sample because `log_std` may have changed between two gradient steps
      if self.use_sde:
        self.actor.reset_noise()

      # Action by the current actor for the sampled state
      actions_pi, log_prob = self.actor.action_log_prob(replay_data.observations)
      log_prob = log_prob.reshape(-1, 1)

      ent_coef_loss = None
      if self.ent_coef_optimizer is not None and self.log_ent_coef is not None:
        # Important: detach the variable from the graph
        # so we don't change it with other losses
        # see https://github.com/rail-berkeley/softlearning/issues/60
        ent_coef = th.exp(self.log_ent_coef.detach())
        ent_coef_loss = -(self.log_ent_coef *
                          (log_prob + self.target_entropy).detach()).mean()
        ent_coef_loss_sum += ent_coef_loss.detach()
        n_ent_coef_loss += 1
      else:
        ent_coef = self.ent_coef_tensor

      # ent_coef is shape [1] (log_ent_coef is a 1-element parameter); reduce
      # it so the accumulators stay 0-d.
      _ec = ent_coef.detach().mean()
      ent_coef_sum += _ec
      ent_coef_max = th.maximum(ent_coef_max, _ec)

      # Optimize entropy coefficient
      if ent_coef_loss is not None and self.ent_coef_optimizer is not None:
        self.ent_coef_optimizer.zero_grad()
        ent_coef_loss.backward()
        self.ent_coef_optimizer.step()
        self._clamp_entropy_temps()  # min_alpha/max_alpha floor/ceiling

      with th.no_grad():
        # Select action according to policy
        next_actions, next_log_prob = self.actor.action_log_prob(
          replay_data.next_observations
        )
        # Compute the next Q values: min over all critics targets
        next_q_values = th.cat(
          self.critic_target(replay_data.next_observations, next_actions), dim=1
        )
        next_q_values, _ = th.min(next_q_values, dim=1, keepdim=True)
        # SOFT next-state value -- ONLY for the cumulative backup.
        #
        # For the AVOID / REACH_AVOID (Hamilton-Jacobi) backups the critic target
        # is the PURE Isaacs/HJ backup: ISAACS (Hsu et al. 2023) eq. 8a is
        #     y = (1-g) g' + g * min{g', Q(x',u',d')}
        # with NO -alpha*log pi in the bootstrap; entropy is confined to the
        # actor loss (eq. 8b), where it is an annealed exploration regulariser.
        # Both reference codebases gate identically -- base_block.py adds the
        # entropy term only `if self.mode == 'performance'` (== our CUMULATIVE),
        # and ISAACS passes entropy_motives=0. A soft term inside the min/max
        # taxes the value whose ZERO LEVEL SET is the safety certificate,
        # warping {V>=0} (optimistically at the boundary) and voiding the RSS'21
        # under-approximation guarantee. This library previously subtracted it
        # for ALL modes -- a port error, contradicting this file's own
        # "matches safe_adaptation_dev" docstring.
        if self._MODE == backups.CUMULATIVE:
          next_q_values = next_q_values - ent_coef * next_log_prob.reshape(-1, 1)

        target_q_values = self._bellman_target(replay_data, next_q_values)

      # Get current Q-values estimates for each critic network
      # using action from the replay buffer
      current_q_values = self.critic(replay_data.observations, replay_data.actions)

      # Compute critic loss
      critic_loss = 0.5 * sum(
        F.mse_loss(current_q, target_q_values) for current_q in current_q_values
      )
      assert isinstance(critic_loss, th.Tensor)  # for type checker
      _cl = critic_loss.detach()  # type: ignore[union-attr]
      critic_loss_sum += _cl
      critic_loss_max = th.maximum(critic_loss_max, _cl)

      # Optimize the critic
      self.critic.optimizer.zero_grad()
      critic_loss.backward()
      self.critic.optimizer.step()

      # Compute actor loss -- the actor MAXIMIZES the value of self._MODE.
      # Alternative: actor_loss = th.mean(log_prob - qf1_pi)
      # Min over all critic networks
      q_values_pi = th.cat(self.critic(replay_data.observations, actions_pi), dim=1)
      min_qf_pi, _ = th.min(q_values_pi, dim=1, keepdim=True)
      actor_loss = (ent_coef*log_prob - min_qf_pi).mean()
      actor_loss_sum += actor_loss.detach()

      # Optimize the actor
      self.actor.optimizer.zero_grad()
      actor_loss.backward()
      self.actor.optimizer.step()

      # Update target networks
      if gradient_step % self.target_update_interval == 0:
        self._polyak(self.critic.parameters(), self.critic_target.parameters(), self.tau)
        # Copy running stats, see GH issue #996 (tau=1.0 == a copy)
        self._polyak(self.batch_norm_stats, self.batch_norm_stats_target, 1.0)

    self._n_updates += gradient_steps

    self.logger.record("train/n_updates", self._n_updates, exclude="tensorboard")
    # record_mean, not record -- see AbstractSAC._record_window_max: the dump
    # cadence is ~49 train() calls, and record() keeps only the last of them.
    n = max(gradient_steps, 1)
    self.logger.record_mean("train/ent_coef", (ent_coef_sum / n).item())
    self.logger.record_mean("train/actor_loss", (actor_loss_sum / n).item())
    self.logger.record_mean("train/critic_loss", (critic_loss_sum / n).item())
    if n_ent_coef_loss:
      self.logger.record_mean("train/ent_coef_loss",
                              (ent_coef_loss_sum / n_ent_coef_loss).item())
    # Onset detectors: a mean over the window hides a spike that starts inside
    # it. These two are the pair that moved first in the E057 divergence.
    self._record_window_max("train/critic_loss_max", critic_loss_max.item())
    self._record_window_max("train/ent_coef_max", ent_coef_max.item())

train

train(gradient_steps, batch_size)

Largely follows the original SAC train method from stable_baselines3. The TD target is the one for self._MODE (:meth:_bellman_target), so this one body serves every single-player mode in the library.

Source code in safety_sb3/sac_1p.py

def train(self, gradient_steps: int, batch_size: int) -> None:
  """Largely follows the original SAC train method from stable_baselines3.
  The TD target is the one for ``self._MODE`` (:meth:`_bellman_target`), so
  this one body serves every single-player mode in the library.
  """
  # Switch to train mode (this affects batch norm / dropout)
  self.policy.set_training_mode(True)
  # Update optimizers learning rate
  optimizers = [self.actor.optimizer, self.critic.optimizer]
  if self.ent_coef_optimizer is not None:
    optimizers += [self.ent_coef_optimizer]

  # Update learning rate according to lr schedule
  self._update_learning_rate(optimizers)

  # Logging scalars accumulate ON DEVICE; one sync per train() call, not four
  # per gradient step. Each `.item()` drains the CUDA queue and blocks the CPU
  # while Python re-queues the next update's ~300 kernel launches -- measured:
  # t_grad is FLAT from batch 512 to 16384 (11.47 -> 14.20 ms), the textbook
  # launch-bound signature, and at gradient_steps=64 the loop was paying 256
  # stalls per collect cycle. It is also a HARD PREREQUISITE for CUDA-graph
  # capture: `.item()` is a capture error.
  _z = th.zeros((), device=self.device)
  _ninf = th.full((), -float("inf"), device=self.device)
  ent_coef_loss_sum, ent_coef_sum = _z.clone(), _z.clone()
  actor_loss_sum, critic_loss_sum = _z.clone(), _z.clone()
  critic_loss_max, ent_coef_max = _ninf.clone(), _ninf.clone()
  n_ent_coef_loss = 0

  for gradient_step in range(gradient_steps):
    # Sample replay buffer
    replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env)

    # We need to sample because `log_std` may have changed between two gradient steps
    if self.use_sde:
      self.actor.reset_noise()

    # Action by the current actor for the sampled state
    actions_pi, log_prob = self.actor.action_log_prob(replay_data.observations)
    log_prob = log_prob.reshape(-1, 1)

    ent_coef_loss = None
    if self.ent_coef_optimizer is not None and self.log_ent_coef is not None:
      # Important: detach the variable from the graph
      # so we don't change it with other losses
      # see https://github.com/rail-berkeley/softlearning/issues/60
      ent_coef = th.exp(self.log_ent_coef.detach())
      ent_coef_loss = -(self.log_ent_coef *
                        (log_prob + self.target_entropy).detach()).mean()
      ent_coef_loss_sum += ent_coef_loss.detach()
      n_ent_coef_loss += 1
    else:
      ent_coef = self.ent_coef_tensor

    # ent_coef is shape [1] (log_ent_coef is a 1-element parameter); reduce
    # it so the accumulators stay 0-d.
    _ec = ent_coef.detach().mean()
    ent_coef_sum += _ec
    ent_coef_max = th.maximum(ent_coef_max, _ec)

    # Optimize entropy coefficient
    if ent_coef_loss is not None and self.ent_coef_optimizer is not None:
      self.ent_coef_optimizer.zero_grad()
      ent_coef_loss.backward()
      self.ent_coef_optimizer.step()
      self._clamp_entropy_temps()  # min_alpha/max_alpha floor/ceiling

    with th.no_grad():
      # Select action according to policy
      next_actions, next_log_prob = self.actor.action_log_prob(
        replay_data.next_observations
      )
      # Compute the next Q values: min over all critics targets
      next_q_values = th.cat(
        self.critic_target(replay_data.next_observations, next_actions), dim=1
      )
      next_q_values, _ = th.min(next_q_values, dim=1, keepdim=True)
      # SOFT next-state value -- ONLY for the cumulative backup.
      #
      # For the AVOID / REACH_AVOID (Hamilton-Jacobi) backups the critic target
      # is the PURE Isaacs/HJ backup: ISAACS (Hsu et al. 2023) eq. 8a is
      #     y = (1-g) g' + g * min{g', Q(x',u',d')}
      # with NO -alpha*log pi in the bootstrap; entropy is confined to the
      # actor loss (eq. 8b), where it is an annealed exploration regulariser.
      # Both reference codebases gate identically -- base_block.py adds the
      # entropy term only `if self.mode == 'performance'` (== our CUMULATIVE),
      # and ISAACS passes entropy_motives=0. A soft term inside the min/max
      # taxes the value whose ZERO LEVEL SET is the safety certificate,
      # warping {V>=0} (optimistically at the boundary) and voiding the RSS'21
      # under-approximation guarantee. This library previously subtracted it
      # for ALL modes -- a port error, contradicting this file's own
      # "matches safe_adaptation_dev" docstring.
      if self._MODE == backups.CUMULATIVE:
        next_q_values = next_q_values - ent_coef * next_log_prob.reshape(-1, 1)

      target_q_values = self._bellman_target(replay_data, next_q_values)

    # Get current Q-values estimates for each critic network
    # using action from the replay buffer
    current_q_values = self.critic(replay_data.observations, replay_data.actions)

    # Compute critic loss
    critic_loss = 0.5 * sum(
      F.mse_loss(current_q, target_q_values) for current_q in current_q_values
    )
    assert isinstance(critic_loss, th.Tensor)  # for type checker
    _cl = critic_loss.detach()  # type: ignore[union-attr]
    critic_loss_sum += _cl
    critic_loss_max = th.maximum(critic_loss_max, _cl)

    # Optimize the critic
    self.critic.optimizer.zero_grad()
    critic_loss.backward()
    self.critic.optimizer.step()

    # Compute actor loss -- the actor MAXIMIZES the value of self._MODE.
    # Alternative: actor_loss = th.mean(log_prob - qf1_pi)
    # Min over all critic networks
    q_values_pi = th.cat(self.critic(replay_data.observations, actions_pi), dim=1)
    min_qf_pi, _ = th.min(q_values_pi, dim=1, keepdim=True)
    actor_loss = (ent_coef*log_prob - min_qf_pi).mean()
    actor_loss_sum += actor_loss.detach()

    # Optimize the actor
    self.actor.optimizer.zero_grad()
    actor_loss.backward()
    self.actor.optimizer.step()

    # Update target networks
    if gradient_step % self.target_update_interval == 0:
      self._polyak(self.critic.parameters(), self.critic_target.parameters(), self.tau)
      # Copy running stats, see GH issue #996 (tau=1.0 == a copy)
      self._polyak(self.batch_norm_stats, self.batch_norm_stats_target, 1.0)

  self._n_updates += gradient_steps

  self.logger.record("train/n_updates", self._n_updates, exclude="tensorboard")
  # record_mean, not record -- see AbstractSAC._record_window_max: the dump
  # cadence is ~49 train() calls, and record() keeps only the last of them.
  n = max(gradient_steps, 1)
  self.logger.record_mean("train/ent_coef", (ent_coef_sum / n).item())
  self.logger.record_mean("train/actor_loss", (actor_loss_sum / n).item())
  self.logger.record_mean("train/critic_loss", (critic_loss_sum / n).item())
  if n_ent_coef_loss:
    self.logger.record_mean("train/ent_coef_loss",
                            (ent_coef_loss_sum / n_ent_coef_loss).item())
  # Onset detectors: a mean over the window hides a spike that starts inside
  # it. These two are the pair that moved first in the E057 divergence.
  self._record_window_max("train/critic_loss_max", critic_loss_max.item())
  self._record_window_max("train/ent_coef_max", ent_coef_max.item())

SafetySAC1P

Bases: AbstractSAC1P

SAC with the avoid (safety) Bellman backup — Fisac et al. 2019::

V(s) = (1-γ)·g + γ·min(g, V')        terminal: V(s) = g

g(s) is the safety margin (it rides on the reward field); V' is SAC's soft next-state value. V(s) >= 0 iff the state is in the safe set. For a task with a target to reach use :class:ReachAvoidSAC1P; the two operators are not interchangeable (see :mod:safety_sb3.backups).

Source code in safety_sb3/sac_1p.py

class SafetySAC1P(AbstractSAC1P):
  """SAC with the **avoid** (safety) Bellman backup — Fisac et al. 2019::

      V(s) = (1-γ)·g + γ·min(g, V')        terminal: V(s) = g

  ``g(s)`` is the safety margin (it rides on the ``reward`` field); ``V'`` is
  SAC's soft next-state value. ``V(s) >= 0`` iff the state is in the safe set.
  For a task with a *target* to reach use :class:`ReachAvoidSAC1P`; the two
  operators are not interchangeable (see :mod:`safety_sb3.backups`).
  """

  _MODE = backups.AVOID
  _tensor_store_l = False  # no l(s) in an avoid problem

ReachAvoidSAC1P

Bases: AbstractSAC1P

SAC with the reach-avoid Bellman backup — Hsu et al. RSS'21::

non-terminal:  y = (1-γ)·min(l, g) + γ·min(g, max(l, V'))
terminal:      y = min(l, g)

g(s) is the safety/avoid margin (rides on the reward field) and l(s) is the target/reach margin, supplied by the env via info["l_x"] and stored by :class:~safety_sb3.buffers_replay.ReachAvoidReplayBuffer — which the base defaults to and validates off this _MODE. Matches the discounted reach-avoid update in safe_adaptation_dev (utils/train.get_bellman_update, mode='reach-avoid', terminal_type='all').

Tensor path (TensorVecEnv): the device-resident :class:~safety_sb3.tensor_replay.TensorReplayBuffer is built with store_l=True (l comes from step_tensor's l_x, not infos).

Source code in safety_sb3/sac_1p.py

class ReachAvoidSAC1P(AbstractSAC1P):
  """SAC with the **reach-avoid** Bellman backup — Hsu et al. RSS'21::

      non-terminal:  y = (1-γ)·min(l, g) + γ·min(g, max(l, V'))
      terminal:      y = min(l, g)

  ``g(s)`` is the safety/avoid margin (rides on the ``reward`` field) and
  ``l(s)`` is the target/reach margin, supplied by the env via ``info["l_x"]``
  and stored by
  :class:`~safety_sb3.buffers_replay.ReachAvoidReplayBuffer` — which the base
  defaults to and validates off this ``_MODE``. Matches the discounted
  reach-avoid update in ``safe_adaptation_dev``
  (``utils/train.get_bellman_update``, ``mode='reach-avoid'``,
  ``terminal_type='all'``).

  Tensor path (``TensorVecEnv``): the device-resident
  :class:`~safety_sb3.tensor_replay.TensorReplayBuffer` is built with
  ``store_l=True`` (``l`` comes from ``step_tensor``'s ``l_x``, not infos).
  """

  _MODE = backups.REACH_AVOID
  _tensor_store_l = True  # tensor buffer stores l(s)

CumulativeSAC1P

Bases: AbstractSAC1P

Ordinary reward-maximizing SAC — y = r + γ·V'.

Not a safety learner: the reward field carries a reward, not a margin, and V >= 0 means nothing. It exists so a nominal baseline runs through every line of the same actor/critic/entropy code as the safety learners — the standard control for "is the safety operator doing the work, or is it just SAC?".

Source code in safety_sb3/sac_1p.py

class CumulativeSAC1P(AbstractSAC1P):
  """Ordinary reward-maximizing SAC — ``y = r + γ·V'``.

  Not a safety learner: the ``reward`` field carries a reward, not a margin, and
  ``V >= 0`` means nothing. It exists so a *nominal* baseline runs through every
  line of the same actor/critic/entropy code as the safety learners — the
  standard control for "is the safety operator doing the work, or is it just
  SAC?".
  """

  _MODE = backups.CUMULATIVE
  _tensor_store_l = False

safety_sb3.sac_2p

Two-player (adversarial) SAC — the 2P half of the SAC family.

AbstractSAC                 (sac_base.py)
└─ AbstractSAC2P            two actors, two entropy temps, ONE joint critic
    ├─ SafetySAC2P          _MODE = AVOID          (ISAACS proper)
    └─ ReachAvoidSAC2P      _MODE = REACH_AVOID    (Gameplay Filters)

Control actor (max-player) and disturbance actor (min-player) share one twin critic over the full concatenated action Q(s, [a_ctrl, a_dstb]). Each actor has its own entropy coefficient / target entropy (ctrl: -ctrl_dim, dstb: -dstb_dim). Uses :class:~safety_sb3.policies.TwoPlayerSACPolicy (two sub-space actors + one full critic).

PURE Hamilton-Jacobi value in the critic target — NO entropy::

V'  = min(Q1', Q2')                          (no soft term)
y   = backups.target(mode, g, V', ...)       (avoid | reach-avoid)

The per-actor temperatures α_ctrl, α_dstb appear ONLY in the ctrl/dstb actor losses (exploration; annealed away over training), never in the critic target. This is ISAACS eq. 8a exactly (Hsu et al. 2023, entropy_motives=0), and it is what both reference codebases do (base_block.py adds entropy to the target only if self.mode == 'performance'). A soft term in the target would tax the game value whose zero level set is the online safety certificate, warping {V>=0}. (An earlier version of this file put - α_ctrl·logπ_ctrl + α_dstb·logπ_dstb in V' and mislabelled it "as ISAACS formulates it"; that was a port error, now corrected — see the note at the target computation.)

This is the minimax game as ISAACS formulates it, and it is NOT the same object as :mod:safety_sb3.ppo_2p even though both hold two actors. Because SAC's critic takes the action, this one Q is the game value and both actors differentiate through it; nothing needs scheduling, and replay makes the learner indifferent to whose policy produced the data. The on-policy two-player learner has two state-only V(s) nets, two rollout buffers and a phase machine, and is an alternating best-response approximation. See :mod:safety_sb3.ppo_2p for the full comparison and for why there is deliberately no shared AbstractTwoPlayer base.

AbstractSAC2P

Bases: AbstractSAC

Two-player off-policy game; the backup is chosen by _MODE.

Source code in safety_sb3/sac_2p.py

class AbstractSAC2P(AbstractSAC):
  """Two-player off-policy game; the backup is chosen by ``_MODE``."""

  policy_aliases = {"MlpPolicy": TwoPlayerSACPolicy,
                    "MultiInputPolicy": TwoPlayerSACPolicy}

  def __init__(
    self,
    policy,
    env,
    *,
    ctrl_action_dim: int | None = None,  # None only on load (restored from policy_kwargs)
    ctrl_update_period: int = 1,
    dstb_update_period: int = 1,
    # --- per-network / per-actor learning rates ---
    # Each is None -> falls back to the shared ``learning_rate``.
    # ``dstb_learning_rate`` is the dstb ACTOR lr; ``ent_coef_lr`` /
    # ``dstb_ent_coef_lr`` are the ctrl / dstb ENTROPY (alpha) optimizer lrs.
    # The ctrl actor always uses the shared ``learning_rate``.
    critic_learning_rate: float | None = None,
    dstb_learning_rate: float | None = None,
    ent_coef_lr: float | None = None,
    dstb_ent_coef_lr: float | None = None,
    # --- optional StepLR decay for the ctrl/dstb/critic optimizers ---
    # OFF by default (constant lr). When on, each lr decays by ``lr_decay`` every
    # ``lr_period`` env-steps toward ``lr_end`` (reference StepLR). The entropy
    # (alpha) lrs stay constant -- an alpha-lr schedule is a follow-up.
    lr_schedule: bool = False,
    lr_period: int = 1_000_000,
    lr_decay: float = 0.1,
    lr_end: float = 0.0,
    # --- league ---
    use_leaderboard: bool = False,
    leaderboard_eval_env=None,
    save_top_k_ctrl: int = 5,
    save_top_k_dstb: int = 5,
    softmax_rationality: float = 3.0,
    leaderboard_freq: int = 10_000,
    n_eval_episodes: int = 10,
    leaderboard_dir: str = "sac_2p_leaderboard",
    **kwargs,
  ) -> None:
    self.ctrl_action_dim = None if ctrl_action_dim is None else int(ctrl_action_dim)
    self.ctrl_update_period = int(ctrl_update_period)
    self.dstb_update_period = int(dstb_update_period)
    # Raw per-network lr args (None = fall back to shared learning_rate); the
    # numeric lrs are resolved in _setup_model once self.lr_schedule exists.
    self._critic_lr_arg = critic_learning_rate
    self._dstb_lr_arg = dstb_learning_rate
    self._ent_coef_lr_arg = ent_coef_lr
    self._dstb_ent_coef_lr_arg = dstb_ent_coef_lr
    # StepLR decay config for the ctrl/dstb/critic optimizers.
    self._lr_schedule_on = bool(lr_schedule)
    self._lr_period = max(1, int(lr_period))
    self._lr_decay = float(lr_decay)
    self._lr_end = float(lr_end)
    self.use_leaderboard = bool(use_leaderboard)
    self._lb_eval_env = leaderboard_eval_env
    self._lb_cfg = dict(
      save_top_k_ctrl=save_top_k_ctrl,
      save_top_k_dstb=save_top_k_dstb,
      softmax_rationality=softmax_rationality,
      model_dir=leaderboard_dir,
    )
    self.leaderboard_freq = int(leaderboard_freq)
    self.n_eval_episodes = int(n_eval_episodes)
    self._next_lb_step = self.leaderboard_freq
    self._rollout_dstb = None  # scratch dstb actor for the current rollout (None = current)
    policy_kwargs = dict(kwargs.pop("policy_kwargs", None) or {})
    if self.ctrl_action_dim is not None:
      policy_kwargs["ctrl_action_dim"] = self.ctrl_action_dim
    super().__init__(policy, env, policy_kwargs=policy_kwargs, **kwargs)

  def _setup_model(self) -> None:
    # On load, ctrl_action_dim is restored into policy_kwargs before _setup_model.
    if self.ctrl_action_dim is None:
      self.ctrl_action_dim = self.policy_kwargs.get("ctrl_action_dim")
    assert self.ctrl_action_dim is not None, (
      "ctrl_action_dim is required (the number of leading control action dims)."
    )
    super()._setup_model()  # sets ctrl entropy (target = -full_dim) and aliases
    # Resolve the per-network / per-actor lrs now that self.lr_schedule exists.
    # None -> shared learning_rate. ``_ctrl_lr`` is the shared value; the ctrl
    # entropy lr (``_ent_coef_lr``) is read by AbstractSAC._reset_entropy_temp.
    shared_lr = float(self.lr_schedule(1))
    self._ctrl_lr = shared_lr
    self._critic_lr = shared_lr if self._critic_lr_arg is None else float(self._critic_lr_arg)
    self._dstb_lr = shared_lr if self._dstb_lr_arg is None else float(self._dstb_lr_arg)
    self._ent_coef_lr = shared_lr if self._ent_coef_lr_arg is None else float(self._ent_coef_lr_arg)
    self._dstb_ent_coef_lr = (
      shared_lr if self._dstb_ent_coef_lr_arg is None else float(self._dstb_ent_coef_lr_arg))
    # SB3 only defines ent_coef_tensor for FIXED ent_coef; default it for "auto".
    if not hasattr(self, "ent_coef_tensor"):
      self.ent_coef_tensor = None
    self.dstb_actor = self.policy.dstb_actor
    dstb_dim = self.policy.dstb_action_dim

    # Per-actor target entropies (SB3 set ctrl to -full_dim).
    self.target_entropy = float(-self.ctrl_action_dim)
    self.dstb_target_entropy = float(-dstb_dim)

    # Disturbance entropy coefficient (mirror SB3's ctrl setup).
    self.dstb_log_ent_coef = None
    self.dstb_ent_coef_optimizer = None
    self.dstb_ent_coef_tensor = None
    if isinstance(self.ent_coef, str) and self.ent_coef.startswith("auto"):
      init_value = 1.0
      if "_" in self.ent_coef:
        init_value = float(self.ent_coef.split("_")[1])
      self.dstb_log_ent_coef = th.log(
        th.ones(1, device=self.device) * init_value
      ).requires_grad_(True)
      self.dstb_ent_coef_optimizer = th.optim.Adam(
        [self.dstb_log_ent_coef], lr=self._dstb_ent_coef_lr
      )
    else:
      self.dstb_ent_coef_tensor = th.tensor(float(self.ent_coef), device=self.device)

    # Snapshot the dstb init entropy temperature so a gamma jump can reset it
    # too (the ctrl init is snapshotted in AbstractSAC._setup_entropy_bounds).
    self._init_dstb_log_ent_coef = (
      None if self.dstb_log_ent_coef is None
      else self.dstb_log_ent_coef.detach().clone())

    # TwoPlayerSACPolicy._build built actor/dstb_actor/critic at the shared lr,
    # and SB3 built the ctrl ent_coef optimizer at the shared lr; stamp the
    # dedicated lrs so they hold even before the first train() step (the dstb ent
    # optimizer was already built at its dedicated lr above). The ctrl actor
    # keeps the shared lr.
    update_learning_rate(self.critic.optimizer, self._critic_lr)
    update_learning_rate(self.dstb_actor.optimizer, self._dstb_lr)
    if self.ent_coef_optimizer is not None:
      update_learning_rate(self.ent_coef_optimizer, self._ent_coef_lr)

    if self.use_leaderboard:
      self._leaderboard = Leaderboard(seed=self.seed or 0, **self._lb_cfg)
      self._scratch_ctrl = copy.deepcopy(self.policy.actor)
      self._scratch_dstb = copy.deepcopy(self.policy.dstb_actor)
      # The league's estimator. Everything about scoring lives in
      # leaderboard.py; this class only decides WHEN to refresh the board.
      self._league_eval = LeagueEvaluator(
        env=self._lb_eval_env,
        device=self.device,
        n_eval_episodes=self.n_eval_episodes,
        dstb_action_dim=self.policy.dstb_action_dim,
        unscale_action=self.policy.unscale_action,
        success=self._league_success,
        obs_normalizer=lambda: (self.env if hasattr(self.env, "normalize_obs")
                                else None),
      )

  # --- the league's win condition -------------------------------------------
  # The avoid game has no target set, so survival IS the win. ReachAvoidSAC2P
  # overrides this with the extra reach requirement. Passing the rule to the
  # evaluator (rather than a mode flag) is what lets both modes share one
  # scorer without either of them carrying the other's machinery.

  @staticmethod
  def _league_success(safe, reached):
    return safe

  # --- league: roll out against a sampled past disturbance ---
  def _resample_rollout_dstb(self) -> None:
    step = self._leaderboard.sample_dstb_step()
    if step is None:
      self._rollout_dstb = None  # use the current dstb actor
    else:
      self._leaderboard.load_actor(self._scratch_dstb, "dstb", step)
      self._scratch_dstb.set_training_mode(False)
      self._rollout_dstb = self._scratch_dstb

  def collect_rollouts(self, env, callback, train_freq, replay_buffer, action_noise=None, learning_starts=0, log_interval=None):
    if self.use_leaderboard and self.num_timesteps >= learning_starts:
      self._resample_rollout_dstb()
    out = super().collect_rollouts(
      env, callback, train_freq, replay_buffer,
      action_noise=action_noise, learning_starts=learning_starts, log_interval=log_interval,
    )
    if (
      self.use_leaderboard
      and self._lb_eval_env is not None
      and self.num_timesteps >= self._next_lb_step
    ):
      self._leaderboard_step()
      self._next_lb_step += self.leaderboard_freq
    return out

  def _leaderboard_step(self) -> None:
    """Refresh the live frontier of the board, then admit/evict.

    Only the current ctrl's row and the current dstb's column are recomputed —
    archived-vs-archived cells are carried over by :meth:`Leaderboard.prune`, so
    each refresh costs ``nc + nd + 2`` evaluations rather than a full matrix.
    """
    lb = self._leaderboard
    ev = self._league_eval
    step = self.num_timesteps
    nc, nd, kc, kd = len(lb.ctrl_steps), len(lb.dstb_steps), lb.kc, lb.kd
    ctrl_cur, dstb_cur = self.policy.actor, self.policy.dstb_actor
    # current ctrl (row kc) vs each opponent
    for j in range(nd):
      lb.load_actor(self._scratch_dstb, "dstb", lb.dstb_steps[j])
      lb.set_score(kc, j, ev.score(ctrl_cur, self._scratch_dstb))
    lb.set_score(kc, kd, ev.score(ctrl_cur, dstb_cur))  # current dstb
    lb.set_score(kc, kd + 1, ev.score(ctrl_cur, None))  # dummy
    # each saved ctrl vs current dstb (col kd)
    for i in range(nc):
      lb.load_actor(self._scratch_ctrl, "ctrl", lb.ctrl_steps[i])
      lb.set_score(i, kd, ev.score(self._scratch_ctrl, dstb_cur))
    lb.prune(step, ctrl_cur, dstb_cur)
    self.logger.record("leaderboard/n_ctrl", len(lb.ctrl_steps))
    self.logger.record("leaderboard/n_dstb", len(lb.dstb_steps))

  # --- rollout: sample ctrl + dstb, concatenate ---
  def _sample_action(self, learning_starts, action_noise=None, n_envs=1):
    if self.num_timesteps < learning_starts and not (
      self.use_sde and self.use_sde_at_warmup
    ):
      unscaled = np.array([self.action_space.sample() for _ in range(n_envs)])
      scaled = self.policy.scale_action(unscaled)
    else:
      obs_tensor, _ = self.policy.obs_to_tensor(self._last_obs)
      dstb_net = (
        self._rollout_dstb if self._rollout_dstb is not None else self.policy.dstb_actor
      )
      with th.no_grad():
        ctrl = self.policy.actor(obs_tensor, deterministic=False)
        dstb = dstb_net(obs_tensor, deterministic=False)
      scaled = th.cat([ctrl, dstb], dim=1).cpu().numpy()  # actors output in [-1, 1]

    if action_noise is not None:
      scaled = np.clip(scaled + action_noise(), -1, 1)
    buffer_action = scaled
    action = self.policy.unscale_action(scaled)
    return action, buffer_action

  # --- tensor-path rollout: same ctrl+dstb concatenation, on device ---
  def _tensor_policy_actions(self, obs: th.Tensor) -> th.Tensor:
    """Two-player tensor collect: sample ctrl + dstb and concatenate to the
    env's ``ctrl_dim + dstb_dim`` action. The GPU-resident analog of
    :meth:`_sample_action` (numpy path); without it the base single-player
    collect samples ctrl only and mismatches the env's action space. The dstb
    is the league-sampled opponent when active, else the current dstb actor
    (mirrors the numpy path). Actors output in [-1, 1] and the env action space
    is [-1, 1], so no unscaling is needed — the collect loop clamps to bounds.
    """
    dstb_net = (
      self._rollout_dstb if getattr(self, "_rollout_dstb", None) is not None
      else self.policy.dstb_actor
    )
    with th.no_grad():
      ctrl = self.policy.actor(obs, deterministic=False)
      dstb = dstb_net(obs, deterministic=False)
    return th.cat([ctrl, dstb], dim=1)

  # --- entropy-coefficient helper ---
  def _alpha(self, log_coef, optimizer, coef_tensor, target_entropy, log_prob):
    if log_coef is not None and optimizer is not None:
      coef = th.exp(log_coef.detach())
      loss = -(log_coef * (log_prob + target_entropy).detach()).mean()
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()
      # min_alpha/max_alpha floor/ceiling -- both actors share bounds.
      if self._log_min_alpha is not None or self._log_max_alpha is not None:
        with th.no_grad():
          log_coef.clamp_(min=self._log_min_alpha, max=self._log_max_alpha)
    else:
      coef = coef_tensor
    return coef

  def _reset_entropy_temp(self) -> None:
    """A gamma jump resets BOTH actors' entropy temperature (reference resets
    ctrl AND dstb alpha on every jump)."""
    super()._reset_entropy_temp()  # ctrl
    dlec = getattr(self, "dstb_log_ent_coef", None)
    if dlec is None or getattr(self, "_init_dstb_log_ent_coef", None) is None:
      return
    with th.no_grad():
      dlec.data.copy_(self._init_dstb_log_ent_coef)
    if getattr(self, "dstb_ent_coef_optimizer", None) is not None:
      dstb_lr = getattr(self, "_dstb_ent_coef_lr", None)
      if dstb_lr is None:
        dstb_lr = float(self.lr_schedule(1))
      self.dstb_ent_coef_optimizer = th.optim.Adam([dlec], lr=dstb_lr)

  # --- per-network learning-rate control ------------------------------------
  def _steplr_value(self, base_lr: float) -> float:
    """Reference StepLR: ``base * lr_decay ** (env_steps // lr_period)``,
    floored at ``lr_end``."""
    num_decay = int(self.num_timesteps) // self._lr_period
    return max(base_lr * (self._lr_decay ** num_decay), self._lr_end)

  def _dedicated_lr_specs(self):
    """Optimizers with their own lr, as ``{id(opt): (base_lr, steplr_on)}``.

    The critic and dstb actor always carry a dedicated (possibly StepLR-decayed)
    lr. The ctrl actor is only intercepted when StepLR is on -- otherwise it
    falls through to SB3's shared-schedule ``_update_learning_rate`` (unchanged
    behaviour). The entropy (alpha) optimizers hold a constant dedicated lr; an
    alpha-lr StepLR is a follow-up."""
    specs = {
      id(self.critic.optimizer): (self._critic_lr, self._lr_schedule_on),
      id(self.dstb_actor.optimizer): (self._dstb_lr, self._lr_schedule_on),
    }
    if self._lr_schedule_on:
      specs[id(self.actor.optimizer)] = (self._ctrl_lr, True)
    if self.ent_coef_optimizer is not None:
      specs[id(self.ent_coef_optimizer)] = (self._ent_coef_lr, False)
    if self.dstb_ent_coef_optimizer is not None:
      specs[id(self.dstb_ent_coef_optimizer)] = (self._dstb_ent_coef_lr, False)
    return specs

  def _update_learning_rate(self, optimizers) -> None:
    """Stop the blanket SB3 overwrite from collapsing every optimizer onto the
    single shared ``lr_schedule``. Optimizers with a dedicated lr keep it (or its
    StepLR-decayed value); anything else routes through SB3's default update so
    the shared ctrl-actor schedule is preserved."""
    if not isinstance(optimizers, list):
      optimizers = [optimizers]
    specs = self._dedicated_lr_specs()
    shared = []
    for opt in optimizers:
      spec = specs.get(id(opt))
      if spec is None:
        shared.append(opt)
        continue
      base_lr, steplr_on = spec
      update_learning_rate(opt, self._steplr_value(base_lr) if steplr_on else base_lr)
    if shared:
      super()._update_learning_rate(shared)

  def train(self, gradient_steps: int, batch_size: int) -> None:
    self.policy.set_training_mode(True)
    opts = [self.actor.optimizer, self.dstb_actor.optimizer, self.critic.optimizer]
    if self.ent_coef_optimizer is not None:
      opts.append(self.ent_coef_optimizer)
    if self.dstb_ent_coef_optimizer is not None:
      opts.append(self.dstb_ent_coef_optimizer)
    self._update_learning_rate(opts)

    # Accumulate on device; one sync per train(). See AbstractSAC1P.train for
    # the measurement and why this is a prerequisite for CUDA-graph capture.
    _z = th.zeros((), device=self.device)
    critic_loss_sum, ctrl_loss_sum, dstb_loss_sum = _z.clone(), _z.clone(), _z.clone()
    critic_loss_max = th.full((), -float("inf"), device=self.device)
    n_critic = n_ctrl = n_dstb = 0
    for step in range(gradient_steps):
      rd = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env)

      ctrl_pi, ctrl_logp = self.actor.action_log_prob(rd.observations)
      ctrl_logp = ctrl_logp.reshape(-1, 1)
      dstb_pi, dstb_logp = self.dstb_actor.action_log_prob(rd.observations)
      dstb_logp = dstb_logp.reshape(-1, 1)

      ctrl_ent = self._alpha(
        self.log_ent_coef, self.ent_coef_optimizer, self.ent_coef_tensor,
        self.target_entropy, ctrl_logp,
      )
      dstb_ent = self._alpha(
        self.dstb_log_ent_coef, self.dstb_ent_coef_optimizer,
        self.dstb_ent_coef_tensor, self.dstb_target_entropy, dstb_logp,
      )

      # --- critic update (soft max-min next value) ---
      with th.no_grad():
        # log-probs of the next actions are intentionally discarded: entropy is
        # not in the safety critic target (see the pure-HJ backup below).
        next_ctrl, _ = self.actor.action_log_prob(rd.next_observations)
        next_dstb, _ = self.dstb_actor.action_log_prob(rd.next_observations)
        next_action = th.cat([next_ctrl, next_dstb], dim=1)
        next_q = th.cat(self.critic_target(rd.next_observations, next_action), dim=1)
        next_q, _ = th.min(next_q, dim=1, keepdim=True)
        # PURE Isaacs/HJ backup -- NO entropy in the critic target. ISAACS
        # (Hsu et al. 2023) passes entropy_motives=0 for the safety value
        # (eq. 8a: y = (1-g)g' + g*min{g', Q'}); the per-actor temperatures live
        # ONLY in the ctrl/dstb ACTOR losses below (eq. 8b). AbstractSAC2P is
        # never CUMULATIVE (there is no two-player cumulative game -- see the MAP
        # law in registry.algo_name), so the soft term is dropped
        # unconditionally, not gated. The previous
        # `- ctrl_ent*logp + dstb_ent*logp` taxed the game value whose zero
        # level set is the online safety certificate.
        # Backup for THIS learner's mode -- see safety_sb3.backups.
        # ReachAvoidSAC2P -> reach-avoid (eq. 6a); SafetySAC2P -> avoid (eq. 7).
        target_q = self._bellman_target(rd, next_q)

      current_q = self.critic(rd.observations, rd.actions)
      critic_loss = 0.5 * sum(F.mse_loss(cq, target_q) for cq in current_q)
      _cl = critic_loss.detach()
      critic_loss_sum += _cl
      critic_loss_max = th.maximum(critic_loss_max, _cl)
      n_critic += 1
      self.critic.optimizer.zero_grad()
      critic_loss.backward()
      self.critic.optimizer.step()

      # --- ctrl actor update (MAX): maximize Q(s, [a_ctrl, detach a_dstb]) ---
      if step % self.ctrl_update_period == 0:
        with th.no_grad():
          dstb_aux = self.dstb_actor(rd.observations, deterministic=False)
        q_pi = th.cat(
          self.critic(rd.observations, th.cat([ctrl_pi, dstb_aux], dim=1)), dim=1
        )
        min_q, _ = th.min(q_pi, dim=1, keepdim=True)
        ctrl_loss = (ctrl_ent * ctrl_logp - min_q).mean()
        ctrl_loss_sum += ctrl_loss.detach(); n_ctrl += 1
        self.actor.optimizer.zero_grad()
        ctrl_loss.backward()
        self.actor.optimizer.step()

      # --- dstb actor update (MIN): minimize Q(s, [detach a_ctrl, a_dstb]) ---
      if step % self.dstb_update_period == 0:
        with th.no_grad():
          ctrl_aux = self.actor(rd.observations, deterministic=False)
        q_pi = th.cat(
          self.critic(rd.observations, th.cat([ctrl_aux, dstb_pi], dim=1)), dim=1
        )
        min_q, _ = th.min(q_pi, dim=1, keepdim=True)
        dstb_loss = (dstb_ent * dstb_logp + min_q).mean()
        dstb_loss_sum += dstb_loss.detach(); n_dstb += 1
        self.dstb_actor.optimizer.zero_grad()
        dstb_loss.backward()
        self.dstb_actor.optimizer.step()

      if step % self.target_update_interval == 0:
        self._polyak(
          self.critic.parameters(), self.critic_target.parameters(), self.tau
        )
        self._polyak(self.batch_norm_stats, self.batch_norm_stats_target, 1.0)

    self._n_updates += gradient_steps
    self.logger.record("train/n_updates", self._n_updates, exclude="tensorboard")
    # record_mean + window max -- see AbstractSAC._record_window_max.
    if n_critic:
      self.logger.record_mean("train/critic_loss", (critic_loss_sum / n_critic).item())
      self._record_window_max("train/critic_loss_max", critic_loss_max.item())
    if n_ctrl:
      self.logger.record_mean("train/ctrl_actor_loss", (ctrl_loss_sum / n_ctrl).item())
    if n_dstb:
      self.logger.record_mean("train/dstb_actor_loss", (dstb_loss_sum / n_dstb).item())
    # Per-actor entropy temperature (alpha) + gamma; ctrl_ent/dstb_ent hold the
    # last step's values.
    self.logger.record("train/ent_coef_ctrl", float(ctrl_ent.mean()))
    self.logger.record("train/ent_coef_dstb", float(dstb_ent.mean()))
    self.logger.record("train/gamma", float(self.gamma))

  def _excluded_save_params(self):
    # League runtime objects are rebuilt by _setup_model on load; the eval env in
    # particular is unpicklable (e.g. mjlab holds a mujoco MjSpec).
    return super()._excluded_save_params() + [
      "dstb_actor",
      "_lb_eval_env",
      "_leaderboard",
      "_league_eval",
      "_scratch_ctrl",
      "_scratch_dstb",
      "_rollout_dstb",
    ]

  def _get_torch_save_params(self):
    state_dicts, others = super()._get_torch_save_params()
    state_dicts += ["dstb_actor.optimizer"]
    if self.dstb_ent_coef_optimizer is not None:
      state_dicts += ["dstb_ent_coef_optimizer"]
      others += ["dstb_log_ent_coef"]
    return state_dicts, others

SafetySAC2P

Bases: AbstractSAC2P

Two-player off-policy avoid game — ISAACS proper (Hsu et al. 2022, eq. 7)::

V(s) = (1-γ)·g + γ·max_ctrl min_dstb min(g, V')

the robust-invariance value under a worst-case disturbance. No target set, no l — the paper has neither.

This is the class for an adversarial avoid task ("stay standing against a worst-case force"). Do NOT emulate it by giving :class:ReachAvoidSAC2P a degenerate l: no l reduces the reach-avoid operator to avoid (:mod:safety_sb3.backups proves the two conditions are contradictory). l is ignored if present, so an l-carrying replay buffer is harmless.

Source code in safety_sb3/sac_2p.py

class SafetySAC2P(AbstractSAC2P):
  """Two-player off-policy **avoid** game — ISAACS proper (Hsu et al. 2022,
  eq. 7)::

      V(s) = (1-γ)·g + γ·max_ctrl min_dstb min(g, V')

  the robust-invariance value under a worst-case disturbance. No target set, no
  ``l`` — the paper has neither.

  This is the class for an adversarial *avoid* task ("stay standing against a
  worst-case force"). Do NOT emulate it by giving :class:`ReachAvoidSAC2P` a
  degenerate ``l``: no ``l`` reduces the reach-avoid operator to avoid
  (:mod:`safety_sb3.backups` proves the two conditions are contradictory).
  ``l`` is ignored if present, so an ``l``-carrying replay buffer is harmless.
  """

  _MODE = backups.AVOID

ReachAvoidSAC2P

Bases: AbstractSAC2P

Two-player off-policy reach-avoid game — Gameplay Filters (Hsu et al. 2024, eq. 6a), which extends ISAACS to reach-avoid: anchor min(l, g).

Needs a target margin l; the base defaults the replay buffer to the l-carrying one off this _MODE. For a two-player avoid game (no target set) use :class:SafetySAC2P.

Source code in safety_sb3/sac_2p.py

class ReachAvoidSAC2P(AbstractSAC2P):
  """Two-player off-policy **reach-avoid** game — Gameplay Filters (Hsu et al.
  2024, eq. 6a), which extends ISAACS to reach-avoid: anchor ``min(l, g)``.

  Needs a target margin ``l``; the base defaults the replay buffer to the
  ``l``-carrying one off this ``_MODE``. For a two-player *avoid* game (no
  target set) use :class:`SafetySAC2P`.
  """

  _MODE = backups.REACH_AVOID

  @staticmethod
  def _league_success(safe, reached):
    """Reach-avoid: the controller must have reached the target AND stayed safe."""
    return safe & reached

A2C and DQN

safety_sb3.a2c

A2C, written without a fixed Bellman backup.

AbstractA2C
├─ SafetyA2C1P        _MODE = AVOID
├─ ReachAvoidA2C1P    _MODE = REACH_AVOID   (+ _ReachAvoidPlumbing)
└─ CumulativeA2C1P    _MODE = CUMULATIVE

A2C is single-player only — there is no AbstractA2C2P. Nothing in the two-player construction is A2C-specific, but nobody has needed it, and the library does not ship learners it has not exercised.

Like PPO, A2C carries its backup in the rollout buffer, so a learner picks its problem by declaring _MODE and the matching buffer pair follows. Unlike PPO, A2C does not own its rollout loop: it uses SB3's stock OnPolicyAlgorithm.collect_rollouts unchanged, so there is no place to gate the timeout value-bootstrap and no place to hand the buffer its per-step extras.

The extras still get through, because SB3's loop calls _update_info_buffer once per step between env.step() and rollout_buffer.add() — exactly the point the reach-avoid buffer needs — and :meth:AbstractA2C._update_info_buffer forwards from there. That is the whole cost of reach-avoid A2C.

.. warning:: A2C here does NOT gate PPO's timeout value-bootstrap, because it does not override the collect loop. When episodes hit the time limit, SB3 adds gamma * V(terminal) to the reward — and in the two safety modes the reward is the physical margin g(s), so that corrupts the backup. Prefer the PPO-family learners for safety work; A2C is kept for parity and for cheap baselines.

AbstractA2C

Bases: GammaAnnealMixin, A2C

A2C whose Bellman backup is chosen by _MODE.

gamma_anneal (ON by default) anneals the discount 0.99 -> 0.9999 over the first 50% of training (read off rollout_buffer.gamma in GAE); applied via _update_current_progress_remaining each iteration. See gamma_anneal.py.

Source code in safety_sb3/a2c.py

class AbstractA2C(GammaAnnealMixin, A2C):
    """A2C whose Bellman backup is chosen by ``_MODE``.

    ``gamma_anneal`` (ON by default) anneals the discount 0.99 -> 0.9999 over the
    first 50% of training (read off ``rollout_buffer.gamma`` in GAE); applied via
    ``_update_current_progress_remaining`` each iteration. See ``gamma_anneal.py``.
    """

    #: the Bellman operator this learner converges to; concretes set it
    _MODE = backups.AVOID

    def __init__(self, *args, rollout_buffer_class=None,
                 rollout_buffer_kwargs=None, gamma_anneal=True, **kwargs):
        backups.check_mode(self._MODE)
        if rollout_buffer_class is None:
            rollout_buffer_class, _ = rollout_buffer_classes(self._MODE)
        super().__init__(
            *args, rollout_buffer_class=rollout_buffer_class,
            rollout_buffer_kwargs=rollout_buffer_kwargs, **kwargs
        )
        self._setup_gamma_anneal(gamma_anneal)

    def _setup_model(self) -> None:
        # Builds policy, optimizer, and rollout buffer -> safe to check.
        super()._setup_model()
        check_rollout_buffer_mode(self, self.rollout_buffer)

    def _update_info_buffer(self, infos, dones=None) -> None:
        """SB3's per-step infos hook, borrowed to feed the rollout buffer.

        This is the one point in the stock on-policy loop that sees ``infos``
        and runs after ``env.step()`` but before ``rollout_buffer.add()`` —
        which is precisely the contract ``record_extras`` needs. The PPO family
        owns its loop and calls ``record_extras`` outright; A2C does not, so it
        piggybacks here. No-op for every buffer whose operator needs no extras.
        """
        super()._update_info_buffer(infos, dones)
        buffer = getattr(self, "rollout_buffer", None)
        if buffer is not None:
            buffer.record_extras(infos)

SafetyA2C1P

Bases: AbstractA2C

A2C with the avoid backup — V(s) = (1-γ)·g + γ·min(g, V').

Source code in safety_sb3/a2c.py

class SafetyA2C1P(AbstractA2C):
    """A2C with the **avoid** backup — ``V(s) = (1-γ)·g + γ·min(g, V')``."""

    _MODE = backups.AVOID

ReachAvoidA2C1P

Bases: _ReachAvoidPlumbing, AbstractA2C

A2C with the reach-avoid backup — V(s) = (1-γ)·min(l, g) + γ·min(g, max(l, V')).

l arrives via info["l_x"] and is captured by the buffer (see the module docstring for how it reaches there without an A2C rollout loop).

Source code in safety_sb3/a2c.py

class ReachAvoidA2C1P(_ReachAvoidPlumbing, AbstractA2C):
    """A2C with the **reach-avoid** backup —
    ``V(s) = (1-γ)·min(l, g) + γ·min(g, max(l, V'))``.

    ``l`` arrives via ``info["l_x"]`` and is captured by the buffer (see the
    module docstring for how it reaches there without an A2C rollout loop).
    """

    _MODE = backups.REACH_AVOID

CumulativeA2C1P

Bases: AbstractA2C

Ordinary reward-maximizing A2C — V(s) = r + γ·V(s'), i.e. stock A2C reached through this library's mode dispatch. Not a safety learner.

Source code in safety_sb3/a2c.py

class CumulativeA2C1P(AbstractA2C):
    """Ordinary reward-maximizing A2C — ``V(s) = r + γ·V(s')``, i.e. stock A2C
    reached through this library's mode dispatch. Not a safety learner."""

    _MODE = backups.CUMULATIVE

safety_sb3.dqn

DQN, written without a fixed Bellman backup.

AbstractDQN
├─ SafetyDQN1P        _MODE = AVOID
└─ CumulativeDQN1P    _MODE = CUMULATIVE

There is deliberately no ReachAvoidDQN1P. MAP would predict one, and it is the single place in the taxonomy where the product does not close: SB3's discrete-action ReplayBuffer carries no target margin l(s), and the reach-avoid operator cannot be evaluated without it. That is a missing capability, not a naming gap, so :class:AbstractDQN refuses mode='reach-avoid' loudly at construction rather than silently computing something else. Use :class:~safety_sb3.sac_1p.ReachAvoidSAC1P (continuous actions) or add an l-carrying discrete replay buffer.

DQN is also single-player only — a two-player discrete game would need a joint action-value table over both players' actions, which is a different learner.

AbstractDQN

Bases: GammaAnnealMixin, DQN

DQN whose TD target is the backup for self._MODE, dispatched through :mod:safety_sb3.backups like every other learner in the library.

gamma_anneal (ON by default) anneals the discount 0.99 -> 0.9999 over the first 50% of training (read as self.gamma in the TD target of train()); the numpy off-policy loop applies it via _update_current_progress_remaining. See gamma_anneal.py.

Parameters:

Name	Type	Description	Default
mode	str | None	backup to converge to; defaults to the class's _MODE.	None

Source code in safety_sb3/dqn.py

class AbstractDQN(GammaAnnealMixin, DQN):
    """DQN whose TD target is the backup for ``self._MODE``, dispatched through
    :mod:`safety_sb3.backups` like every other learner in the library.

    ``gamma_anneal`` (ON by default) anneals the discount 0.99 -> 0.9999 over the
    first 50% of training (read as ``self.gamma`` in the TD target of
    ``train()``); the numpy off-policy loop applies it via
    ``_update_current_progress_remaining``. See ``gamma_anneal.py``.

    :param mode: backup to converge to; defaults to the class's ``_MODE``.
    """

    _MODE = backups.AVOID

    def __init__(self, *args, mode: str | None = None, gamma_anneal=True,
                 **kwargs):
        self._MODE = backups.check_mode(self._MODE if mode is None else mode)
        if self._MODE == backups.REACH_AVOID:
            raise ValueError(
                f"{type(self).__name__} cannot run mode='reach-avoid': its "
                "replay buffer stores no target margin l(s), so the reach-avoid "
                "operator is not computable here. Use ReachAvoidSAC1P, or add an "
                "l-carrying discrete replay buffer.")
        super().__init__(*args, **kwargs)
        self._setup_gamma_anneal(gamma_anneal)

    def train(self, gradient_steps: int, batch_size: int) -> None:
        """Largely follows the original DQN train method from stable_baselines3.
        The TD target is the backup for ``self._MODE`` (safety_sb3.backups).
        """
        # Switch to train mode (this affects batch norm / dropout)
        self.policy.set_training_mode(True)
        # Update learning rate according to schedule
        self._update_learning_rate(self.policy.optimizer)

        losses = []
        for _ in range(gradient_steps):
            # Sample replay buffer
            replay_data = self.replay_buffer.sample(
                batch_size, env=self._vec_normalize_env
            )  # type: ignore[union-attr]

            with th.no_grad():
                # Compute the next Q-values using the target network
                next_q_values = self.q_net_target(replay_data.next_observations)
                # Follow greedy policy: use the one with the highest value
                next_q_values, _ = next_q_values.max(dim=1)
                # Avoid potential broadcast issue
                next_q_values = next_q_values.reshape(-1, 1)

                # 1-step TD target for this learner's mode -- defined in
                # safety_sb3.backups. For AVOID this is the algebraically
                # identical rearrangement of the form this file used to inline:
                #   (1 - g*nt)*gs + g*nt*min(gs, V')
                #   == nt*((1-g)*gs + g*min(gs, V')) + (1-nt)*gs
                # i.e. the full gs is still returned at terminal states.
                gs = replay_data.rewards  # g(s) (a reward, in CUMULATIVE mode)
                not_done = 1.0 - replay_data.dones
                target_q_values = backups.target(
                    self._MODE, gs, next_q_values, not_done, self.gamma)

            # Get current Q-values estimates
            current_q_values = self.q_net(replay_data.observations)

            # Retrieve the q-values for the actions from the replay buffer
            current_q_values = th.gather(current_q_values, dim=1, index=replay_data.actions.long())

            # Compute Huber loss (less sensitive to outliers)
            loss = F.smooth_l1_loss(current_q_values, target_q_values)
            losses.append(loss.item())

            # Optimize the policy
            self.policy.optimizer.zero_grad()
            loss.backward()
            # Clip gradient norm
            th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
            self.policy.optimizer.step()

        # Increase update counter
        self._n_updates += gradient_steps

        self.logger.record("train/n_updates", self._n_updates, exclude="tensorboard")
        self.logger.record("train/loss", np.mean(losses))

train

train(gradient_steps, batch_size)

Largely follows the original DQN train method from stable_baselines3. The TD target is the backup for self._MODE (safety_sb3.backups).

Source code in safety_sb3/dqn.py

def train(self, gradient_steps: int, batch_size: int) -> None:
    """Largely follows the original DQN train method from stable_baselines3.
    The TD target is the backup for ``self._MODE`` (safety_sb3.backups).
    """
    # Switch to train mode (this affects batch norm / dropout)
    self.policy.set_training_mode(True)
    # Update learning rate according to schedule
    self._update_learning_rate(self.policy.optimizer)

    losses = []
    for _ in range(gradient_steps):
        # Sample replay buffer
        replay_data = self.replay_buffer.sample(
            batch_size, env=self._vec_normalize_env
        )  # type: ignore[union-attr]

        with th.no_grad():
            # Compute the next Q-values using the target network
            next_q_values = self.q_net_target(replay_data.next_observations)
            # Follow greedy policy: use the one with the highest value
            next_q_values, _ = next_q_values.max(dim=1)
            # Avoid potential broadcast issue
            next_q_values = next_q_values.reshape(-1, 1)

            # 1-step TD target for this learner's mode -- defined in
            # safety_sb3.backups. For AVOID this is the algebraically
            # identical rearrangement of the form this file used to inline:
            #   (1 - g*nt)*gs + g*nt*min(gs, V')
            #   == nt*((1-g)*gs + g*min(gs, V')) + (1-nt)*gs
            # i.e. the full gs is still returned at terminal states.
            gs = replay_data.rewards  # g(s) (a reward, in CUMULATIVE mode)
            not_done = 1.0 - replay_data.dones
            target_q_values = backups.target(
                self._MODE, gs, next_q_values, not_done, self.gamma)

        # Get current Q-values estimates
        current_q_values = self.q_net(replay_data.observations)

        # Retrieve the q-values for the actions from the replay buffer
        current_q_values = th.gather(current_q_values, dim=1, index=replay_data.actions.long())

        # Compute Huber loss (less sensitive to outliers)
        loss = F.smooth_l1_loss(current_q_values, target_q_values)
        losses.append(loss.item())

        # Optimize the policy
        self.policy.optimizer.zero_grad()
        loss.backward()
        # Clip gradient norm
        th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
        self.policy.optimizer.step()

    # Increase update counter
    self._n_updates += gradient_steps

    self.logger.record("train/n_updates", self._n_updates, exclude="tensorboard")
    self.logger.record("train/loss", np.mean(losses))

SafetyDQN1P

Bases: AbstractDQN

DQN with the avoid backup — Q(s,a) = (1-γ)·g + γ·min(g, max_a' Q') (Fisac et al. 2019). g(s) rides on the reward channel.

Source code in safety_sb3/dqn.py

class SafetyDQN1P(AbstractDQN):
    """DQN with the **avoid** backup — ``Q(s,a) = (1-γ)·g + γ·min(g, max_a' Q')``
    (Fisac et al. 2019). ``g(s)`` rides on the reward channel."""

    _MODE = backups.AVOID

CumulativeDQN1P

Bases: AbstractDQN

Ordinary DQN — Q(s,a) = r + γ·max_a' Q'. Not a safety learner; it is here so a reward-maximizing discrete baseline shares this library's code.

Source code in safety_sb3/dqn.py

class CumulativeDQN1P(AbstractDQN):
    """Ordinary DQN — ``Q(s,a) = r + γ·max_a' Q'``. Not a safety learner; it is
    here so a reward-maximizing discrete baseline shares this library's code."""

    _MODE = backups.CUMULATIVE

The reach-avoid mixin

safety_sb3.reach_avoid_mixin

The reach-avoid plumbing for the on-policy family — a MIXIN, not a base.

Why a mixin. MAP is a product: Mode x Algorithm x Players. Single inheritance can linearise only one axis, so one of them has to be composition, and the players axis is the one that carries the code — the two on-policy update loops are hundreds of lines apart, while the reach-avoid delta below is a couple of dozen. Making reach-avoid a parent of the loops would therefore mean duplicating a loop per mode; making it a mixin means it composes with both::

class ReachAvoidPPO1P(_ReachAvoidPlumbing, AbstractPPO1P): _MODE = REACH_AVOID
class ReachAvoidPPO2P(_ReachAvoidPlumbing, AbstractPPO2P): _MODE = REACH_AVOID
class SafetyPPO1P(AbstractPPO1P):                          _MODE = AVOID

and — the part that matters for correctness — the avoid learners never receive it at all. Before v0.4.0 the avoid classes inherited the reach-avoid class and switched its machinery back off through an _is_reach_avoid predicate; every piece of l-handling then had to be guarded, and one missed guard silently fed an avoid learner a target margin. There is nothing to switch off now.

What is left here is deliberately small. The capture of l(s) itself used to live in the algorithm (a _record_step_extras hook on the rollout loop) and was the bulk of this file; it now lives in the rollout buffers, which is where the rest of the l plumbing already was — see :mod:safety_sb3.buffers_rollout. What remains is one constructor argument and the reach half of the two-player league's win condition.

SAC needs no equivalent: its mode delta is a handful of lines that already sit inside :class:~safety_sb3.sac_base.AbstractSAC and dispatch on _MODE.

The two-player league

safety_sb3.leaderboard

The league: a policy archive shared by both two-player families.

Mirrors safe_adaptation_dev's leaderboard: a checkpoint archive of the best control (max-player) and disturbance (min-player) actors, scored against the opponents on the board, with worst-performer eviction and softmax-rationality sampling of past disturbances into rollouts (the "strategy-space regularization" that damps cycling).

Score matrix board has shape (kc+1, kd+2): * rows 0..kc-1 = saved ctrl checkpoints; row -1 = the current ctrl. * cols 0..kd-1 = saved dstb checkpoints; col -2 = the current dstb; col -1 = the dummy (no-disturbance) opponent. board[i, j] = success metric of ctrl i vs dstb j (higher = ctrl better / dstb weaker). Control is a maximizer; disturbance a minimizer.

Both two-player families use :class:Leaderboard, but they FILL it differently and that difference is deliberate, not an oversight:

• the off-policy (SAC) league scores with dedicated evaluation episodes on a separate eval env (:class:LeagueEvaluator below), writes cells with :meth:Leaderboard.set_score (an overwrite), and samples ONE archived opponent per rollout via :meth:Leaderboard.sample_dstb_step;
• the on-policy (PPO) league has no eval env at all. It scores from the training rollout outcomes it already produced, which are noisy, so it accumulates them with :meth:Leaderboard.ema_score, and it faces a whole POPULATION of opponents at once by assigning them to env slices via :meth:Leaderboard.sample_dstb_slices.

Neither is a special case of the other: dedicated-eval + overwrite + one opponent is not the same estimator as training-outcome + EMA + population, and on-policy learning cannot reuse another opponent's data the way replay can. So this module offers both vocabularies rather than forcing one.

Leaderboard

Source code in safety_sb3/leaderboard.py

class Leaderboard:
  def __init__(
    self,
    save_top_k_ctrl: int,
    save_top_k_dstb: int,
    softmax_rationality: float,
    model_dir: str,
    seed: int = 0,
  ) -> None:
    self.kc = int(save_top_k_ctrl)
    self.kd = int(save_top_k_dstb)
    self.rationality = float(softmax_rationality)
    self.dir = model_dir
    os.makedirs(self.dir, exist_ok=True)
    self.rng = np.random.default_rng(seed)
    self.ctrl_steps: list[int] = []  # parallel to board rows 0..kc-1
    self.dstb_steps: list[int] = []  # parallel to board cols 0..kd-1
    self.board = np.full((self.kc + 1, self.kd + 2), np.nan, dtype=float)

  # --- checkpoint io ---
  def _path(self, kind: str, step: int) -> str:
    return os.path.join(self.dir, f"{kind}_{step}.pt")

  def save_actor(self, actor: th.nn.Module, kind: str, step: int) -> None:
    th.save(actor.state_dict(), self._path(kind, step))

  def load_actor(self, actor: th.nn.Module, kind: str, step: int) -> th.nn.Module:
    actor.load_state_dict(th.load(self._path(kind, step), map_location="cpu"))
    return actor

  def _remove(self, kind: str, step: int) -> None:
    p = self._path(kind, step)
    if os.path.exists(p):
      os.remove(p)

  # --- score matrix ---
  def set_score(self, ctrl_idx: int, dstb_idx: int, metric: float) -> None:
    """Overwrite a cell — used when the score came from a dedicated eval run
    (contrast :meth:`ema_score`)."""
    self.board[ctrl_idx, dstb_idx] = metric

  def prune(self, step: int, ctrl_actor: th.nn.Module, dstb_actor: th.nn.Module) -> None:
    """Admit the current ctrl/dstb (row -1 / col -2) into the archive, evicting
    the worst performer when full.  Matches ``safe_adaptation_dev.prune_leaderboard``."""
    # --- control (maximizer): evict the lowest-average-metric checkpoint ---
    if len(self.ctrl_steps) == self.kc:
      ctrl_avg = np.nanmean(self.board, axis=1)  # over all dstb cols
      worst = int(np.argmin(ctrl_avg))
      if worst != self.kc:  # current beats some saved ctrl -> replace it
        self._remove("ctrl", self.ctrl_steps[worst])
        self.ctrl_steps[worst] = step
        self.board[worst] = self.board[-1]
        self.save_actor(ctrl_actor, "ctrl", step)
    else:
      self.ctrl_steps.append(step)
      self.save_actor(ctrl_actor, "ctrl", step)

    # --- disturbance (minimizer): evict the highest-average-metric checkpoint ---
    if len(self.dstb_steps) == self.kd:
      dstb_avg = np.nanmean(self.board[:, :-1], axis=0)  # exclude dummy col
      worst = int(np.argmax(dstb_avg))
      if worst != self.kd:  # current is more challenging than some saved dstb
        self._remove("dstb", self.dstb_steps[worst])
        self.dstb_steps[worst] = step
        self.board[:, worst] = self.board[:, -2]
        self.save_actor(dstb_actor, "dstb", step)
    else:
      self.dstb_steps.append(step)
      self.save_actor(dstb_actor, "dstb", step)

  # --- rollout sampling (softmax over how effective each dstb is) ---
  def sample_dstb_step(self) -> int | None:
    """Return a saved dstb step to roll out against, or ``None`` for the current
    dstb / dummy.  Probability ∝ exp(-rationality · avg_metric): a disturbance
    that drives the control's reach-avoid value *down* (low metric) is sampled
    more often.  Falls back to the current dstb when the board is empty."""
    n = len(self.dstb_steps)
    if n == 0 or not self.ctrl_steps:
      return None
    # choices: the n saved dstb cols, plus the dummy (-1); col -2 = "current".
    cols = list(range(n)) + [self.board.shape[1] - 1]
    with np.errstate(invalid="ignore"):
      logit = np.nanmean(self.board[: len(self.ctrl_steps)][:, cols], axis=0)
      fill = float(np.nanmean(self.board)) if np.isfinite(self.board).any() else 0.0
    logit = np.nan_to_num(logit, nan=fill)
    p = np.exp(-self.rationality * logit)
    p = p / p.sum()
    pick = int(self.rng.choice(len(cols), p=p))
    if cols[pick] == self.board.shape[1] - 1:
      return None  # dummy / current dstb (no archived checkpoint)
    return self.dstb_steps[pick]

  # --- on-policy (PPO 2P) additions -----------------------------------------

  def ema_score(self, ctrl_idx: int, dstb_idx: int, metric: float,
                beta: float = 0.1) -> None:
    """EMA score update — the on-policy board is filled from noisy per-rollout
    TRAINING outcomes rather than from dedicated eval runs, so cells are
    accumulated instead of overwritten (contrast :meth:`set_score`)."""
    old = self.board[ctrl_idx, dstb_idx]
    if np.isnan(old):
      self.board[ctrl_idx, dstb_idx] = metric
    else:
      self.board[ctrl_idx, dstb_idx] = (1.0 - beta) * old + beta * metric

  def sample_dstb_slices(self, n_slices: int) -> list[int]:
    """Opponent assignment for one on-policy rollout over env SLICES.

    Returns a list of length ``n_slices`` with entries: ``-3`` = ZERO (dummy
    opponent; also anchors nominal behavior), ``-2`` = RANDOM, ``-1`` =
    CURRENT dstb, ``>= 0`` = index into ``dstb_steps`` (archived checkpoint).
    Slices 0 and 1 are always ZERO and RANDOM; the rest are softmax-sampled
    with ``p ∝ exp(-rationality * avg_metric)`` over {archived, CURRENT}.
    """
    out = [-3, -2]
    n_free = max(n_slices - 2, 0)
    n_arch = len(self.dstb_steps)
    if n_arch == 0:
      return (out + [-1] * n_free)[:n_slices]
    cols = list(range(n_arch)) + [self.board.shape[1] - 2]
    with np.errstate(invalid="ignore"):
      logit = np.nanmean(self.board[:, cols], axis=0)
      fill = float(np.nanmean(self.board)) if np.isfinite(self.board).any() else 0.0
    logit = np.nan_to_num(logit, nan=fill)
    p = np.exp(-self.rationality * logit)
    p = p / p.sum()
    for _ in range(n_free):
      pick = int(self.rng.choice(len(cols), p=p))
      out.append(-1 if cols[pick] == self.board.shape[1] - 2 else pick)
    return out[:n_slices]

set_score

set_score(ctrl_idx, dstb_idx, metric)

Overwrite a cell — used when the score came from a dedicated eval run (contrast :meth:ema_score).

Source code in safety_sb3/leaderboard.py

def set_score(self, ctrl_idx: int, dstb_idx: int, metric: float) -> None:
  """Overwrite a cell — used when the score came from a dedicated eval run
  (contrast :meth:`ema_score`)."""
  self.board[ctrl_idx, dstb_idx] = metric

prune

prune(step, ctrl_actor, dstb_actor)

Admit the current ctrl/dstb (row -1 / col -2) into the archive, evicting the worst performer when full. Matches safe_adaptation_dev.prune_leaderboard.

Source code in safety_sb3/leaderboard.py

def prune(self, step: int, ctrl_actor: th.nn.Module, dstb_actor: th.nn.Module) -> None:
  """Admit the current ctrl/dstb (row -1 / col -2) into the archive, evicting
  the worst performer when full.  Matches ``safe_adaptation_dev.prune_leaderboard``."""
  # --- control (maximizer): evict the lowest-average-metric checkpoint ---
  if len(self.ctrl_steps) == self.kc:
    ctrl_avg = np.nanmean(self.board, axis=1)  # over all dstb cols
    worst = int(np.argmin(ctrl_avg))
    if worst != self.kc:  # current beats some saved ctrl -> replace it
      self._remove("ctrl", self.ctrl_steps[worst])
      self.ctrl_steps[worst] = step
      self.board[worst] = self.board[-1]
      self.save_actor(ctrl_actor, "ctrl", step)
  else:
    self.ctrl_steps.append(step)
    self.save_actor(ctrl_actor, "ctrl", step)

  # --- disturbance (minimizer): evict the highest-average-metric checkpoint ---
  if len(self.dstb_steps) == self.kd:
    dstb_avg = np.nanmean(self.board[:, :-1], axis=0)  # exclude dummy col
    worst = int(np.argmax(dstb_avg))
    if worst != self.kd:  # current is more challenging than some saved dstb
      self._remove("dstb", self.dstb_steps[worst])
      self.dstb_steps[worst] = step
      self.board[:, worst] = self.board[:, -2]
      self.save_actor(dstb_actor, "dstb", step)
  else:
    self.dstb_steps.append(step)
    self.save_actor(dstb_actor, "dstb", step)

sample_dstb_step

sample_dstb_step()

Return a saved dstb step to roll out against, or None for the current dstb / dummy. Probability ∝ exp(-rationality · avg_metric): a disturbance that drives the control's reach-avoid value down (low metric) is sampled more often. Falls back to the current dstb when the board is empty.

Source code in safety_sb3/leaderboard.py

def sample_dstb_step(self) -> int | None:
  """Return a saved dstb step to roll out against, or ``None`` for the current
  dstb / dummy.  Probability ∝ exp(-rationality · avg_metric): a disturbance
  that drives the control's reach-avoid value *down* (low metric) is sampled
  more often.  Falls back to the current dstb when the board is empty."""
  n = len(self.dstb_steps)
  if n == 0 or not self.ctrl_steps:
    return None
  # choices: the n saved dstb cols, plus the dummy (-1); col -2 = "current".
  cols = list(range(n)) + [self.board.shape[1] - 1]
  with np.errstate(invalid="ignore"):
    logit = np.nanmean(self.board[: len(self.ctrl_steps)][:, cols], axis=0)
    fill = float(np.nanmean(self.board)) if np.isfinite(self.board).any() else 0.0
  logit = np.nan_to_num(logit, nan=fill)
  p = np.exp(-self.rationality * logit)
  p = p / p.sum()
  pick = int(self.rng.choice(len(cols), p=p))
  if cols[pick] == self.board.shape[1] - 1:
    return None  # dummy / current dstb (no archived checkpoint)
  return self.dstb_steps[pick]

ema_score

ema_score(ctrl_idx, dstb_idx, metric, beta=0.1)

EMA score update — the on-policy board is filled from noisy per-rollout TRAINING outcomes rather than from dedicated eval runs, so cells are accumulated instead of overwritten (contrast :meth:set_score).

Source code in safety_sb3/leaderboard.py

def ema_score(self, ctrl_idx: int, dstb_idx: int, metric: float,
              beta: float = 0.1) -> None:
  """EMA score update — the on-policy board is filled from noisy per-rollout
  TRAINING outcomes rather than from dedicated eval runs, so cells are
  accumulated instead of overwritten (contrast :meth:`set_score`)."""
  old = self.board[ctrl_idx, dstb_idx]
  if np.isnan(old):
    self.board[ctrl_idx, dstb_idx] = metric
  else:
    self.board[ctrl_idx, dstb_idx] = (1.0 - beta) * old + beta * metric

sample_dstb_slices

sample_dstb_slices(n_slices)

Opponent assignment for one on-policy rollout over env SLICES.

Returns a list of length n_slices with entries: -3 = ZERO (dummy opponent; also anchors nominal behavior), -2 = RANDOM, -1 = CURRENT dstb, >= 0 = index into dstb_steps (archived checkpoint). Slices 0 and 1 are always ZERO and RANDOM; the rest are softmax-sampled with p ∝ exp(-rationality * avg_metric) over {archived, CURRENT}.

Source code in safety_sb3/leaderboard.py

def sample_dstb_slices(self, n_slices: int) -> list[int]:
  """Opponent assignment for one on-policy rollout over env SLICES.

  Returns a list of length ``n_slices`` with entries: ``-3`` = ZERO (dummy
  opponent; also anchors nominal behavior), ``-2`` = RANDOM, ``-1`` =
  CURRENT dstb, ``>= 0`` = index into ``dstb_steps`` (archived checkpoint).
  Slices 0 and 1 are always ZERO and RANDOM; the rest are softmax-sampled
  with ``p ∝ exp(-rationality * avg_metric)`` over {archived, CURRENT}.
  """
  out = [-3, -2]
  n_free = max(n_slices - 2, 0)
  n_arch = len(self.dstb_steps)
  if n_arch == 0:
    return (out + [-1] * n_free)[:n_slices]
  cols = list(range(n_arch)) + [self.board.shape[1] - 2]
  with np.errstate(invalid="ignore"):
    logit = np.nanmean(self.board[:, cols], axis=0)
    fill = float(np.nanmean(self.board)) if np.isfinite(self.board).any() else 0.0
  logit = np.nan_to_num(logit, nan=fill)
  p = np.exp(-self.rationality * logit)
  p = p / p.sum()
  for _ in range(n_free):
    pick = int(self.rng.choice(len(cols), p=p))
    out.append(-1 if cols[pick] == self.board.shape[1] - 2 else pick)
  return out[:n_slices]

LeagueEvaluator

Scores one (ctrl, dstb) pair by rolling out dedicated eval episodes.

This is the off-policy league's estimator, lifted out of the algorithm so that everything about the league lives in this module. It picks the cheapest of three equivalent rollout paths automatically:

• step_tensor on a GPU-resident env — everything on device, no numpy VecEnv and no per-step host<->device sync. ~10-50x faster than the numpy path and the main league-throughput fix (profiling: the numpy vectorized path cost ~100s per refresh);
• a parallel numpy VecEnv — one batch = num_envs first-episodes;
• a single raw gym.Env loop.

All three share the success semantics, which the caller supplies:

Parameters:

Name	Type	Description	Default
success		(safe, reached) -> won, evaluated elementwise. The avoid game passes safe (survival IS the win — there is no target set); the reach-avoid game passes safe & reached. Passing the rule in is what lets the two modes share one evaluator without a mode flag inside it.	required
obs_normalizer		zero-arg callable returning the LIVE training normalizer (or None). Only the tensor path needs it: it rolls out on a RAW tensor env, so observations must be normalized with the same running stats the actors were trained under. The margins g/l are physical (unnormalized) and drive safe/reached either way.	None

Source code in safety_sb3/leaderboard.py

class LeagueEvaluator:
  """Scores one ``(ctrl, dstb)`` pair by rolling out dedicated eval episodes.

  This is the off-policy league's estimator, lifted out of the algorithm so that
  everything about the league lives in this module. It picks the cheapest of
  three equivalent rollout paths automatically:

  * ``step_tensor`` on a GPU-resident env — everything on device, no numpy
    ``VecEnv`` and no per-step host<->device sync. ~10-50x faster than the numpy
    path and the main league-throughput fix (profiling: the numpy vectorized path
    cost ~100s per refresh);
  * a parallel numpy ``VecEnv`` — one batch = ``num_envs`` first-episodes;
  * a single raw ``gym.Env`` loop.

  All three share the success semantics, which the caller supplies:

  :param success: ``(safe, reached) -> won``, evaluated elementwise. The avoid
      game passes ``safe`` (survival IS the win — there is no target set); the
      reach-avoid game passes ``safe & reached``. Passing the rule in is what
      lets the two modes share one evaluator without a mode flag inside it.
  :param obs_normalizer: zero-arg callable returning the LIVE training
      normalizer (or ``None``). Only the tensor path needs it: it rolls out on a
      RAW tensor env, so observations must be normalized with the same running
      stats the actors were trained under. The margins ``g``/``l`` are physical
      (unnormalized) and drive safe/reached either way.
  """

  MAX_STEPS = 400

  def __init__(self, env, device, n_eval_episodes: int, dstb_action_dim: int,
               unscale_action, success, obs_normalizer=None) -> None:
    self.env = env
    self.device = device
    self.n_eval_episodes = int(n_eval_episodes)
    self.dn = int(dstb_action_dim)
    self.unscale_action = unscale_action
    self.success = success
    self.obs_normalizer = obs_normalizer

  @th.no_grad()
  def score(self, ctrl_actor, dstb_actor) -> float:
    """Success rate of ``ctrl`` vs ``dstb`` (``None`` = dummy / no disturbance)."""
    env = self.env
    if getattr(env, "is_tensor_env", False):
      return self._score_tensor(env, ctrl_actor, dstb_actor)
    if hasattr(env, "num_envs") and hasattr(env, "step_async"):
      return self._score_vec(env, ctrl_actor, dstb_actor)

    dn = self.dn
    succ = 0
    for _ in range(self.n_eval_episodes):
      obs, _ = env.reset()
      done = reached = False
      safe = True
      while not done:
        ot = th.as_tensor(np.asarray(obs), dtype=th.float32,
                          device=self.device).reshape(1, -1)
        c = ctrl_actor(ot, deterministic=True).cpu().numpy()[0]
        d = (np.zeros(dn, np.float32) if dstb_actor is None
             else dstb_actor(ot, deterministic=True).cpu().numpy()[0])
        scaled = np.concatenate([c, d]).astype(np.float32)  # [-1, 1]
        action = self.unscale_action(scaled[None])[0]
        obs, g, term, trunc, info = env.step(action)
        done = bool(term or trunc)
        if g < 0:
          safe = False
        if float(info.get("l_x", -1.0)) >= 0:
          reached = True
      succ += int(self.success(safe, reached))
    return succ / max(self.n_eval_episodes, 1)

  def _score_vec(self, env, ctrl_actor, dstb_actor) -> float:
    """Parallel success over ``num_envs`` * ``n_eval_episodes`` first-episodes."""
    dn = self.dn
    n_env = env.num_envs
    total_succ, total = 0, 0
    for _ in range(max(1, self.n_eval_episodes)):
      obs = env.reset()
      ep_safe = np.ones(n_env, dtype=bool)
      ep_reached = np.zeros(n_env, dtype=bool)
      done_once = np.zeros(n_env, dtype=bool)
      for _ in range(self.MAX_STEPS):
        ot = th.as_tensor(np.asarray(obs), dtype=th.float32, device=self.device)
        c = ctrl_actor(ot, deterministic=True).cpu().numpy()
        d = (np.zeros((n_env, dn), np.float32) if dstb_actor is None
             else dstb_actor(ot, deterministic=True).cpu().numpy())
        env.step_async(np.concatenate([c, d], axis=1).astype(np.float32))
        obs, g, dones, infos = env.step_wait()
        active = ~done_once
        ep_safe &= ~(active & (np.asarray(g) < 0))
        lx = np.array([i.get("l_x", -1.0) for i in infos], dtype=np.float32)
        ep_reached |= active & (lx >= 0)
        done_once |= active & np.asarray(dones, dtype=bool)
        if done_once.all():
          break
      hits = self.success(ep_safe, ep_reached)
      total_succ += int(hits.sum())
      total += n_env
    return total_succ / max(total, 1)

  @th.no_grad()
  def _score_tensor(self, env, ctrl_actor, dstb_actor) -> float:
    """GPU-resident twin of :meth:`_score_vec`. Actors output in [-1, 1] and the
    env clamps, so no unscaling is needed here."""
    dn = self.dn
    dev = env.device
    n = env.num_envs
    norm = self.obs_normalizer() if self.obs_normalizer is not None else None
    total_succ, total = 0, 0
    for _ in range(max(1, self.n_eval_episodes)):
      raw = env.reset()
      if not th.is_tensor(raw):
        raw = th.as_tensor(np.asarray(raw), dtype=th.float32, device=dev)
      ep_safe = th.ones(n, dtype=th.bool, device=dev)
      ep_reached = th.zeros(n, dtype=th.bool, device=dev)
      done_once = th.zeros(n, dtype=th.bool, device=dev)
      for _ in range(self.MAX_STEPS):
        obs = norm.normalize_obs(raw) if norm is not None else raw
        c = ctrl_actor(obs, deterministic=True)
        d = (th.zeros((n, dn), device=dev) if dstb_actor is None
             else dstb_actor(obs, deterministic=True))
        raw, g, dones, _timeouts, l_x = env.step_tensor(th.cat([c, d], dim=1))
        active = ~done_once
        ep_safe &= ~(active & (g.reshape(n) < 0))
        ep_reached |= active & (l_x.reshape(n) >= 0)
        done_once |= active & dones.reshape(n).bool()
        if bool(done_once.all()):
          break
      hits = self.success(ep_safe, ep_reached)
      total_succ += int(hits.sum().item())
      total += n
    return total_succ / max(total, 1)

score

score(ctrl_actor, dstb_actor)

Success rate of ctrl vs dstb (None = dummy / no disturbance).

Source code in safety_sb3/leaderboard.py

@th.no_grad()
def score(self, ctrl_actor, dstb_actor) -> float:
  """Success rate of ``ctrl`` vs ``dstb`` (``None`` = dummy / no disturbance)."""
  env = self.env
  if getattr(env, "is_tensor_env", False):
    return self._score_tensor(env, ctrl_actor, dstb_actor)
  if hasattr(env, "num_envs") and hasattr(env, "step_async"):
    return self._score_vec(env, ctrl_actor, dstb_actor)

  dn = self.dn
  succ = 0
  for _ in range(self.n_eval_episodes):
    obs, _ = env.reset()
    done = reached = False
    safe = True
    while not done:
      ot = th.as_tensor(np.asarray(obs), dtype=th.float32,
                        device=self.device).reshape(1, -1)
      c = ctrl_actor(ot, deterministic=True).cpu().numpy()[0]
      d = (np.zeros(dn, np.float32) if dstb_actor is None
           else dstb_actor(ot, deterministic=True).cpu().numpy()[0])
      scaled = np.concatenate([c, d]).astype(np.float32)  # [-1, 1]
      action = self.unscale_action(scaled[None])[0]
      obs, g, term, trunc, info = env.step(action)
      done = bool(term or trunc)
      if g < 0:
        safe = False
      if float(info.get("l_x", -1.0)) >= 0:
        reached = True
    succ += int(self.success(safe, reached))
  return succ / max(self.n_eval_episodes, 1)

Policies

safety_sb3.policies

Policies. Currently one: the two-player SAC policy.

:class:TwoPlayerSACPolicy is SAC-only, and the name says so deliberately. It exists because the SAC-family two-player game needs ONE critic over the JOINT action Q(s, [a_ctrl, a_dstb]) — so both actors and that shared critic have to be built together, inside a single policy object. The PPO-family two-player learner has no use for it: PPO's critic is state-only (V(s)), so each player gets an independent, ordinary single-player policy built over its own action sub-space, and :class:~safety_sb3.ppo_2p.AbstractPPO2P simply holds two of them. Naming this a bare TwoPlayerPolicy would overclaim.

It subclasses SB3's :class:SACPolicy so the two-player learners stay plug-and-play for existing SB3 code: model.predict / save / load behave normally and return the control action. Internally it holds two actors over disjoint sub-action spaces and one twin critic over the FULL (concatenated) action:

• self.actor — control actor (max-player), over the first ctrl_action_dim action dims. This is the deployable safety controller; predict uses it.
• self.dstb_actor — disturbance actor (min-player), over the remaining dims.
• self.critic / self.critic_target — Q(s, [a_ctrl, a_dstb]) over the full action space.

The env's action_space is the concatenation Box(ctrl_dim + dstb_dim); the env splits it and applies ctrl as control + dstb as a bounded disturbance. ctrl_action_dim tells the policy where to split (pass via policy_kwargs).

TwoPlayerSACPolicy

Bases: SACPolicy

Source code in safety_sb3/policies.py

class TwoPlayerSACPolicy(SACPolicy):
  def __init__(self, *args, ctrl_action_dim: int, **kwargs) -> None:
    # Must be set before super().__init__ → _build().
    self.ctrl_action_dim = int(ctrl_action_dim)
    super().__init__(*args, **kwargs)

  # --- helpers ---
  def _split_action_space(self) -> tuple[spaces.Box, spaces.Box]:
    full = self.action_space
    assert isinstance(full, spaces.Box) and full.shape is not None
    c = self.ctrl_action_dim
    assert 0 < c < full.shape[0], (
      f"ctrl_action_dim={c} must be in (0, {full.shape[0]}) for a "
      f"concatenated [ctrl, dstb] action space."
    )
    ctrl = spaces.Box(full.low[:c], full.high[:c], dtype=full.dtype)  # type: ignore[index]
    dstb = spaces.Box(full.low[c:], full.high[c:], dtype=full.dtype)  # type: ignore[index]
    return ctrl, dstb

  def _make_actor_for(self, action_space: spaces.Box) -> Actor:
    actor_kwargs = self._update_features_extractor(self.actor_kwargs, None)
    actor_kwargs["action_space"] = action_space
    return Actor(**actor_kwargs).to(self.device)

  # --- build two actors + full critic ---
  def _build(self, lr_schedule: Schedule) -> None:
    ctrl_space, dstb_space = self._split_action_space()
    self._ctrl_space, self._dstb_space = ctrl_space, dstb_space

    # Control actor == self.actor (used by predict / deployment).
    self.actor = self._make_actor_for(ctrl_space)
    self.actor.optimizer = self.optimizer_class(
      self.actor.parameters(), lr=lr_schedule(1), **self.optimizer_kwargs
    )
    # Disturbance actor (min-player).
    self.dstb_actor = self._make_actor_for(dstb_space)
    self.dstb_actor.optimizer = self.optimizer_class(
      self.dstb_actor.parameters(), lr=lr_schedule(1), **self.optimizer_kwargs
    )

    # Twin critic over the FULL concatenated action (own features extractor).
    self.critic = self.make_critic(features_extractor=None)
    self.critic.optimizer = self.optimizer_class(
      list(self.critic.parameters()), lr=lr_schedule(1), **self.optimizer_kwargs
    )
    self.critic_target = self.make_critic(features_extractor=None)
    self.critic_target.load_state_dict(self.critic.state_dict())
    self.critic_target.set_training_mode(False)

  def _get_constructor_parameters(self) -> dict[str, Any]:
    data = super()._get_constructor_parameters()
    data.update(ctrl_action_dim=self.ctrl_action_dim)
    return data

  def set_training_mode(self, mode: bool) -> None:
    self.actor.set_training_mode(mode)
    self.dstb_actor.set_training_mode(mode)
    self.critic.set_training_mode(mode)
    self.training = mode

  @property
  def dstb_action_dim(self) -> int:
    assert self.action_space.shape is not None
    return int(self.action_space.shape[0]) - self.ctrl_action_dim

  def predict(self, observation, state=None, episode_start=None, deterministic=False):
    """Deployment: return the CONTROL action only (the safety controller, no
    adversary), in the control sub-space bounds.  Overrides SB3's predict, which
    would reshape the ctrl-dim output to the full concatenated action space."""
    self.set_training_mode(False)
    obs_tensor, vectorized = self.obs_to_tensor(observation)
    with __import__("torch").no_grad():
      actions = self.actor(obs_tensor, deterministic=deterministic)
    actions = actions.cpu().numpy().reshape((-1, self.ctrl_action_dim))
    low, high = self._ctrl_space.low, self._ctrl_space.high
    actions = low + 0.5 * (actions + 1.0) * (high - low)  # [-1,1] -> ctrl bounds
    actions = np.clip(actions, low, high)
    if not vectorized:
      actions = actions.squeeze(axis=0)
    return actions, state

predict

predict(
    observation,
    state=None,
    episode_start=None,
    deterministic=False,
)

Deployment: return the CONTROL action only (the safety controller, no adversary), in the control sub-space bounds. Overrides SB3's predict, which would reshape the ctrl-dim output to the full concatenated action space.

Source code in safety_sb3/policies.py

def predict(self, observation, state=None, episode_start=None, deterministic=False):
  """Deployment: return the CONTROL action only (the safety controller, no
  adversary), in the control sub-space bounds.  Overrides SB3's predict, which
  would reshape the ctrl-dim output to the full concatenated action space."""
  self.set_training_mode(False)
  obs_tensor, vectorized = self.obs_to_tensor(observation)
  with __import__("torch").no_grad():
    actions = self.actor(obs_tensor, deterministic=deterministic)
  actions = actions.cpu().numpy().reshape((-1, self.ctrl_action_dim))
  low, high = self._ctrl_space.low, self._ctrl_space.high
  actions = low + 0.5 * (actions + 1.0) * (high - low)  # [-1,1] -> ctrl bounds
  actions = np.clip(actions, low, high)
  if not vectorized:
    actions = actions.squeeze(axis=0)
  return actions, state

Discount (gamma) annealing

safety_sb3.gamma_anneal

Discount-factor (gamma) annealing for the Safety Bellman backups.

WHY THIS EXISTS

The reach-avoid / avoid Safety Bellman operator

V = (1 - gamma) * min(l, g)  +  gamma * min(g, max(l, V'))

is a contraction with a sharp fixed point only in the limit gamma -> 1. At a fixed gamma = 0.99 the (1 - gamma) anchor keeps bleeding the immediate margin into the value everywhere, so the learned zero-level set (the safe/unsafe boundary we filter on) is smeared and the reach term never becomes crisp. Training the operator directly at gamma -> 1 is ill-conditioned (the contraction modulus -> 1, bootstrapping dominates, the critic diverges early). The reach-avoid literature (Fisac 2019; Hsu 2021; the ISAACS safe_adaptation_dev codebase) fixes this with DISCOUNT ANNEALING: start contractive at gamma = 0.99 for a well-conditioned early fixed point, then anneal gamma -> 1 so the boundary sharpens as the critic stabilises.

TWO SCHEDULERS (both provided; StepGammaAnneal is the DEFAULT)

• :class:StepGammaAnneal — REFERENCE-FAITHFUL discrete jumps (StepLRMargin in safe_adaptation_dev/agent/scheduler.py): the gap (goal - gamma) is multiplied by ratio at fixed training-fraction milestones. Default 0.99 -> 0.999 (at 20%) -> 0.9999 (at 40%), then hold. This is is_stepwise so each jump triggers an ENTROPY-TEMPERATURE (alpha) RESET in the SAC learners (the Q-scale shifts discontinuously at a jump, so the tuned alpha is stale — the reference resets both actors' alpha on every gamma jump; see GammaAnnealMixin._on_gamma_jump / the SAC family).

• :class:GeometricGammaAnneal — a smooth (continuous) log-space interpolation reaching end at anneal_frac then holding. No discontinuity, so no alpha reset. Available via gamma_anneal=GeometricGammaAnneal(...).

ON by default in every learner; gamma_anneal=False disables it, a callable frac -> gamma (optionally exposing .is_stepwise) is a custom schedule. All knobs (start/end/ratio/period/anneal fraction) are constructor arguments. Gamma is logged as train/gamma (+ train/gamma_jump on a step).

StepGammaAnneal

Discrete-jump gamma schedule (reference StepLRMargin).

At training fraction f, the number of elapsed jumps is n = floor(f / period_frac) and

gamma(f) = goal - (goal - init) * ratio**n     (clamped to <= end)

Defaults (init 0.99, ratio 0.1, period_frac 0.20, end 0.9999) reproduce the reference progression the user specified: 0.99 -> 0.999 @20% -> 0.9999 @40%, then hold. is_stepwise = True so the SAC learners reset alpha on each jump.

Parameters:

Name	Type	Description	Default
init	float	starting gamma (default 0.99).	0.99
end	float	final gamma held after the last jump (default 0.9999).	0.9999
ratio	float	multiplicative decay of the gap (goal - gamma) per jump (default 0.1 -> one extra nine per jump).	0.1
period_frac	float	training fraction between jumps (default 0.20; the user's "a jump every 10-20% of allotted time").	0.2
goal	float	the limit the gap closes toward — 1.0 for HJ reachability.	1.0

Source code in safety_sb3/gamma_anneal.py

class StepGammaAnneal:
    """Discrete-jump gamma schedule (reference ``StepLRMargin``).

    At training fraction ``f``, the number of elapsed jumps is
    ``n = floor(f / period_frac)`` and

        gamma(f) = goal - (goal - init) * ratio**n     (clamped to <= end)

    Defaults (init 0.99, ratio 0.1, period_frac 0.20, end 0.9999) reproduce the
    reference progression the user specified: 0.99 -> 0.999 @20% -> 0.9999 @40%,
    then hold.  ``is_stepwise = True`` so the SAC learners reset alpha on each jump.

    :param init: starting gamma (default 0.99).
    :param end: final gamma held after the last jump (default 0.9999).
    :param ratio: multiplicative decay of the gap ``(goal - gamma)`` per jump
        (default 0.1 -> one extra nine per jump).
    :param period_frac: training fraction between jumps (default 0.20; the user's
        "a jump every 10-20% of allotted time").
    :param goal: the limit the gap closes toward — 1.0 for HJ reachability.
    """

    is_stepwise = True

    def __init__(self, init: float = 0.99, end: float = 0.9999, ratio: float = 0.1,
                 period_frac: float = 0.20, goal: float = 1.0):
        init, end, ratio, period_frac, goal = (
            float(init), float(end), float(ratio), float(period_frac), float(goal))
        if not 0.0 < period_frac <= 1.0:
            raise ValueError(f"period_frac must be in (0, 1], got {period_frac}")
        if not 0.0 < ratio < 1.0:
            raise ValueError(f"ratio must be in (0, 1), got {ratio}")
        if not init < goal or not end < goal:
            raise ValueError(f"init and end must be < goal={goal}; got {init}, {end}")
        self.init, self.end, self.ratio = init, end, ratio
        self.period_frac, self.goal = period_frac, goal
        self._gap0 = goal - init
        self._degenerate = init >= end

    def __call__(self, frac: float) -> float:
        f = min(max(float(frac), 0.0), 1.0)
        if self._degenerate:
            return self.init
        n = int(f / self.period_frac)              # elapsed jumps
        gamma = self.goal - self._gap0 * (self.ratio ** n)
        return min(gamma, self.end)                # clamp at the final value

    def __repr__(self) -> str:
        return (f"StepGammaAnneal(init={self.init}, end={self.end}, "
                f"ratio={self.ratio}, period_frac={self.period_frac})")

GeometricGammaAnneal

Smooth (continuous) gamma schedule — log-space interpolation of the gap.

gamma(0) = init; gamma(anneal_frac) = end; constant end after. Equal fractions of progress multiply (goal - gamma) by a constant factor. is_stepwise = False -> no alpha reset (the change is continuous).

Parameters:

Name	Type	Description	Default
init	float	gamma at fraction 0 (default 0.99).	0.99
end	float	gamma reached at anneal_frac and held after (default 0.9999).	0.9999
goal	float	the limit the gap closes toward, 1.0 for HJ reachability.	1.0
anneal_frac	float	training fraction at which end is reached (default 0.5).	0.5

Source code in safety_sb3/gamma_anneal.py

class GeometricGammaAnneal:
    """Smooth (continuous) gamma schedule — log-space interpolation of the gap.

    ``gamma(0) = init``; ``gamma(anneal_frac) = end``; constant ``end`` after.
    Equal fractions of progress multiply ``(goal - gamma)`` by a constant factor.
    ``is_stepwise = False`` -> no alpha reset (the change is continuous).

    :param init: gamma at fraction 0 (default 0.99).
    :param end: gamma reached at ``anneal_frac`` and held after (default 0.9999).
    :param goal: the limit the gap closes toward, ``1.0`` for HJ reachability.
    :param anneal_frac: training fraction at which ``end`` is reached (default 0.5).
    """

    is_stepwise = False

    def __init__(self, init: float = 0.99, end: float = 0.9999,
                 goal: float = 1.0, anneal_frac: float = 0.5):
        init, end, goal, anneal_frac = float(init), float(end), float(goal), float(anneal_frac)
        if not 0.0 < anneal_frac <= 1.0:
            raise ValueError(f"anneal_frac must be in (0, 1], got {anneal_frac}")
        if not init < goal or not end < goal:
            raise ValueError(f"init and end must be < goal={goal}, got init={init}, end={end}")
        self.init, self.end, self.goal, self.anneal_frac = init, end, goal, anneal_frac
        self._gap0 = goal - init   # initial gap, e.g. 0.01
        self._gap1 = goal - end    # final gap,   e.g. 0.0001
        self._degenerate = init >= end

    def __call__(self, frac: float) -> float:
        f = min(max(float(frac), 0.0), 1.0)
        if self._degenerate:
            return self.init
        if f >= self.anneal_frac:
            return self.end
        t = f / self.anneal_frac                       # 0 .. 1 across the ramp
        gap = self._gap0 * (self._gap1 / self._gap0) ** t   # geometric in log space
        return self.goal - gap

    def __repr__(self) -> str:
        return (f"GeometricGammaAnneal(init={self.init}, end={self.end}, "
                f"goal={self.goal}, anneal_frac={self.anneal_frac})")

GammaAnnealMixin

Mix-in that anneals self.gamma over training for every learner.

Wiring:

• _update_current_progress_remaining re-applies the schedule each iteration -- covers every NUMPY / on-policy path (SB3 calls it in both learn loops).
• The GPU-resident collect loops bypass that call, so the PPO and SAC families (1P and 2P alike) also call _apply_gamma_anneal() at the top of their collect_rollouts. _apply is idempotent.

Gamma is applied at consumption time: rollout_buffer.gamma for PPO GAE, self.gamma for the SAC/DQN TD target. On a DISCRETE jump (StepGammaAnneal) _on_gamma_jump fires -- a no-op here (PPO), overridden by the SAC learners to reset the entropy temperature(s), mirroring the reference.

Source code in safety_sb3/gamma_anneal.py

class GammaAnnealMixin:
    """Mix-in that anneals ``self.gamma`` over training for every learner.

    Wiring:

    * ``_update_current_progress_remaining`` re-applies the schedule each
      iteration -- covers every NUMPY / on-policy path (SB3 calls it in both
      learn loops).
    * The GPU-resident collect loops bypass that call, so the PPO and SAC
      families (1P and 2P alike) also call ``_apply_gamma_anneal()``
      at the top of their ``collect_rollouts``.  ``_apply`` is idempotent.

    Gamma is applied at *consumption* time: ``rollout_buffer.gamma`` for PPO GAE,
    ``self.gamma`` for the SAC/DQN TD target.  On a DISCRETE jump (``StepGammaAnneal``)
    ``_on_gamma_jump`` fires -- a no-op here (PPO), overridden by the SAC learners
    to reset the entropy temperature(s), mirroring the reference.
    """

    _gamma_schedule: GammaSchedule | None = None
    _gamma_last: float | None = None

    def _setup_gamma_anneal(self, gamma_anneal) -> None:
        """Resolve the ``gamma_anneal`` constructor arg to a schedule (or None).

        ``True`` (default) -> the default DISCRETE-JUMP schedule seeded from the
        current ``self.gamma``; ``False``/``None`` -> disabled (constant gamma); a
        callable ``frac -> gamma`` -> used verbatim.
        """
        if gamma_anneal is False or gamma_anneal is None:
            self._gamma_schedule = None
        elif gamma_anneal is True:
            self._gamma_schedule = make_default_gamma_schedule(init=float(self.gamma))
        elif callable(gamma_anneal):
            self._gamma_schedule = gamma_anneal
        else:
            raise TypeError(
                "gamma_anneal must be a bool or a callable(frac)->gamma, got "
                f"{type(gamma_anneal).__name__}")
        self._gamma_last = None

    def _apply_gamma_anneal(self) -> None:
        """Recompute gamma from the current training fraction, propagate it, and
        fire ``_on_gamma_jump`` when a stepwise schedule crosses a jump."""
        sched = getattr(self, "_gamma_schedule", None)
        if sched is None:
            return
        total = getattr(self, "_total_timesteps", 0) or 0
        frac = 0.0 if total <= 0 else min(1.0, max(0.0, self.num_timesteps / float(total)))
        g = float(sched(frac))
        old = self._gamma_last
        self.gamma = g
        # On-policy: GAE reads gamma off the buffer(s). Two-player PPO keeps a
        # second (dstb) buffer -- update whichever exist.
        for name in ("rollout_buffer", "dstb_rollout_buffer"):
            buf = getattr(self, name, None)
            if buf is not None:
                buf.gamma = g
        logger = getattr(self, "logger", None)
        if logger is not None:
            logger.record("train/gamma", g)
        # Discrete jump -> alpha reset hook (stepwise schedules only).
        if (getattr(sched, "is_stepwise", False) and old is not None
                and abs(g - old) > 1e-12):
            if logger is not None:
                logger.record("train/gamma_jump", g)
            self._on_gamma_jump(old, g)
        self._gamma_last = g

    def _on_gamma_jump(self, old_gamma: float, new_gamma: float) -> None:
        """Hook fired when a DISCRETE (stepwise) gamma jump occurs. No-op for
        PPO/DQN; SAC learners override to reset the entropy temperature(s)."""

    def _update_current_progress_remaining(self, num_timesteps: int,
                                           total_timesteps: int) -> None:
        super()._update_current_progress_remaining(num_timesteps, total_timesteps)
        self._apply_gamma_anneal()

make_default_gamma_schedule

make_default_gamma_schedule(init=0.99)

The default on-by-default schedule: REFERENCE-FAITHFUL discrete jumps (init -> 0.999 @20% -> 0.9999 @40%, hold), alpha reset on each jump.

Source code in safety_sb3/gamma_anneal.py

def make_default_gamma_schedule(init: float = 0.99) -> StepGammaAnneal:
    """The default on-by-default schedule: REFERENCE-FAITHFUL discrete jumps
    (``init`` -> 0.999 @20% -> 0.9999 @40%, hold), alpha reset on each jump."""
    return StepGammaAnneal(init=init, end=0.9999, ratio=0.1, period_frac=0.20)

Environments and normalization

The GPU-resident tensor path. Subclass TensorVecEnv and implement step_tensor; wrap in TensorVecNormalize for on-device observation normalization. See the environment contract.

safety_sb3.tensor_env

GPU-resident vectorized-env contract for massive-parallel safety training.

Standard SB3 VecEnv is numpy end-to-end: with thousands of GPU-simulated envs (mjlab / Isaac) every step pays a device->host->device bounce for obs, rewards and infos, and the rollout buffer re-uploads at train time. The tensor path keeps EVERYTHING on device:

env.step_tensor(actions) -> (obs, reward_g, dones, timeouts, l_x)

all torch.Tensor on env.device. The learners detect is_tensor_env and switch to a torch-native collect_rollouts + :mod:safety_sb3.tensor_buffers (identical backup math, no numpy anywhere on the hot path).

Reward semantics are unchanged: the reward IS the safety margin g(s); l_x is the target margin for reach-avoid (zeros for avoid-only envs). Reward normalization does not exist on this path by design (it would corrupt the Safety Bellman backup); observation normalization is provided by :class:TensorVecNormalize (an rsl_rl-style running normalizer on device).

Subclass :class:TensorVecEnv for your simulator bridge; it also subclasses SB3's VecEnv so BaseAlgorithm accepts it (the numpy step API raises).

TensorVecEnv

Bases: VecEnv

Torch-native VecEnv. Implement reset and step_tensor.

Source code in safety_sb3/tensor_env.py

class TensorVecEnv(VecEnv):
  """Torch-native VecEnv. Implement ``reset`` and ``step_tensor``."""

  is_tensor_env = True
  render_mode = None

  def __init__(self, num_envs: int, observation_space: spaces.Space,
               action_space: spaces.Space, device: str = "cuda:0"):
    super().__init__(num_envs, observation_space, action_space)
    self.device = device

  # --- the tensor API ------------------------------------------------------
  def reset(self) -> th.Tensor:  # type: ignore[override]
    raise NotImplementedError

  def step_tensor(
    self, actions: th.Tensor
  ) -> tuple[th.Tensor, th.Tensor, th.Tensor, th.Tensor, th.Tensor]:
    """-> (obs, reward_g, dones, timeouts, l_x); all on ``self.device``.

    ``dones`` = terminated | truncated (env auto-resets internally);
    ``timeouts`` = truncated & ~terminated; ``l_x`` = target margin (zeros
    for avoid-only envs)."""
    raise NotImplementedError

  def metrics(self) -> dict[str, float]:
    """Optional per-rollout scalars (curriculum levels, task metrics) that the
    algorithm forwards into its logger. Never train on these."""
    return {}

  # --- numpy VecEnv API: not supported on the tensor path ------------------
  def step_async(self, actions):
    raise RuntimeError(
      "TensorVecEnv has no numpy step API — pair it with a safety_sb3 "
      "algorithm (the learners detect is_tensor_env)."
    )

  def step_wait(self):
    raise RuntimeError("TensorVecEnv: use step_tensor().")

  def close(self) -> None:
    pass

  def _indices(self, indices):
    if indices is None:
      return range(self.num_envs)
    return [indices] if isinstance(indices, int) else indices

  def get_attr(self, attr_name, indices=None):
    return [getattr(self, attr_name, None) for _ in self._indices(indices)]

  def set_attr(self, attr_name, value, indices=None):
    setattr(self, attr_name, value)

  def env_method(self, method_name, *args, indices=None, **kwargs):
    return [None for _ in self._indices(indices)]

  def env_is_wrapped(self, wrapper_class, indices=None):
    return [False for _ in self._indices(indices)]

step_tensor

step_tensor(actions)

-> (obs, reward_g, dones, timeouts, l_x); all on self.device.

dones = terminated | truncated (env auto-resets internally); timeouts = truncated & ~terminated; l_x = target margin (zeros for avoid-only envs).

Source code in safety_sb3/tensor_env.py

def step_tensor(
  self, actions: th.Tensor
) -> tuple[th.Tensor, th.Tensor, th.Tensor, th.Tensor, th.Tensor]:
  """-> (obs, reward_g, dones, timeouts, l_x); all on ``self.device``.

  ``dones`` = terminated | truncated (env auto-resets internally);
  ``timeouts`` = truncated & ~terminated; ``l_x`` = target margin (zeros
  for avoid-only envs)."""
  raise NotImplementedError

metrics

metrics()

Optional per-rollout scalars (curriculum levels, task metrics) that the algorithm forwards into its logger. Never train on these.

Source code in safety_sb3/tensor_env.py

def metrics(self) -> dict[str, float]:
  """Optional per-rollout scalars (curriculum levels, task metrics) that the
  algorithm forwards into its logger. Never train on these."""
  return {}

TensorVecNormalize

Bases: TensorVecEnv

Running observation normalizer on device (rsl_rl-parity).

Wraps a :class:TensorVecEnv; normalizes observations with running mean/var updated during training (freeze with training=False). There is deliberately NO reward normalization on this path — the reward is the physical margin g(s).

Source code in safety_sb3/tensor_env.py

class TensorVecNormalize(TensorVecEnv):
  """Running observation normalizer on device (rsl_rl-parity).

  Wraps a :class:`TensorVecEnv`; normalizes observations with running
  mean/var updated during training (freeze with ``training=False``). There is
  deliberately NO reward normalization on this path — the reward is the
  physical margin g(s).
  """

  def __init__(self, venv: TensorVecEnv, epsilon: float = 1e-8,
               clip_obs: float = 10.0):
    # clip_obs MUST match SB3 VecNormalize (10.0): policies warm-started from
    # the numpy path were trained on +-10-clipped obs; a wider clip feeds them
    # 10x out-of-distribution spikes (raycast at gap edges) and silently
    # breaks the transferred behavior.
    super().__init__(venv.num_envs, venv.observation_space, venv.action_space,
                     venv.device)
    self.venv = venv
    self.training = True
    self.epsilon = float(epsilon)
    self.clip_obs = float(clip_obs)
    dim = int(np.prod(venv.observation_space.shape))
    dev = venv.device
    self.obs_mean = th.zeros(dim, device=dev)
    self.obs_var = th.ones(dim, device=dev)
    self.count = th.tensor(1e-4, device=dev)

  # --- normalization -------------------------------------------------------
  def _update(self, obs: th.Tensor) -> None:
    batch_mean = obs.mean(dim=0)
    batch_var = obs.var(dim=0, unbiased=False)
    batch_count = obs.shape[0]
    delta = batch_mean - self.obs_mean
    tot = self.count + batch_count
    self.obs_mean += delta * batch_count / tot
    m_a = self.obs_var * self.count
    m_b = batch_var * batch_count
    m2 = m_a + m_b + delta.square() * self.count * batch_count / tot
    self.obs_var = m2 / tot
    self.count = tot

  def normalize_obs(self, obs: th.Tensor) -> th.Tensor:
    return th.clamp(
      (obs - self.obs_mean) / th.sqrt(self.obs_var + self.epsilon),
      -self.clip_obs, self.clip_obs,
    )

  def normalize_obs_np(self, obs: np.ndarray) -> np.ndarray:
    """Numpy convenience for eval/video harnesses."""
    t = th.as_tensor(obs, dtype=th.float32, device=self.device)
    return self.normalize_obs(t).cpu().numpy()

  # --- env API --------------------------------------------------------------
  def reset(self) -> th.Tensor:
    obs = self.venv.reset()
    if self.training:
      self._update(obs)
    return self.normalize_obs(obs)

  def step_tensor(self, actions: th.Tensor):
    obs, r, dones, timeouts, l_x = self.venv.step_tensor(actions)
    if self.training:
      self._update(obs)
    return self.normalize_obs(obs), r, dones, timeouts, l_x

  def metrics(self) -> dict[str, float]:
    return self.venv.metrics()

  @property
  def executed_action(self):
    """What the WRAPPED env actually applied last step, or None.

    An env may execute something other than the action it was handed -- a
    training-time safety filter substituting its fallback is the case this
    exists for -- and the off-policy collector reads it back so the replay
    buffer holds the transition that really happened. This class has no
    ``__getattr__``, so without an explicit forward the readback would find
    nothing on the wrapper, silently keep the PROPOSED action, and the buffer
    would fill with actions that were never applied. That failure is invisible
    from the outside: training runs, the curves look plausible, and the critic
    is fit to transitions that never occurred.
    """
    return getattr(self.venv, "executed_action", None)

  def close(self) -> None:
    self.venv.close()

  # --- persistence ----------------------------------------------------------
  def save(self, path: str) -> None:
    th.save({"obs_mean": self.obs_mean, "obs_var": self.obs_var,
             "count": self.count}, path)

  @classmethod
  def load(cls, path: str, venv: TensorVecEnv,
           device: Optional[str] = None) -> "TensorVecNormalize":
    obj = cls(venv)
    state = th.load(path, map_location=device or venv.device, weights_only=True)
    obj.obs_mean = state["obs_mean"].to(venv.device)
    obj.obs_var = state["obs_var"].to(venv.device)
    obj.count = state["count"].to(venv.device)
    return obj

executed_action property

executed_action

What the WRAPPED env actually applied last step, or None.

An env may execute something other than the action it was handed -- a training-time safety filter substituting its fallback is the case this exists for -- and the off-policy collector reads it back so the replay buffer holds the transition that really happened. This class has no __getattr__, so without an explicit forward the readback would find nothing on the wrapper, silently keep the PROPOSED action, and the buffer would fill with actions that were never applied. That failure is invisible from the outside: training runs, the curves look plausible, and the critic is fit to transitions that never occurred.

normalize_obs_np

normalize_obs_np(obs)

Numpy convenience for eval/video harnesses.

Source code in safety_sb3/tensor_env.py

def normalize_obs_np(self, obs: np.ndarray) -> np.ndarray:
  """Numpy convenience for eval/video harnesses."""
  t = th.as_tensor(obs, dtype=th.float32, device=self.device)
  return self.normalize_obs(t).cpu().numpy()

Callbacks

safety_sb3.eval_callbacks

Periodic SAFE-RATE / SUCCESS-RATE evaluation for the tensor training stack.

A standalone SB3 :class:~stable_baselines3.common.callbacks.BaseCallback that, every eval_freq env-steps, rolls out the current policy on a GPU-resident eval env and records three scalars to the SB3 logger (which syncs to wandb):

• eval/safe_rate — fraction of episodes that NEVER entered the failure set (never saw g < 0).
• eval/success_rate — for reach-avoid tasks: fraction that reached the target (ever l_x >= 0) AND stayed safe; for avoid-only tasks it equals safe_rate (the reach term does not apply).
• eval/ep_len_mean — mean length of the tracked first episodes.

The definitions mirror the reference audit target safe_adaptation_dev (utils/eval.py evaluate_zero_sum: safe_rate = mean(results != -1), ra_rate = mean(results == 1); result codes success=1 / failure=-1 / timeout=0) and the league eval in :mod:safety_sb3.leaderboard (:class:~safety_sb3.leaderboard.LeagueEvaluator: safe = never g < 0, reached = ever l_x >= 0).

Rollout accounting is parallel: the eval env runs num_envs episodes at once, we track each env's FIRST episode (safe until its first g < 0, reached if it ever hit l_x >= 0 before that env's first done), and accumulate ceil(n_rollouts / num_envs) batches so at least n_rollouts first-episodes are scored. Everything stays on device — the eval env speaks the tensor API obs, g, dones, timeouts, l_x = env.step_tensor(actions).

SafeSuccessRateEvalCallback

Bases: BaseCallback

Periodic safe-rate / success-rate eval on a :class:TensorVecEnv.

Parameters:

Name	Type	Description	Default
eval_env		a GPU-resident tensor env (safety_sb3.tensor_env TensorVecEnv, optionally wrapped in TensorVecNormalize). Must expose num_envs, reset() -> obs and step_tensor(actions) -> (obs, g, dones, timeouts, l_x) (all torch tensors on eval_env.device). Should be SEPARATE from the training env.	required
n_rollouts	int	minimum number of first-episodes to score per eval. The callback runs ceil(n_rollouts / num_envs) parallel batches.	100
eval_freq	int	run an eval every eval_freq env-steps (compared against self.num_timesteps).	1000000
reach_avoid	bool	True for a reach-avoid task (success = reached AND safe); False for avoid-only (success == safe_rate).	True
max_ep_steps	int	hard cap on rollout length (guards against episodes that never terminate). A capped, never-violated episode counts as safe.	400
deterministic	bool	use deterministic policy actions during eval.	True

Source code in safety_sb3/eval_callbacks.py

class SafeSuccessRateEvalCallback(BaseCallback):
  """Periodic safe-rate / success-rate eval on a :class:`TensorVecEnv`.

  :param eval_env: a GPU-resident tensor env (``safety_sb3.tensor_env``
    ``TensorVecEnv``, optionally wrapped in ``TensorVecNormalize``). Must expose
    ``num_envs``, ``reset() -> obs`` and
    ``step_tensor(actions) -> (obs, g, dones, timeouts, l_x)`` (all torch
    tensors on ``eval_env.device``). Should be SEPARATE from the training env.
  :param n_rollouts: minimum number of first-episodes to score per eval. The
    callback runs ``ceil(n_rollouts / num_envs)`` parallel batches.
  :param eval_freq: run an eval every ``eval_freq`` env-steps (compared against
    ``self.num_timesteps``).
  :param reach_avoid: ``True`` for a reach-avoid task (success = reached AND
    safe); ``False`` for avoid-only (success == safe_rate).
  :param max_ep_steps: hard cap on rollout length (guards against episodes that
    never terminate). A capped, never-violated episode counts as safe.
  :param deterministic: use deterministic policy actions during eval.
  """

  def __init__(
    self,
    eval_env,
    n_rollouts: int = 100,
    eval_freq: int = 1_000_000,
    reach_avoid: bool = True,
    max_ep_steps: int = 400,
    deterministic: bool = True,
    verbose: int = 0,
  ) -> None:
    super().__init__(verbose)
    self.eval_env = eval_env
    self.n_rollouts = int(n_rollouts)
    self.eval_freq = int(eval_freq)
    self.reach_avoid = bool(reach_avoid)
    self.max_ep_steps = int(max_ep_steps)
    self.deterministic = bool(deterministic)
    # Fire the first eval once `eval_freq` steps have elapsed.
    self._next_eval = self.eval_freq

  # --- SB3 hook ------------------------------------------------------------
  def _on_step(self) -> bool:
    if self.eval_freq > 0 and self.num_timesteps >= self._next_eval:
      # Advance past any missed windows so a batched step-increment (tensor
      # path adds `num_envs` per step) does not re-fire immediately.
      while self.num_timesteps >= self._next_eval:
        self._next_eval += self.eval_freq
      self._safe_eval()
    return True

  def _safe_eval(self) -> None:
    """Run one eval, never letting a bad env/policy shape crash training."""
    try:
      safe_rate, success_rate, ep_len_mean, n = self._run_eval()
    except Exception as exc:  # noqa: BLE001 -- eval must never kill training
      warnings.warn(
        f"[SafeSuccessRateEvalCallback] eval at step {self.num_timesteps} "
        f"skipped: {type(exc).__name__}: {exc}"
      )
      return
    self.logger.record("eval/safe_rate", safe_rate)
    self.logger.record("eval/success_rate", success_rate)
    self.logger.record("eval/ep_len_mean", ep_len_mean)
    if self.verbose:
      tag = "reach-avoid" if self.reach_avoid else "avoid-only"
      print(
        f"[eval @ {self.num_timesteps}] ({tag}, n={n}) "
        f"safe_rate={safe_rate:.3f} success_rate={success_rate:.3f} "
        f"ep_len_mean={ep_len_mean:.1f}"
      )

  # --- rollout accounting --------------------------------------------------
  @th.no_grad()
  def _run_eval(self):
    env = self.eval_env
    n_env = int(getattr(env, "num_envs"))
    device = getattr(env, "device", self.model.device)
    policy = self.model.policy
    two_player = hasattr(policy, "dstb_actor")
    n_batches = max(1, math.ceil(self.n_rollouts / max(n_env, 1)))

    # Freeze the observation normalizer (if any) and switch the policy to eval
    # mode; both are restored in `finally`.
    prev_training = getattr(env, "training", None)
    if prev_training is not None:
      env.training = False
    prev_policy_training = getattr(policy, "training", None)
    policy.set_training_mode(False)

    safe_hits = 0
    succ_hits = 0
    len_sum = 0.0
    total = 0
    try:
      for _ in range(n_batches):
        ep_safe, ep_reached, ep_len = self._rollout_batch(
          env, policy, two_player, n_env, device
        )
        succ = ep_safe & ep_reached if self.reach_avoid else ep_safe
        safe_hits += int(ep_safe.sum().item())
        succ_hits += int(succ.sum().item())
        len_sum += float(ep_len.sum().item())
        total += n_env
    finally:
      if prev_training is not None:
        env.training = prev_training
      if prev_policy_training:
        policy.set_training_mode(True)

    denom = max(total, 1)
    return safe_hits / denom, succ_hits / denom, len_sum / denom, total

  def _rollout_batch(self, env, policy, two_player, n_env, device):
    """One parallel batch of `num_envs` first-episodes.

    Returns per-env (ep_safe, ep_reached, ep_len) for THIS env's first episode;
    an env stops being tracked at its first `done` (the env auto-resets, but we
    ignore everything after the first termination)."""
    obs = env.reset()
    ep_safe = th.ones(n_env, dtype=th.bool, device=device)
    ep_reached = th.zeros(n_env, dtype=th.bool, device=device)
    done_once = th.zeros(n_env, dtype=th.bool, device=device)
    ep_len = th.zeros(n_env, dtype=th.float32, device=device)

    for _ in range(self.max_ep_steps):
      actions = self._policy_actions(policy, obs, two_player)
      obs, g, dones, _timeouts, l_x = env.step_tensor(actions)
      active = ~done_once  # envs still on their first episode this step
      g = th.as_tensor(g, device=device).reshape(n_env)
      l_x = th.as_tensor(l_x, device=device).reshape(n_env)
      dones = th.as_tensor(dones, device=device).reshape(n_env).bool()
      ep_safe &= ~(active & (g < 0))
      ep_reached |= active & (l_x >= 0)
      ep_len += active.float()  # counts the terminating step too
      done_once |= active & dones
      if bool(done_once.all()):
        break
    return ep_safe, ep_reached, ep_len

  def _policy_actions(self, policy, obs, two_player):
    """Deterministic policy action for the tensor env.

    Single-player: the control actor's output. Two-player (has ``dstb_actor``):
    the control + disturbance actors composed into the env's
    ``ctrl_dim + dstb_dim`` action, mirroring
    ``sac_2p.AbstractSAC2P._tensor_policy_actions``.
    Actors output in ``[-1, 1]`` and the tensor env's action space is
    ``[-1, 1]`` (the env clamps), so no unscaling is needed on this path. The
    eval env already returns normalized observations when it is a
    ``TensorVecNormalize`` wrapper, so ``obs`` is fed to the actor as-is."""
    ctrl = policy.actor(obs, deterministic=self.deterministic)
    if not two_player:
      return ctrl
    dstb = policy.dstb_actor(obs, deterministic=self.deterministic)
    return th.cat([ctrl, dstb], dim=1)

safety_sb3.callbacks

Callbacks for stable safety-policy fine-tuning.

Empirically, on-policy training against margin-only objectives (safety or reach-avoid backups) carries no gait-/action-quality gradient, so the usual PPO entropy dynamics can inflate the action std until a converged motor skill erodes. Capping the std is a one-line remedy that preserved fine-tuned capability where uncapped runs lost it. Pair with SB3's native target_kl early-stop for the actor epochs.

StdCapCallback

Bases: BaseCallback

Clamp the policy's action std at the start of every rollout.

Works with any SB3 on-policy algorithm whose policy exposes a state- independent log_std parameter (the default MlpPolicy for Box action spaces). Policies without log_std are silently ignored.

Parameters:

Name	Type	Description	Default
max_std	float	hard upper bound on the (post-clamp) action std.	required

Source code in safety_sb3/callbacks.py

class StdCapCallback(BaseCallback):
  """Clamp the policy's action std at the start of every rollout.

  Works with any SB3 on-policy algorithm whose policy exposes a state-
  independent ``log_std`` parameter (the default ``MlpPolicy`` for Box
  action spaces). Policies without ``log_std`` are silently ignored.

  :param max_std: hard upper bound on the (post-clamp) action std.
  """

  def __init__(self, max_std: float, verbose: int = 0):
    super().__init__(verbose)
    if max_std <= 0:
      raise ValueError(f"max_std must be positive, got {max_std}")
    self.max_log_std = math.log(max_std)

  def _on_rollout_start(self) -> None:
    if hasattr(self.model.policy, "log_std"):
      with th.no_grad():
        self.model.policy.log_std.clamp_(max=self.max_log_std)

  def _on_step(self) -> bool:
    return True

Buffers

safety_sb3.buffers_rollout

On-policy rollout buffers — one per Mode, player-agnostic.

Rollout buffers carry the M of MAP and nothing else. A buffer does not know or care whether one player or two produced the data in it, so there is no 1P/2P in these names: :class:~safety_sb3.ppo_2p.AbstractPPO2P builds two instances of the very same class, one per player.

Each buffer computes the backup for its _MODE through :func:safety_sb3.backups.target, so the operator is a parameter rather than a class body; mode= overrides the class default per instance.

Each buffer also captures its own extras. The safety margin g(s) always rides on the reward channel, but the reach-avoid operator additionally needs the target margin l(s), which SB3 delivers only in infos. Rather than teach every algorithm to thread l, the collect loop hands each step's infos to the buffer via :meth:record_extras and the buffer keeps whatever its own operator needs — nothing, for the two modes that have no l. That is why the reach-avoid learners need almost no reach-avoid code.

SafetyRolloutBuffer

Bases: RolloutBuffer

Rollout buffer for the avoid backup (Fisac et al. 2019).

g(s) rides on the reward channel; the operator is V(s) = (1-gamma)*g + gamma*min(g, V(s')) with terminal target g. See :mod:safety_sb3.backups, which defines it.

Source code in safety_sb3/buffers_rollout.py

class SafetyRolloutBuffer(RolloutBuffer):
    """Rollout buffer for the **avoid** backup (Fisac et al. 2019).

    ``g(s)`` rides on the reward channel; the operator is
    ``V(s) = (1-gamma)*g + gamma*min(g, V(s'))`` with terminal target ``g``.
    See :mod:`safety_sb3.backups`, which defines it.
    """

    _MODE = backups.AVOID

    def __init__(self, *args, mode: str | None = None, **kwargs):
        super().__init__(*args, **kwargs)
        self._MODE = backups.check_mode(self._MODE if mode is None else mode)

    def record_extras(self, infos: list) -> None:
        """Capture per-step extras from ``infos`` into the slot ``add()`` will fill.

        Called by the on-policy collect loop immediately before ``add()``.
        No-op here: the avoid operator needs only ``g``, which arrives on the
        reward channel. :class:`ReachAvoidRolloutBuffer` overrides it to keep
        ``l(s)``.
        """

    def _target(self, step: int, v_next: np.ndarray,
                not_done: np.ndarray) -> np.ndarray:
        l_x = getattr(self, "l_x", None)
        return backups.target(
            self._MODE, self.rewards[step], v_next, not_done, self.gamma,
            l=None if l_x is None else l_x[step],
            terminal_type=getattr(self, "terminal_type", "all"))

    def compute_returns_and_advantage(self, last_values: th.Tensor, dones: np.ndarray) -> None:
        """SB3's GAE loop with this buffer's Bellman backup substituted for the
        one-step TD target inside ``delta`` (see :meth:`_target`)."""
        # Convert to numpy
        last_values = last_values.clone().cpu().numpy().flatten()  # type: ignore[assignment]

        last_gae_lam = 0
        for step in reversed(range(self.buffer_size)):
            if step == self.buffer_size - 1:
                next_non_terminal = 1.0 - dones.astype(np.float32)
                v_next = last_values
            else:
                next_non_terminal = 1.0 - self.episode_starts[step + 1]
                v_next = self.values[step + 1]

            target = self._target(step, v_next, next_non_terminal)
            delta = target - self.values[step]
            last_gae_lam = delta + self.gamma * self.gae_lambda * next_non_terminal * last_gae_lam
            self.advantages[step] = last_gae_lam

        self.returns = self.advantages + self.values

record_extras

record_extras(infos)

Capture per-step extras from infos into the slot add() will fill.

Called by the on-policy collect loop immediately before add(). No-op here: the avoid operator needs only g, which arrives on the reward channel. :class:ReachAvoidRolloutBuffer overrides it to keep l(s).

Source code in safety_sb3/buffers_rollout.py

def record_extras(self, infos: list) -> None:
    """Capture per-step extras from ``infos`` into the slot ``add()`` will fill.

    Called by the on-policy collect loop immediately before ``add()``.
    No-op here: the avoid operator needs only ``g``, which arrives on the
    reward channel. :class:`ReachAvoidRolloutBuffer` overrides it to keep
    ``l(s)``.
    """

compute_returns_and_advantage

compute_returns_and_advantage(last_values, dones)

SB3's GAE loop with this buffer's Bellman backup substituted for the one-step TD target inside delta (see :meth:_target).

Source code in safety_sb3/buffers_rollout.py

def compute_returns_and_advantage(self, last_values: th.Tensor, dones: np.ndarray) -> None:
    """SB3's GAE loop with this buffer's Bellman backup substituted for the
    one-step TD target inside ``delta`` (see :meth:`_target`)."""
    # Convert to numpy
    last_values = last_values.clone().cpu().numpy().flatten()  # type: ignore[assignment]

    last_gae_lam = 0
    for step in reversed(range(self.buffer_size)):
        if step == self.buffer_size - 1:
            next_non_terminal = 1.0 - dones.astype(np.float32)
            v_next = last_values
        else:
            next_non_terminal = 1.0 - self.episode_starts[step + 1]
            v_next = self.values[step + 1]

        target = self._target(step, v_next, next_non_terminal)
        delta = target - self.values[step]
        last_gae_lam = delta + self.gamma * self.gae_lambda * next_non_terminal * last_gae_lam
        self.advantages[step] = last_gae_lam

    self.returns = self.advantages + self.values

ReachAvoidRolloutBuffer

Bases: SafetyRolloutBuffer

Rollout buffer for the reach-avoid backup (Hsu et al. RSS'21 eq. 15 / Gameplay Filters eq. 6a)::

V(s) = (1-gamma)*min(l, g) + gamma*min(g, max(l, V(s')))

Stores the per-step target margin l(s) beside the safety margin g(s) (which rides on rewards). l is read out of info["l_x"] by :meth:record_extras — the buffer captures it itself, which is what keeps the reach-avoid learners free of l-plumbing.

See :mod:safety_sb3.backups for why the anchor is min(l, g) and not g, and why avoid is not expressible by degenerating l.

Parameters:

Name	Type	Description	Default
terminal_type	str	"all" (default) -> terminal target min(l, g); "g" -> g. See :func:safety_sb3.backups.reach_avoid_target.	'all'

Source code in safety_sb3/buffers_rollout.py

class ReachAvoidRolloutBuffer(SafetyRolloutBuffer):
    """Rollout buffer for the **reach-avoid** backup (Hsu et al. RSS'21 eq. 15 /
    Gameplay Filters eq. 6a)::

        V(s) = (1-gamma)*min(l, g) + gamma*min(g, max(l, V(s')))

    Stores the per-step target margin ``l(s)`` beside the safety margin ``g(s)``
    (which rides on ``rewards``). ``l`` is read out of ``info["l_x"]`` by
    :meth:`record_extras` — the buffer captures it itself, which is what keeps
    the reach-avoid learners free of ``l``-plumbing.

    See :mod:`safety_sb3.backups` for why the anchor is ``min(l, g)`` and not
    ``g``, and why avoid is not expressible by degenerating ``l``.

    :param terminal_type: ``"all"`` (default) -> terminal target ``min(l, g)``;
        ``"g"`` -> ``g``. See :func:`safety_sb3.backups.reach_avoid_target`.
    """

    _MODE = backups.REACH_AVOID

    def __init__(self, *args, terminal_type: str = "all", **kwargs):
        super().__init__(*args, **kwargs)
        self.terminal_type = backups.check_terminal_type(terminal_type)

    def reset(self) -> None:
        super().reset()
        self.l_x = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)

    def record_extras(self, infos: list) -> None:
        """Keep this step's target margin ``l(s)``, read from ``info["l_x"]``."""
        self.l_x[self.pos] = np.array(
            [float(info.get("l_x", 0.0)) for info in infos], dtype=np.float32)

record_extras

record_extras(infos)

Keep this step's target margin l(s), read from info["l_x"].

Source code in safety_sb3/buffers_rollout.py

def record_extras(self, infos: list) -> None:
    """Keep this step's target margin ``l(s)``, read from ``info["l_x"]``."""
    self.l_x[self.pos] = np.array(
        [float(info.get("l_x", 0.0)) for info in infos], dtype=np.float32)

CumulativeRolloutBuffer

Bases: SafetyRolloutBuffer

Rollout buffer for the cumulative (ordinary RL) backup V(s) = r + gamma*V(s') — which makes the loop above exactly stock SB3 GAE.

Not a safety buffer: the reward channel carries a reward, not a margin, and the resulting value has no zero-level-set meaning. It exists so a reward-maximizing baseline runs through every line of the same code as the safety learners (see :mod:safety_sb3.backups).

Source code in safety_sb3/buffers_rollout.py

class CumulativeRolloutBuffer(SafetyRolloutBuffer):
    """Rollout buffer for the **cumulative** (ordinary RL) backup
    ``V(s) = r + gamma*V(s')`` — which makes the loop above exactly stock SB3 GAE.

    Not a safety buffer: the reward channel carries a *reward*, not a margin, and
    the resulting value has no zero-level-set meaning. It exists so a
    reward-maximizing baseline runs through every line of the same code as the
    safety learners (see :mod:`safety_sb3.backups`).
    """

    _MODE = backups.CUMULATIVE

rollout_buffer_classes

rollout_buffer_classes(mode)

(numpy_cls, tensor_cls) implementing mode's backup.

The on-policy learners look their buffers up here from their own _MODE instead of naming them as class attributes, so a buffer cannot drift out of step with the mode its owner declares. (It used to: the two-player avoid learner hardcoded the reach-avoid buffer for its min player and so trained that player on the wrong game.) It also makes every concrete learner in the library a genuine one-liner — _MODE = ... and nothing else.

Source code in safety_sb3/buffers_rollout.py

def rollout_buffer_classes(mode: str) -> tuple[type, type]:
    """``(numpy_cls, tensor_cls)`` implementing ``mode``'s backup.

    The on-policy learners look their buffers up here from their own ``_MODE``
    instead of naming them as class attributes, so a buffer cannot drift out of
    step with the mode its owner declares. (It used to: the two-player avoid
    learner hardcoded the reach-avoid buffer for its min player and so trained
    that player on the wrong game.) It also makes every concrete learner in the
    library a genuine one-liner — ``_MODE = ...`` and nothing else.
    """
    from .buffers_tensor import (TensorCumulativeRolloutBuffer,
                                 TensorReachAvoidRolloutBuffer,
                                 TensorSafetyRolloutBuffer)
    return {
        backups.AVOID: (SafetyRolloutBuffer, TensorSafetyRolloutBuffer),
        backups.REACH_AVOID: (ReachAvoidRolloutBuffer,
                              TensorReachAvoidRolloutBuffer),
        backups.CUMULATIVE: (CumulativeRolloutBuffer,
                             TensorCumulativeRolloutBuffer),
    }[backups.check_mode(mode)]

safety_sb3.buffers_tensor

GPU-resident rollout buffers — torch twins of :mod:safety_sb3.buffers_rollout.

Identical math, no numpy on the hot path. One class per Mode, player-agnostic (see the numpy module for why buffers carry no 1P/2P). Both families call the same operators from :mod:safety_sb3.backups; tests/test_backups.py asserts exact parity.

get() yields standard RolloutBufferSamples whose fields are device tensors, so stock PPO.train() consumes them unchanged. values / returns are exposed as numpy properties (PPO's explained-variance logging touches them once per update).

:meth:TensorSafetyRolloutBuffer.record_extras is the tensor twin of the numpy buffers' infos capture: on this path the env returns l_x directly from step_tensor, so the buffer receives the tensor rather than a list of dicts.

TensorSafetyRolloutBuffer

Torch rollout buffer with the AVOID (safety) Bellman backup.

Source code in safety_sb3/buffers_tensor.py

class TensorSafetyRolloutBuffer:
  """Torch rollout buffer with the AVOID (safety) Bellman backup."""

  is_tensor_buffer = True

  def __init__(self, buffer_size: int, observation_space: spaces.Space,
               action_space: spaces.Space, device: str = "cuda:0",
               gae_lambda: float = 0.95, gamma: float = 0.99,
               n_envs: int = 1, mode: str | None = None, **_ignored):
    self._MODE = backups.check_mode(self._MODE if mode is None else mode)
    self.buffer_size = int(buffer_size)
    self.n_envs = int(n_envs)
    self.device = device
    self.gamma = float(gamma)
    self.gae_lambda = float(gae_lambda)
    self.obs_dim = int(np.prod(observation_space.shape))
    self.act_dim = int(np.prod(action_space.shape))
    self.reset()

  def reset(self) -> None:
    T, N, dev = self.buffer_size, self.n_envs, self.device
    self.observations = th.zeros(T, N, self.obs_dim, device=dev)
    self.actions = th.zeros(T, N, self.act_dim, device=dev)
    self.rewards = th.zeros(T, N, device=dev)
    self.episode_starts = th.zeros(T, N, device=dev)
    self._values = th.zeros(T, N, device=dev)
    self.log_probs = th.zeros(T, N, device=dev)
    self.advantages = th.zeros(T, N, device=dev)
    self._returns = th.zeros(T, N, device=dev)
    self.pos = 0
    self.full = False

  # numpy views for PPO.train()'s explained-variance logging.
  @property
  def values(self) -> np.ndarray:
    return self._values.detach().cpu().numpy()

  @property
  def returns(self) -> np.ndarray:
    return self._returns.detach().cpu().numpy()

  def add(self, obs: th.Tensor, actions: th.Tensor, rewards: th.Tensor,
          episode_starts: th.Tensor, values: th.Tensor,
          log_probs: th.Tensor) -> None:
    p = self.pos
    self.observations[p] = obs.reshape(self.n_envs, self.obs_dim)
    self.actions[p] = actions.reshape(self.n_envs, self.act_dim)
    self.rewards[p] = rewards.reshape(self.n_envs)
    self.episode_starts[p] = episode_starts.reshape(self.n_envs).float()
    self._values[p] = values.reshape(self.n_envs)
    self.log_probs[p] = log_probs.reshape(self.n_envs)
    self.pos += 1
    if self.pos == self.buffer_size:
      self.full = True

  def record_extras(self, l_x: th.Tensor) -> None:
    """Capture per-step extras into the slot ``add()`` will fill.

    Tensor twin of the numpy buffers' ``record_extras(infos)``: the GPU-resident
    env hands back ``l_x`` from ``step_tensor`` directly, so there are no infos
    to parse. No-op here — the avoid operator needs only ``g``.
    """

  # --- backup ---------------------------------------------------------------
  #: backup this buffer computes; the ``mode=`` ctor kwarg overrides it
  _MODE = backups.AVOID

  def _target(self, step: int, v_next: th.Tensor,
              not_done: th.Tensor) -> th.Tensor:
    """Dispatch on ``self._MODE`` — twin of the numpy buffer's ``_target``."""
    l_x = getattr(self, "l_x", None)
    return backups.target(
      self._MODE, self.rewards[step], v_next, not_done, self.gamma,
      l=None if l_x is None else l_x[step],
      terminal_type=getattr(self, "terminal_type", "all"))

  @th.no_grad()
  def compute_returns_and_advantage(self, last_values: th.Tensor,
                                    dones: th.Tensor) -> None:
    last_values = last_values.reshape(self.n_envs).to(self.device)
    dones = dones.reshape(self.n_envs).float().to(self.device)
    last_gae = th.zeros(self.n_envs, device=self.device)
    for step in reversed(range(self.buffer_size)):
      if step == self.buffer_size - 1:
        next_non_terminal = 1.0 - dones
        v_next = last_values
      else:
        next_non_terminal = 1.0 - self.episode_starts[step + 1]
        v_next = self._values[step + 1]
      target = self._target(step, v_next, next_non_terminal)
      delta = target - self._values[step]
      last_gae = delta + self.gamma * self.gae_lambda * next_non_terminal * last_gae
      self.advantages[step] = last_gae
    self._returns = self.advantages + self._values

  # --- sampling ---------------------------------------------------------------
  def get(self, batch_size: Optional[int] = None
          ) -> Generator[RolloutBufferSamples, None, None]:
    assert self.full, "buffer not full"
    total = self.buffer_size * self.n_envs
    obs = self.observations.reshape(total, self.obs_dim)
    act = self.actions.reshape(total, self.act_dim)
    val = self._values.reshape(total)
    logp = self.log_probs.reshape(total)
    adv = self.advantages.reshape(total)
    ret = self._returns.reshape(total)
    idx = th.randperm(total, device=self.device)
    if batch_size is None:
      batch_size = total
    for start in range(0, total, batch_size):
      b = idx[start:start + batch_size]
      yield RolloutBufferSamples(
        observations=obs[b], actions=act[b], old_values=val[b],
        old_log_prob=logp[b], advantages=adv[b], returns=ret[b])

record_extras

record_extras(l_x)

Capture per-step extras into the slot add() will fill.

Tensor twin of the numpy buffers' record_extras(infos): the GPU-resident env hands back l_x from step_tensor directly, so there are no infos to parse. No-op here — the avoid operator needs only g.

Source code in safety_sb3/buffers_tensor.py

def record_extras(self, l_x: th.Tensor) -> None:
  """Capture per-step extras into the slot ``add()`` will fill.

  Tensor twin of the numpy buffers' ``record_extras(infos)``: the GPU-resident
  env hands back ``l_x`` from ``step_tensor`` directly, so there are no infos
  to parse. No-op here — the avoid operator needs only ``g``.
  """

TensorReachAvoidRolloutBuffer

Bases: TensorSafetyRolloutBuffer

Torch rollout buffer with the reach-avoid backup (adds l_x).

Torch twin of :class:safety_sb3.buffers_rollout.ReachAvoidRolloutBuffer; see it for the operator, the anchor rationale, and terminal_type.

Source code in safety_sb3/buffers_tensor.py

class TensorReachAvoidRolloutBuffer(TensorSafetyRolloutBuffer):
  """Torch rollout buffer with the reach-avoid backup (adds ``l_x``).

  Torch twin of :class:`safety_sb3.buffers_rollout.ReachAvoidRolloutBuffer`; see
  it for the operator, the anchor rationale, and ``terminal_type``.
  """

  _MODE = backups.REACH_AVOID

  def __init__(self, *args, terminal_type: str = "all", **kwargs):
    self.terminal_type = backups.check_terminal_type(terminal_type)
    super().__init__(*args, **kwargs)

  def reset(self) -> None:
    super().reset()
    self.l_x = th.zeros(self.buffer_size, self.n_envs, device=self.device)

  def record_extras(self, l_x: th.Tensor) -> None:
    """Keep this step's target margin ``l(s)`` (straight off ``step_tensor``)."""
    self.l_x[self.pos] = l_x.reshape(self.n_envs)

record_extras

record_extras(l_x)

Keep this step's target margin l(s) (straight off step_tensor).

Source code in safety_sb3/buffers_tensor.py

def record_extras(self, l_x: th.Tensor) -> None:
  """Keep this step's target margin ``l(s)`` (straight off ``step_tensor``)."""
  self.l_x[self.pos] = l_x.reshape(self.n_envs)

TensorCumulativeRolloutBuffer

Bases: TensorSafetyRolloutBuffer

Torch twin of :class:safety_sb3.buffers_rollout.CumulativeRolloutBuffer — the ordinary discounted-return backup, i.e. stock GAE on device.

Source code in safety_sb3/buffers_tensor.py

class TensorCumulativeRolloutBuffer(TensorSafetyRolloutBuffer):
  """Torch twin of :class:`safety_sb3.buffers_rollout.CumulativeRolloutBuffer` —
  the ordinary discounted-return backup, i.e. stock GAE on device."""

  _MODE = backups.CUMULATIVE

safety_sb3.buffers_replay

Off-policy replay buffer for the reach-avoid Mode — player-agnostic.

Like the rollout buffers, replay buffers carry the M of MAP and nothing else: both the single-player and two-player reach-avoid SAC learners store into the same class (the two-player game concatenates [a_ctrl, a_dstb] into the ordinary action field, so no extra column is needed).

ReachAvoidReplayBuffer extends SB3's :class:ReplayBuffer with the target margin l(s) that the reach-avoid backup needs::

V(s) = (1-gamma)*min(l, g) + gamma*min(g, max(l, V(s')))

The safety margin g(s) keeps riding on the standard reward field; the env supplies l(s) per step via info["l_x"], and the buffer captures it itself in :meth:ReachAvoidReplayBuffer.add. The avoid and cumulative modes need no extra column and use SB3's stock ReplayBuffer.

ReachAvoidReplayBuffer

Bases: ReplayBuffer

SB3 replay buffer that also stores the per-step target margin l(s).

l(s) is read from info["l_x"] on each add. Assumes the default SAC buffer layout (optimize_memory_usage=False), which stores next_observations explicitly.

Source code in safety_sb3/buffers_replay.py

class ReachAvoidReplayBuffer(ReplayBuffer):
  """SB3 replay buffer that also stores the per-step target margin ``l(s)``.

  ``l(s)`` is read from ``info["l_x"]`` on each ``add``. Assumes the default
  SAC buffer layout (``optimize_memory_usage=False``), which stores
  ``next_observations`` explicitly.
  """

  def __init__(self, *args, **kwargs) -> None:
    super().__init__(*args, **kwargs)
    if self.optimize_memory_usage:
      raise ValueError(
        "ReachAvoidReplayBuffer requires optimize_memory_usage=False "
        "(it stores next_observations explicitly)."
      )
    self.l_x = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)

  def add(self, obs, next_obs, action, reward, done, infos) -> None:
    # Must capture l_x at the CURRENT pos before super().add() advances it.
    for env_i, info in enumerate(infos):
      self.l_x[self.pos, env_i] = float(info.get("l_x", 0.0))
    super().add(obs, next_obs, action, reward, done, infos)

  def sample(self, batch_size: int, env=None) -> ReachAvoidReplayBufferSamples:
    upper = self.buffer_size if self.full else self.pos
    batch_inds = np.random.randint(0, upper, size=batch_size)
    env_indices = np.random.randint(0, self.n_envs, size=batch_size)

    next_obs = self._normalize_obs(
      self.next_observations[batch_inds, env_indices, :], env
    )
    return ReachAvoidReplayBufferSamples(
      observations=self.to_torch(
        self._normalize_obs(self.observations[batch_inds, env_indices, :], env)
      ),
      actions=self.to_torch(self.actions[batch_inds, env_indices, :]),
      next_observations=self.to_torch(next_obs),
      dones=self.to_torch(
        (
          self.dones[batch_inds, env_indices]
          * (1 - self.timeouts[batch_inds, env_indices])
        ).reshape(-1, 1)
      ),
      rewards=self.to_torch(
        self._normalize_reward(
          self.rewards[batch_inds, env_indices].reshape(-1, 1), env
        )
      ),
      l_x=self.to_torch(self.l_x[batch_inds, env_indices].reshape(-1, 1)),
    )