feat(sim,training): actuator-keyed policies, vertical-camera fix, RL parity (evaluate/vec/gym/staged_reward)

Sim bug fixes (found by running policies across embodiments, each pinned):
- robot_action_keys(): policies were keyed by ALL joint names incl. passive/
  mimic followers and missed tendon actuators, so g1/aloha/xarm7/stretch/hands
  silently no-op'd. Now keyed by the robot's actuators (what send_action
  resolves). MuJoCo override + SimEngine ABC default (other backends unchanged).
- spec_builder._target_quat: returned None for top-down/vertical cameras
  (forward parallel to world +Z -> degenerate cross). Falls back to +Y up.
- mujoco render(output_path=...): persist PNG for verifiable renders, with
  path-traversal / shell-metachar guard on the LLM-supplied path.
- eval_policy: auto-resolves the single robot (API parity with run_policy).
- DeclarativeBenchmark.on_episode_start: reset stateful reward terms (no phase
  leak across episodes).

RL training parity (4 gaps closed, CPU-tested, no new required deps):
- predicates.staged_reward: declarative monotonic phase-machine reward
  (StatefulRewardTerm). Curricula authored as DATA via the closed predicate
  registry - never eval, never shipped task code.
- base_algo.evaluate()/load_checkpoint()/_deterministic_action(): deterministic
  eval peer of train(), frozen normalizer, success_rate; works for PPO + SAC.
- SimEnv.n_substeps (default 5): position-target PD needs substeps to track;
  n_substeps=1 was the "nothing learns" bug. Resets stateful reward terms.
- PpoTrainer: vectorized rollout path; num_envs>1 now accepted (was hard-
  rejected). N=1 path byte-identical (dispatch on _vectorized).
- VecSimEnv: N independent SimEnv, one reused thread pool, autoreset with
  terminal_obs capture (load-bearing for GAE bootstrap), homogeneity guard.
- GymSimEnv: gymnasium adapter with correct terminated-vs-truncated split
  (the SB3 TimeLimit footgun fixed at the source) + critic-obs passthrough.

Tests: 38 new CPU tests (vec_env, evaluate, vectorized-ppo, gym_env,
staged_reward) + SO-101 e2e wiring/convergence integ. Updated 3 existing tests
that encoded the old contracts (num_envs==1 rejection, eval_policy robot_name).

Author: cagataycali

strands_robots/simulation/benchmark_spec.py

@@ -236,6 +236,13 @@ def on_episode_start(self, sim: SimEngine, rng: random.Random) -> None:
                     raise RuntimeError(
                         f"DeclarativeBenchmark '{self._name}': load_scene({self._scene!r}) failed: {msg}"
                     )
+        # Reset any stateful reward terms (e.g. a staged_reward phase machine)
+        # so per-episode phase state does not leak across episodes. Stateless
+        # function terms have no reset() and are skipped.
+        for term in self._reward_terms:
+            term_reset = getattr(term, "reset", None)
+            if callable(term_reset):
+                term_reset()
         super().on_episode_start(sim, rng)
 
     def on_step(self, sim: SimEngine, obs: dict[str, Any], action: dict[str, Any]) -> StepInfo:

strands_robots/simulation/mujoco/rendering.py

@@ -571,10 +571,21 @@ def _scale_ctrl_for_actuator(model: Any, ai: int, value: float, mj: Any) -> floa
         return lo + frac * span
 
     def render(
-        self, camera_name: str = "default", width: int | None = None, height: int | None = None
+        self,
+        camera_name: str = "default",
+        width: int | None = None,
+        height: int | None = None,
+        output_path: str | None = None,
     ) -> dict[str, Any]:
         """Render a camera view to a PNG image.
 
+        When ``output_path`` is given the PNG is ALSO written to that file path
+        (parent dirs created) and the saved path is reported in the ``json``
+        block as ``saved_path`` and in the text summary. This lets an agent (or
+        a human) persist a render for independent verification instead of only
+        receiving the bytes inline.
+
+
         Returns an agent-tool dict with ``status`` and a ``content`` list; on
         success the content holds an ``image`` block carrying PNG bytes
         (``{"image": {"format": "png", "source": {"bytes": ...}}}``) plus a
@@ -662,12 +673,36 @@ def render(
             pixel_var = float(_np.var(img))
             pixel_mean = float(_np.mean(img))
 
+            saved_path: str | None = None
+            if output_path:
+                import os as _os
+
+                # Validate against shell/path-traversal injection (LLM-supplied).
+                bad = {";", "|", "$", "`", ">", "<", "\n", "\r", "\x00"}
+                if any(b in output_path for b in bad) or ".." in output_path.split("/"):
+                    return {
+                        "status": "error",
+                        "content": [{"text": f"render: unsafe output_path {output_path!r}"}],
+                    }
+                _dir = _os.path.dirname(_os.path.abspath(output_path))
+                _os.makedirs(_dir, exist_ok=True)
+                with open(output_path, "wb") as _f:
+                    _f.write(png_bytes)
+                saved_path = _os.path.abspath(output_path)
+
+            summary = f"{w}x{h} from '{label}' at t={self._world.sim_time:.3f}s"
+            if saved_path:
+                summary += f" -> saved {saved_path}"
+            json_block = {"pixel_variance": pixel_var, "pixel_mean": pixel_mean, "camera": label}
+            if saved_path:
+                json_block["saved_path"] = saved_path
+
             return {
                 "status": "success",
                 "content": [
-                    {"text": f"{w}x{h} from '{label}' at t={self._world.sim_time:.3f}s"},
+                    {"text": summary},
                     {"image": {"format": "png", "source": {"bytes": png_bytes}}},
-                    {"json": {"pixel_variance": pixel_var, "pixel_mean": pixel_mean, "camera": label}},
+                    {"json": json_block},
                 ],
             }
         except Exception as e:

strands_robots/simulation/predicates.py

@@ -554,6 +554,162 @@ def term(_sim: SimEngine) -> float:
     return term
 
 
+# Stateful reward terms (declarative phase machine)
+#
+# A plain RewardTerm is stateless: ``(SimEngine) -> float``. Some rewards need
+# memory across steps - a pick-place curriculum advances Reach -> Grasp ->
+# Transport -> Place, awards a one-time bonus on each transition, and only ever
+# moves forward. Rather than hardcode any specific task, we expose ONE
+# generic primitive, ``staged_reward``, that composes EXISTING registry
+# predicates into a phase machine. The task itself is then authored as data
+# (a spec dict / YAML) by a human or LLM - never as shipped code, and never via
+# ``eval`` (sub-predicates are compiled through :func:`make_predicate`, the same
+# closed-registry path as every other DSL call).
+
+
+class StatefulRewardTerm:
+    """A reward term that carries per-episode state and must be ``reset()``.
+
+    Duck-typed by consumers: anything with ``__call__(sim) -> float`` AND a
+    zero-arg ``reset()`` is treated as episode-stateful. ``SimEnv.reset`` and
+    ``DeclarativeBenchmark.on_episode_start`` call ``reset()`` on any reward
+    term that has it, so stateless plain-function terms are unaffected.
+    """
+
+    def __call__(self, sim: SimEngine) -> float:  # pragma: no cover - interface
+        raise NotImplementedError
+
+    def reset(self) -> None:  # pragma: no cover - interface
+        raise NotImplementedError
+
+
+class _StagedReward(StatefulRewardTerm):
+    """Monotonic multi-stage (phase-machine) reward built from sub-predicates.
+
+    Each stage declares:
+        - ``reward``: a float-valued registry predicate giving the dense
+          shaping signal while the machine is IN that stage.
+        - ``advance_when``: a bool-valued registry predicate; the FIRST step it
+          returns True the machine awards ``bonus`` once and advances to the
+          next stage. Phases only ever move forward (no regression), matching
+          curriculum semantics and giving a stable, non-oscillating signal.
+        - ``bonus``: a one-time scalar added on the transition out of the stage
+          (default 0.0).
+
+    The last stage has no ``advance_when`` gate (the task is "done" there for
+    reward purposes; episode termination is a separate ``success`` predicate).
+    Per step the emitted reward is ``current_stage.reward(sim) +
+    (bonus if this step advanced else 0.0)``.
+    """
+
+    def __init__(
+        self,
+        stages: list[tuple[RewardTerm, BoolPredicate | None, float]],
+    ) -> None:
+        self._stages = stages
+        self._phase = 0
+
+    def reset(self) -> None:
+        self._phase = 0
+
+    @property
+    def phase(self) -> int:
+        """Current stage index (0-based). Exposed for logging / tests."""
+        return self._phase
+
+    def __call__(self, sim: SimEngine) -> float:
+        if not self._stages:
+            return 0.0
+        phase = min(self._phase, len(self._stages) - 1)
+        reward_fn, advance_fn, bonus = self._stages[phase]
+        r = float(reward_fn(sim))
+        # Advance (and award the one-time bonus) only if there IS a next stage
+        # and this stage declares a gate that now fires.
+        if self._phase < len(self._stages) - 1 and advance_fn is not None and bool(advance_fn(sim)):
+            self._phase += 1
+            return r + float(bonus)
+        return r
+
+
+def _staged_reward(stages: list[Any]) -> RewardTerm:
+    """Factory: compile a declared stage list into a :class:`_StagedReward`.
+
+    This is the single new primitive that turns the stateless DSL into a
+    declarative phase machine. It recursively compiles each stage's ``reward``
+    and ``advance_when`` through :func:`make_predicate`, so the whole thing
+    stays inside the closed-registry / no-``eval`` safety contract: a spec can
+    only ever reference predicates that already exist in the registry.
+
+    Args:
+        stages: Ordered list of stage dicts. Each stage::
+
+            {
+                "reward": {"predicate": <float-term name>, **kwargs},
+                "advance_when": {"predicate": <bool-pred name>, **kwargs},  # omit on last stage
+                "bonus": <float>,   # optional, default 0.0
+            }
+
+    Returns:
+        A callable+resettable :class:`_StagedReward`.
+
+    Raises:
+        ValueError: stages is not a non-empty list, a stage is malformed, a
+            non-final stage omits ``advance_when``, or ``bonus`` is non-numeric.
+        TypeError: surfaced from :func:`make_predicate` for bad sub-kwargs.
+    """
+    if not isinstance(stages, list) or not stages:
+        raise ValueError("staged_reward: 'stages' must be a non-empty list of stage dicts")
+
+    compiled: list[tuple[RewardTerm, BoolPredicate | None, float]] = []
+    n = len(stages)
+    for i, stage in enumerate(stages):
+        if not isinstance(stage, dict):
+            raise ValueError(f"staged_reward: stage[{i}] must be a dict, got {type(stage).__name__}")
+        unknown = set(stage.keys()) - {"reward", "advance_when", "bonus"}
+        if unknown:
+            raise ValueError(
+                f"staged_reward: stage[{i}] has unknown keys {sorted(unknown)}; "
+                "allowed: reward, advance_when, bonus"
+            )
+
+        reward_call = stage.get("reward")
+        if not isinstance(reward_call, dict) or "predicate" not in reward_call:
+            raise ValueError(
+                f"staged_reward: stage[{i}].reward must be a predicate-call dict "
+                "like {predicate: distance_neg, body_a: ..., body_b: ...}"
+            )
+        reward_name = reward_call["predicate"]
+        reward_kwargs = {k: v for k, v in reward_call.items() if k != "predicate"}
+        reward_fn = make_predicate(reward_name, **reward_kwargs)
+
+        advance_call = stage.get("advance_when")
+        advance_fn: BoolPredicate | None
+        if advance_call is None:
+            if i != n - 1:
+                raise ValueError(
+                    f"staged_reward: stage[{i}] is not the final stage and must declare "
+                    "'advance_when' (a bool predicate gating the transition to the next stage)"
+                )
+            advance_fn = None
+        else:
+            if not isinstance(advance_call, dict) or "predicate" not in advance_call:
+                raise ValueError(
+                    f"staged_reward: stage[{i}].advance_when must be a predicate-call dict "
+                    "like {predicate: distance_less_than, body_a: ..., body_b: ..., threshold: ...}"
+                )
+            advance_name = advance_call["predicate"]
+            advance_kwargs = {k: v for k, v in advance_call.items() if k != "predicate"}
+            advance_fn = make_predicate(advance_name, **advance_kwargs)
+
+        bonus_raw = stage.get("bonus", 0.0)
+        if isinstance(bonus_raw, bool) or not isinstance(bonus_raw, (int, float)):
+            raise ValueError(f"staged_reward: stage[{i}].bonus must be a number, got {bonus_raw!r}")
+
+        compiled.append((reward_fn, advance_fn, float(bonus_raw)))
+
+    return _StagedReward(compiled)
+
+
 # Registry
 
 PREDICATE_REGISTRY: dict[str, PredicateFactory] = {
@@ -574,6 +730,8 @@ def term(_sim: SimEngine) -> float:
     "distance_neg": _distance_neg,
     "joint_progress": _joint_progress,
     "constant": _constant,
+    # stateful (phase machine)
+    "staged_reward": _staged_reward,
 }
 
 
@@ -634,6 +792,7 @@ def make_predicate(name: str, **kwargs: Any) -> Callable[[SimEngine], Any]:
     "BoolPredicate",
     "PredicateFactory",
     "RewardTerm",
+    "StatefulRewardTerm",
     "make_predicate",
     "register_predicate",
 ]

strands_robots/training/rl/__init__.py

@@ -22,6 +22,8 @@
 
 from strands_robots.training.rl.base_algo import BaseRLAlgo, RLTrainSpec
 from strands_robots.training.rl.env import SimEnv
+from strands_robots.training.rl.gym_env import GymSimEnv
+from strands_robots.training.rl.vec_env import VecSimEnv
 from strands_robots.training.rl.fast_sac import FastSacTrainer
 from strands_robots.training.rl.normalization import EmpiricalNormalization
 from strands_robots.training.rl.ppo import PpoTrainer
@@ -34,5 +36,7 @@
     "FastSacTrainer",
     "SimpleReplayBuffer",
     "SimEnv",
+    "GymSimEnv",
+    "VecSimEnv",
     "EmpiricalNormalization",
 ]