Code reference

Auto-generated from source docstrings — the single source of truth. See the API guide for the conceptual contract.

Registry

robot_safety_sandbox.registry

Task registry: one place future work looks to run or add benchmark tasks.

from robot_safety_sandbox import make_tensor, list_tasks
env = make_tensor("go2_gap_chain", num_envs=2048)   # -> TensorVecEnv
model = ReachAvoidPPO1P("MlpPolicy", env, normalize_obs=True, ...)

A :class:TaskSpec pins everything a benchmark run needs: the mjlab cfg builder (spawn events + curricula), the reach-avoid margins, the action dims, and the task's mode. Curriculum LINEAGE (landing -> crossing -> chain) is deliberately NOT a field here: it lives in docs/log/experiments.md, which is where it is actually maintained.

Here's a MAP to navigate the codebase — Mode. Algorithm. Players.

M = Mode       Safety | ReachAvoid | Cumulative    the Bellman operator
A = Algorithm  PPO | SAC | A2C | DQN               the RL update rule
P = Players    1P | 2P                             single-player | zero-sum

A learner's NAME is those three letters concatenated in that order — SafetyPPO1P, ReachAvoidSAC2P — and :func:algo_name is that concatenation and nothing else. No lookup table, no per-task override:

M comes from the TASK   its ``mode`` (below), a property of its margins
A comes from the RUN    the trainer family (on_policy -> PPO, off_policy -> SAC)
P comes from the RUN    the ``--adversary`` flag

Every task declares exactly one axis of its own, its mode — the safety_sb3 BACKUP it is trained under (see :data:MODES):

mode="safety" AVOID: V = min(g, gamma V') margins, no l mode="reach-avoid" REACH_AVOID: V = min(g, max(l, gamma V')) margins with l mode="cumulative" CUMULATIVE: V = r + gamma (1-d) V' dense env reward, i.e. the plain-RL TASK policy a filter wraps (margin_fn=None; envs are auto-built in dense mode, trained with STOCK SB3).

A full filter experiment needs both layers: a cumulative task policy (pi_task) and a safety task supplying the certificate V(s) + fallback.

The two SAFETY modes give four learners per algorithm, one per (M, P) cell — the mode is a property of the TASK's margins, the player count a property of the RUN:

                   avoid (no l)      reach-avoid (real l)
    single-player  SafetyPPO1P       ReachAvoidPPO1P
    two-player     SafetyPPO2P       ReachAvoidPPO2P

mode picks the COLUMN, --adversary the row. CUMULATIVE has no P axis at all: it is stock stable_baselines3 ("PPO" / "SAC", no suffix) and there is no two-player cumulative game.

NB the 2P cell is a DIFFERENT ALGORITHM in the two families, not the same game with a different optimizer: *SAC2P is the minimax game on one shared joint-action critic Q(s, [a_ctrl, a_dstb]), while *PPO2P is an alternating best-response approximation with two independent V(s) nets, two rollout buffers and a phase machine. See docs/API.md.

TaskSpec dataclass

Source code in robot_safety_sandbox/registry.py

@dataclass
class TaskSpec:
  task_id: str
  cfg_builder: Callable          # (play: bool) -> ManagerBasedRlEnvCfg
  margin_fn: Optional[Callable] = None  # (env) -> (g, l); None for cumulative
  description: str = ""
  ctrl_dim: int = 12
  dstb_dim: int = 3              # adversary force dims (two-player games)
  # REQUIRED. Which BACKUP values this task: one of MODES. This is the task's
  # whole say in the MAP — it supplies the **M**, and nothing else here names a
  # learner. It also fixes whether the env is built in dense-reward mode and
  # whether margins are required.
  mode: Optional[str] = None
  supports_adversary: bool = False
  # WHEN the episode ends from (g, l); one of END_CRITERIA. "failure" (default)
  # reproduces today's behavior for EVERY registered task — an audit (2026-07-17)
  # found none currently terminate on success. A run may override it via the
  # trainer's --end-criterion flag; the two knobs (this + terminal_type) are
  # orthogonal. Set "reach-avoid" on a task only if it should end on reach.
  end_criterion: str = "failure"
  kwargs: dict = field(default_factory=dict)  # extra bridge kwargs

  def __post_init__(self):
    if self.mode is None:
      raise ValueError(
        f"task '{self.task_id}' declares no mode=; every task must name the "
        f"backup it trains under, one of {MODES}. The learner is DERIVED from "
        f"it (see algo_name), so there is nothing else to declare.")
    if self.mode not in MODES:
      raise ValueError(f"task '{self.task_id}' has mode={self.mode!r}; "
                       f"must be one of {MODES}")
    if self.mode != CUMULATIVE and self.margin_fn is None:
      raise ValueError(
        f"task '{self.task_id}' is mode={self.mode!r} and needs a margin_fn "
        f"(only mode={CUMULATIVE!r} trains without margins, on dense reward)")
    if self.end_criterion not in END_CRITERIA:
      raise ValueError(
        f"task '{self.task_id}' has end_criterion={self.end_criterion!r}; "
        f"must be one of {END_CRITERIA}")

register

register(spec)

Source code in robot_safety_sandbox/registry.py

def register(spec: TaskSpec) -> None:
  if spec.task_id in _REGISTRY:
    raise ValueError(f"task '{spec.task_id}' already registered")
  _REGISTRY[spec.task_id] = spec

spec

spec(task_id)

Source code in robot_safety_sandbox/registry.py

def spec(task_id: str) -> TaskSpec:
  if task_id not in _REGISTRY:
    raise KeyError(
      f"unknown task '{task_id}'. Registered: {list_tasks()}. "
      "(Some tasks require their source repo on sys.path during the "
      "phase-1 compat period — see tasks/*.py and MIGRATION.md.)")
  return _REGISTRY[task_id]

list_tasks

list_tasks(mode=None)

Registered task ids, optionally filtered to one :data:MODES entry.

Source code in robot_safety_sandbox/registry.py

def list_tasks(mode: Optional[str] = None) -> list[str]:
  """Registered task ids, optionally filtered to one :data:`MODES` entry."""
  if mode is not None and mode not in MODES:
    raise ValueError(f"list_tasks: unknown mode={mode!r}; one of {MODES}")
  return sorted(t for t, s in _REGISTRY.items()
                if mode is None or s.mode == mode)

algo_name

algo_name(task_id, adversary=False, family='on_policy')

The learner CLASS NAME for running task_id — the MAP, spelled out.

It is a FORMULA, not a lookup: Mode + Algorithm + Players, concatenated. Nothing in the registry can override it; a task supplies only its mode.

algo_name("go2_stabilize")                            -> ReachAvoidPPO1P
algo_name("go2_stabilize", adversary=True)            -> ReachAvoidPPO2P
algo_name("go2_stabilize", family="off_policy")       -> ReachAvoidSAC1P
algo_name("digit_stabilize_avoid", adversary=True)    -> SafetyPPO2P

Names only — this module never imports safety_sb3, so a cumulative-only install (dense reward + vanilla SB3, see base.py) still imports the registry.

mode=CUMULATIVE returns the bare algorithm ("PPO" / "SAC"): the plain-RL task policy trains with STOCK stable_baselines3, which keeps the checkpoint a vanilla SB3 zip. It has no P axis — there is no two-player cumulative game.

Two pairings are refused rather than resolved:

• a two-player CUMULATIVE run — no such game exists;
• a reach-avoid learner on an avoid-only task (no target set). Pinning l to a constant does NOT make the reach-avoid backup compute the avoid value — a negative l empties the safe set, a non-negative one strips the lookahead — so it has no valid formulation and must not be reachable by accident. See margins.py.

Source code in robot_safety_sandbox/registry.py

def algo_name(task_id: str, adversary: bool = False,
              family: str = "on_policy") -> str:
  """The learner CLASS NAME for running ``task_id`` — the MAP, spelled out.

  It is a FORMULA, not a lookup: **M**ode + **A**lgorithm + **P**layers,
  concatenated. Nothing in the registry can override it; a task supplies only
  its ``mode``.

      algo_name("go2_stabilize")                            -> ReachAvoidPPO1P
      algo_name("go2_stabilize", adversary=True)            -> ReachAvoidPPO2P
      algo_name("go2_stabilize", family="off_policy")       -> ReachAvoidSAC1P
      algo_name("digit_stabilize_avoid", adversary=True)    -> SafetyPPO2P

  Names only — this module never imports safety_sb3, so a cumulative-only
  install (dense reward + vanilla SB3, see base.py) still imports the registry.

  mode=CUMULATIVE returns the bare algorithm (``"PPO"`` / ``"SAC"``): the
  plain-RL task policy trains with STOCK stable_baselines3, which keeps the
  checkpoint a vanilla SB3 zip. It has no P axis — there is no two-player
  cumulative game.

  Two pairings are refused rather than resolved:

  * a two-player CUMULATIVE run — no such game exists;
  * a reach-avoid learner on an avoid-only task (no target set). Pinning l to a
    constant does NOT make the reach-avoid backup compute the avoid value — a
    negative l empties the safe set, a non-negative one strips the lookahead —
    so it has no valid formulation and must not be reachable by accident. See
    margins.py.
  """
  if family not in _ALG:
    raise ValueError(f"unknown family={family!r}; one of {FAMILIES}")
  s = spec(task_id)
  if adversary and not s.supports_adversary:
    raise ValueError(f"task '{task_id}' does not define an adversary")
  if adversary and s.mode == CUMULATIVE:
    raise ValueError(f"task '{task_id}' is mode={CUMULATIVE!r}: there is no "
                     f"two-player cumulative learner")
  # margin_fns built by margins.compose/avoid_only carry has_target; anything
  # else (task-local margin builders) is assumed to declare a real l.
  if s.mode == REACH_AVOID and not getattr(s.margin_fn, "has_target", True):
    raise ValueError(
      f"task '{task_id}' is AVOID-ONLY (its margin_fn declares no target set) "
      f"but declares mode={REACH_AVOID!r}. Avoid is not a reach-avoid instance "
      f"for ANY constant l — declare mode={AVOID!r} (--adversary then gives the "
      f"two-player Safety{_ALG[family]}2P), or give the task a real reach "
      f"margin l. See margins.py.")
  if s.mode == CUMULATIVE:
    return _ALG[family]                      # stock stable_baselines3, 1 player
  return f"{_PREFIX[s.mode]}{_ALG[family]}{'2P' if adversary else '1P'}"

make_tensor

make_tensor(task_id, num_envs=2048, device='cuda:0', adversary=False, end_criterion=None, cfg_overrides=None, **kw)

GPU-resident env (primary path; pair with safety_sb3 PPO learners).

end_criterion (None -> the task's TaskSpec value; else an explicit override, one of :data:END_CRITERIA) sets WHEN the episode ends from (g, l). cfg_overrides (dict) forwards experiment-level env/task params to the task's cfg_builder, overriding the values baked into its registration (e.g. {"gate_close_rate": 0.003}) -- lets a config recipe tune the env without editing the task or the trainer flags.

Source code in robot_safety_sandbox/registry.py

def make_tensor(task_id: str, num_envs: int = 2048, device: str = "cuda:0",
                adversary: bool = False, end_criterion: Optional[str] = None,
                cfg_overrides: Optional[dict] = None, **kw):
  """GPU-resident env (primary path; pair with safety_sb3 PPO learners).

  ``end_criterion`` (None -> the task's TaskSpec value; else an explicit
  override, one of :data:`END_CRITERIA`) sets WHEN the episode ends from (g, l).
  ``cfg_overrides`` (dict) forwards experiment-level env/task params to the
  task's cfg_builder, overriding the values baked into its registration (e.g.
  ``{"gate_close_rate": 0.003}``) -- lets a config recipe tune the env without
  editing the task or the trainer flags.
  """
  from .base import MjlabTensorSafetyEnv
  s = spec(task_id)
  if adversary and not s.supports_adversary:
    raise ValueError(f"task '{task_id}' does not define an adversary")
  kw.setdefault("dense_reward", s.mode == CUMULATIVE)  # cumulative => dense
  ec = end_criterion if end_criterion is not None else s.end_criterion
  return MjlabTensorSafetyEnv(
    num_envs, device, cfg_builder=s.cfg_builder, margin_fn=s.margin_fn,
    ctrl_dim=s.ctrl_dim, dstb_dim=s.dstb_dim, adversary=adversary,
    end_criterion=ec, cfg_overrides=cfg_overrides, **{**s.kwargs, **kw})

make_numpy

make_numpy(task_id, num_envs=64, device='cuda:0', adversary=False, end_criterion=None, cfg_overrides=None, **kw)

Classic SB3 VecEnv (for the SAC family / stock SB3 tooling).

end_criterion (None -> the task's TaskSpec value; else an override) sets WHEN the episode ends from (g, l). cfg_overrides (dict) forwards env/task params to the task's cfg_builder. See :func:make_tensor.

Source code in robot_safety_sandbox/registry.py

def make_numpy(task_id: str, num_envs: int = 64, device: str = "cuda:0",
               adversary: bool = False, end_criterion: Optional[str] = None,
               cfg_overrides: Optional[dict] = None, **kw):
  """Classic SB3 VecEnv (for the SAC family / stock SB3 tooling).

  ``end_criterion`` (None -> the task's TaskSpec value; else an override) sets
  WHEN the episode ends from (g, l). ``cfg_overrides`` (dict) forwards env/task
  params to the task's cfg_builder. See :func:`make_tensor`.
  """
  from .base import MjlabNumpySafetyEnv
  s = spec(task_id)
  if adversary and not s.supports_adversary:
    raise ValueError(f"task '{task_id}' does not define an adversary")
  kw.setdefault("dense_reward", s.mode == CUMULATIVE)  # cumulative => dense
  ec = end_criterion if end_criterion is not None else s.end_criterion
  return MjlabNumpySafetyEnv(
    num_envs, device, cfg_builder=s.cfg_builder, margin_fn=s.margin_fn,
    ctrl_dim=s.ctrl_dim, dstb_dim=s.dstb_dim, adversary=adversary,
    end_criterion=ec, cfg_overrides=cfg_overrides, **{**s.kwargs, **kw})

Margins

robot_safety_sandbox.margins

Composable reach-avoid margin library (batched torch, mjlab scene API).

A task's margin_fn(env) -> (g, l) is composed from these terms:

g (avoid)  : stay out of the failure set    — g < 0 == failure
l (reach)  : arrive in the target set       — l >= 0 == reached

Conventions: margins are signed distances normalized to O(1); the g terminal anchor and clamping live in :mod:base / the buffers, not here. All terms are validated on the Go2 parkour tasks (ported from the reference wrappers).

compose

compose(g_fn, l_fn=None, clamp=CLAMP)

margin_fn from a g term and an l term (clamped for value regression).

l_fn=None declares an AVOID-ONLY task: there is no target set. The bridge still has to hand the learner an l channel (step_tensor returns a 5-tuple), so a zero placeholder is emitted — but the returned margin_fn is tagged has_target = False, and it is ONLY valid under an avoid learner (the Safety* half of the MAP, 1P or 2P), which ignores l. Feeding it to a reach-avoid learner (ReachAvoid*) is the degenerate l_zero case above; the tag exists so that mistake raises instead of training silently — see registry.algo_name.

Source code in robot_safety_sandbox/margins.py

def compose(g_fn, l_fn=None, clamp: float = CLAMP):
  """margin_fn from a g term and an l term (clamped for value regression).

  ``l_fn=None`` declares an AVOID-ONLY task: there is no target set. The bridge
  still has to hand the learner an ``l`` channel (``step_tensor`` returns a
  5-tuple), so a zero placeholder is emitted — but the returned margin_fn is
  tagged ``has_target = False``, and it is ONLY valid under an avoid learner
  (the ``Safety*`` half of the MAP, 1P or 2P), which ignores l. Feeding it to a
  reach-avoid learner (``ReachAvoid*``) is the degenerate ``l_zero`` case above;
  the tag exists so that mistake raises instead of training silently — see
  ``registry.algo_name``.
  """
  def margin_fn(env):
    g = g_fn(env).clamp(-clamp, clamp)
    if l_fn is None:
      return g, torch.zeros_like(g)  # placeholder; avoid learners ignore it
    return g, l_fn(env).clamp(-clamp, clamp)
  margin_fn.has_target = l_fn is not None
  return margin_fn

avoid_only

avoid_only(margin_fn)

Strip the reach target off an existing margin_fn -> avoid-only twin.

For twin pairs that share one (g, l) builder and differ only in the backup: the avoid twin keeps g verbatim and declares no target. Same contract as compose(g_fn) — the l channel is a zero placeholder, has_target is False, and only an avoid learner may consume it.

Source code in robot_safety_sandbox/margins.py

def avoid_only(margin_fn):
  """Strip the reach target off an existing ``margin_fn`` -> avoid-only twin.

  For twin pairs that share one (g, l) builder and differ only in the backup:
  the avoid twin keeps g verbatim and declares no target. Same contract as
  ``compose(g_fn)`` — the l channel is a zero placeholder, ``has_target`` is
  False, and only an avoid learner may consume it.
  """
  def fn(env):
    g, _l = margin_fn(env)
    return g, torch.zeros_like(g)
  fn.has_target = False
  return fn

Bridges

robot_safety_sandbox.base

Base classes: run any mjlab task with safety_sb3 (tensor or numpy path).

The zoo's env contract (matches what every safety_sb3 learner consumes):

reward   = g(s)   the physical safety margin  (NEVER normalize/reshape it)
l_x      = l(s)   the target margin           (zeros for avoid-only tasks)
dones    = terminated | truncated  (mjlab auto-resets internally)
timeouts = truncated & ~terminated (no value bootstrap: g is absolute)

A task is fully specified by two callables (see :mod:registry):

cfg_builder(play: bool) -> ManagerBasedRlEnvCfg   # the mjlab env
margin_fn(env)          -> (g, l) batched tensors # the reach-avoid margins

Everything else here is plumbing: :class:MjlabTensorSafetyEnv exposes the task as a safety_sb3 TensorVecEnv (GPU-resident, ~50k steps/s on a 12GB card at 2048 envs); :class:MjlabNumpySafetyEnv as a classic SB3 VecEnv (for the SAC family / plain SB3 tooling). Curriculum levels and task metrics (mjlab extras['log']) are forwarded via metrics() so training logs always show curriculum progression — silent curriculum stalls were the single biggest source of lost weeks in the reference tasks.

Porting a new task = write a cfg_builder (spawn events + curricula, plain mjlab) + a margin_fn (compose from :mod:margins), then register() it.

MjlabTensorSafetyEnv

Bases: _MjlabCore, TensorVecEnv

GPU-resident bridge (primary): torch end-to-end, no numpy bounce. Pair with a safety_sb3 1P learner, SafetyPPO1P / ReachAvoidPPO1P (auto-detected).

Source code in robot_safety_sandbox/base.py

class MjlabTensorSafetyEnv(_MjlabCore, TensorVecEnv):
  """GPU-resident bridge (primary): torch end-to-end, no numpy bounce.
  Pair with a safety_sb3 1P learner, SafetyPPO1P / ReachAvoidPPO1P
  (auto-detected)."""

  def __init__(self, num_envs=2048, device="cuda:0", *, cfg_builder,
               margin_fn, ctrl_dim=12, dstb_dim=3, ctrl_gain=3.0,
               force_max=50.0, adversary=False, adversary_body="base_link",
               render_mode=None, obs_key=None, dense_reward=False,
               dstb_mode="wrench", dstb_gain=0.25, hybrid_skill=None,
                 latch_margin_fn=None, end_criterion="failure",
               cfg_overrides=None, dense_margins=False):
    if not _HAS_SAFETY_SB3:
      raise ImportError(
        "safety_sb3 is required for the tensor bridge (pip install it or put "
        "the safety-stable-baselines repo on sys.path); the numpy bridge "
        "(make_numpy) works with vanilla stable_baselines3 only.")
    obs_space, act_space = self._init_core(
      num_envs, device, cfg_builder, margin_fn, ctrl_dim=ctrl_dim,
      dstb_dim=dstb_dim, ctrl_gain=ctrl_gain, force_max=force_max,
      adversary=adversary, adversary_body=adversary_body,
      render_mode=render_mode, obs_key=obs_key, dense_reward=dense_reward,
      dstb_mode=dstb_mode, dstb_gain=dstb_gain, hybrid_skill=hybrid_skill,
      latch_margin_fn=latch_margin_fn, end_criterion=end_criterion,
      cfg_overrides=cfg_overrides, dense_margins=dense_margins)
    TensorVecEnv.__init__(self, int(num_envs), obs_space, act_space, device)

  def reset(self) -> torch.Tensor:
    obs_dict, _ = self.mj.reset()
    self._obs_dict = obs_dict
    obs = obs_dict[self.obs_key].float()
    self.mj._zoo_last_obs = obs
    self._hyb_latch.zero_()
    self._last_l = None
    return obs

  def step_tensor(self, actions: torch.Tensor):
    obs, g, terminated, truncated, l = self._mj_step(actions)
    # NaN sanitation (see safety_margin_hook): a NaN obs entering the rollout
    # buffer NaNs the policy update even though the env resets via nan_term.
    obs = torch.nan_to_num(obs, nan=0.0, posinf=0.0, neginf=0.0)
    dones = terminated | truncated
    timeouts = truncated & ~terminated
    return obs, g, dones, timeouts, l

MjlabNumpySafetyEnv

Bases: _MjlabCore, VecEnv

Classic numpy SB3 VecEnv bridge (for the SAC family / stock SB3 tooling). Slower (device<->host each step); prefer the tensor bridge for on-policy training.

Source code in robot_safety_sandbox/base.py

class MjlabNumpySafetyEnv(_MjlabCore, VecEnv):
  """Classic numpy SB3 ``VecEnv`` bridge (for the SAC family / stock SB3
  tooling). Slower (device<->host each step); prefer the tensor bridge for
  on-policy training."""

  def __init__(self, num_envs=64, device="cuda:0", *, cfg_builder, margin_fn,
               ctrl_dim=12, dstb_dim=3, ctrl_gain=3.0, force_max=50.0,
               adversary=False, adversary_body="base_link", render_mode=None,
               obs_key=None, dense_reward=False, dstb_mode="wrench",
               dstb_gain=0.25, hybrid_skill=None, latch_margin_fn=None,
               end_criterion="failure", cfg_overrides=None,
               dense_margins=False):
    obs_space, act_space = self._init_core(
      num_envs, device, cfg_builder, margin_fn, ctrl_dim=ctrl_dim,
      dstb_dim=dstb_dim, ctrl_gain=ctrl_gain, force_max=force_max,
      adversary=adversary, adversary_body=adversary_body,
      render_mode=render_mode, obs_key=obs_key, dense_reward=dense_reward,
      dstb_mode=dstb_mode, dstb_gain=dstb_gain, hybrid_skill=hybrid_skill,
      latch_margin_fn=latch_margin_fn, end_criterion=end_criterion,
      cfg_overrides=cfg_overrides, dense_margins=dense_margins)
    self._device = device
    VecEnv.__init__(self, int(num_envs), obs_space, act_space)
    self._actions = None

  def reset(self):
    obs_dict, _ = self.mj.reset()
    self._obs_dict = obs_dict
    return obs_dict[self.obs_key].float().cpu().numpy()

  def step_async(self, actions):
    self._actions = np.asarray(actions, dtype=np.float32)

  def step_wait(self):
    a = torch.as_tensor(self._actions, device=self._device)
    obs, g, terminated, truncated, l = self._mj_step(a)
    term = terminated.cpu().numpy()
    trunc = truncated.cpu().numpy()
    dones = np.logical_or(term, trunc)
    obs_np = obs.cpu().numpy()
    l_np = l.cpu().numpy()
    infos = []
    for i in range(self.num_envs):
      info = {"l_x": float(l_np[i])}
      if dones[i]:
        info["terminal_observation"] = obs_np[i]
        info["TimeLimit.truncated"] = bool(trunc[i] and not term[i])
      infos.append(info)
    return obs_np, g.cpu().numpy(), dones, infos

  # VecEnv boilerplate
  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)]

build_task_cfg

build_task_cfg(cfg_builder, margin_fn, num_envs, drop_events=('push_robot',), dense=False, end_criterion='failure', cfg_overrides=None, dense_margins=False)

Assemble an mjlab cfg for the zoo: task cfg + the margin hook.

drop_events removes events a learned adversary replaces (default: the random push). dense=True is the STAGE-1 locomotion mode: the reward stays the env's own dense reward stack (_r) and the safety hook is NOT added, so a stock PPO shapes a proper gait (reach-avoid g/l are not used here).

end_criterion (see :data:registry.END_CRITERIA) sets WHEN episodes end, UNIFORMLY across tasks. The task cfg's own terminations already encode the FAILURE set (fell_over / illegal_contact / ... -> the reward hook anchors g<0 there) plus the env timeout, so: * "failure" : add nothing. The task's existing terminations stand as-is -> bit-identical to pre-end_criterion behavior. Reaching the target does not end the episode; the agent reaches deeper (l keeps climbing to the g ceiling). * "reach-avoid" : ALSO add the uniform success DoneTerm (g>=0 AND l>=0) on top of the failure terminations -> reach-and-stop. * "timeout" : strip every non-time_out (failure) termination so ONLY the env timeout ends episodes, and add no success term (pure value-learning / diagnostic). g is still anchored nowhere, so failures live on in the margin, they just don't reset.

Source code in robot_safety_sandbox/base.py

def build_task_cfg(cfg_builder: Callable, margin_fn: Callable, num_envs: int,
                   drop_events: tuple[str, ...] = ("push_robot",),
                   dense: bool = False, end_criterion: str = "failure",
                   cfg_overrides: dict | None = None,
                   dense_margins: bool = False):
  """Assemble an mjlab cfg for the zoo: task cfg + the margin hook.

  ``drop_events`` removes events a learned adversary replaces (default: the
  random push). ``dense=True`` is the STAGE-1 locomotion mode: the reward stays
  the env's own dense reward stack (`_r`) and the safety hook is NOT added, so a
  stock PPO shapes a proper gait (reach-avoid g/l are not used here).

  ``end_criterion`` (see :data:`registry.END_CRITERIA`) sets WHEN episodes end,
  UNIFORMLY across tasks. The task cfg's own terminations already encode the
  FAILURE set (fell_over / illegal_contact / ... -> the reward hook anchors g<0
  there) plus the env timeout, so:
    * "failure"     : add nothing. The task's existing terminations stand as-is
                      -> bit-identical to pre-end_criterion behavior. Reaching
                      the target does not end the episode; the agent reaches
                      deeper (l keeps climbing to the g ceiling).
    * "reach-avoid" : ALSO add the uniform success DoneTerm (g>=0 AND l>=0) on
                      top of the failure terminations -> reach-and-stop.
    * "timeout"     : strip every non-time_out (failure) termination so ONLY the
                      env timeout ends episodes, and add no success term (pure
                      value-learning / diagnostic). g is still anchored nowhere,
                      so failures live on in the margin, they just don't reset.
  """
  # cfg_overrides: experiment-level env/task params forwarded to the cfg_builder
  # (partial call-kwargs override the task registration's baked values, e.g.
  # {"gate_close_rate": 0.003}). Fail loud on a param the cfg_builder rejects.
  try:
    cfg = cfg_builder(play=False, **(cfg_overrides or {}))
  except TypeError as e:
    if cfg_overrides:
      raise SystemExit(
        f"[cfg_overrides] {list(cfg_overrides)} not all accepted by the task's "
        f"cfg_builder: {e}") from e
    raise
  cfg.scene.num_envs = int(num_envs)
  if cfg.events is not None:
    for e in drop_events:
      cfg.events.pop(e, None)
  if dense and dense_margins:
    if margin_fn is None:
      raise ValueError(
        "dense_margins=True needs a margin_fn: it exists to give a dense-reward "
        "task the margins a train-time safety filter monitors. Register the "
        "task with the margin_fn of the twin that will be its fallback.")
    cfg.rewards["zoo_margin_probe"] = RewardTermCfg(
      func=margin_probe_hook, weight=1.0, params={"margin_fn": margin_fn})
  if not dense:
    if margin_fn is None:
      raise ValueError(
        "margin_fn=None with dense=False: a mode='cumulative' task trains on "
        "the dense env reward — build it in dense mode (make_tensor/make_numpy "
        "do this automatically from the task's mode).")
    cfg.rewards["zoo_safety_hook"] = RewardTermCfg(
      func=safety_margin_hook, weight=1.0, params={"margin_fn": margin_fn})
    if end_criterion == "reach-avoid":
      cfg.terminations[SUCCESS_TERM] = TerminationTermCfg(
        func=zoo_reach_success, params={"margin_fn": margin_fn})
    elif end_criterion == "timeout":
      # Keep only the env timeout (time_out=True); drop the failure set so the
      # episode always runs to the horizon (diagnostic / pure value-learning).
      for name in [n for n, t in cfg.terminations.items()
                   if not getattr(t, "time_out", False)]:
        cfg.terminations.pop(name)
  return cfg