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/base.py

@@ -245,6 +245,17 @@ def list_robots(self) -> list[str]:
         """
         ...
 
+    def robot_action_keys(self, robot_name: str) -> list[str]:
+        """DRIVEABLE action keys a policy should emit (default: all joint names).
+
+        Backends with a distinct actuator set (e.g. MuJoCo, where passive/mimic
+        joints have no actuator and tendons are actuators-not-joints) SHOULD
+        override this to return the keys :meth:`send_action` actually resolves.
+        The default mirrors :meth:`robot_joint_names` so existing backends keep
+        their behaviour.
+        """
+        return self.robot_joint_names(robot_name)
+
     @abstractmethod
     def robot_joint_names(self, robot_name: str) -> list[str]:
         """Return ordered joint names for ``robot_name``.
@@ -821,7 +832,7 @@ def run_policy(
             # Caller is responsible for policy.set_robot_state_keys(...) if needed,
             # but we set it here defensively so the semantics match the provider path.
             policy = policy_object
-        policy.set_robot_state_keys(self.robot_joint_names(robot_name))
+        policy.set_robot_state_keys(self.robot_action_keys(robot_name))
         self.bind_policy_sim_context(policy, robot_name)
 
         # Auto-install any action controller this policy needs to run correctly
@@ -1449,14 +1460,17 @@ def eval_policy(
         for cuRobo / MoveIt2 - the issue #300 contract). Without it the eval ran
         such a policy with an empty goal and reported a meaningless success rate.
         """
-        if not robot_name:
-            return {
-                "status": "error",
-                "content": [{"text": "eval_policy requires 'robot_name'."}],
-            }
+        # Resolve an optional robot_name to a concrete one (None -> the single
+        # robot), matching run_policy's ergonomics. Previously eval_policy hard-
+        # required robot_name even with exactly one robot in the world, an API
+        # inconsistency with its sibling run_policy.
         robots = self.list_robots()
         if not robots:
             return {"status": "error", "content": [{"text": "No robots in sim. Add one first."}]}
+        try:
+            robot_name = self._resolve_single_robot(robot_name)
+        except ValueError as exc:
+            return {"status": "error", "content": [{"text": str(exc)}]}
         if robot_name not in robots:
             return {
                 "status": "error",
@@ -1486,7 +1500,7 @@ def eval_policy(
             # set robot_state_keys; we set defensively so semantics match the
             # provider path.
             policy = policy_object
-        policy.set_robot_state_keys(self.robot_joint_names(resolved_robot))
+        policy.set_robot_state_keys(self.robot_action_keys(resolved_robot))
         self.bind_policy_sim_context(policy, resolved_robot)
 
         return PolicyRunner(self).evaluate(
@@ -1637,7 +1651,7 @@ def evaluate_benchmark(
             }
 
         policy = create_policy(policy_provider, **(policy_config or {}))
-        policy.set_robot_state_keys(self.robot_joint_names(resolved_robot))
+        policy.set_robot_state_keys(self.robot_action_keys(resolved_robot))
         self.bind_policy_sim_context(policy, resolved_robot)
 
         return PolicyRunner(self).evaluate(

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/mujoco/simulation.py

@@ -1096,6 +1096,40 @@ def robot_joint_names(self, robot_name: str) -> list[str]:
             return []
         return list(self._world.robots[robot_name].joint_names)
 
+    def robot_action_keys(self, robot_name: str) -> list[str]:
+        """Ordered DRIVEABLE action keys for ``robot_name`` (actuator short names).
+
+        Overrides the ABC default. ``robot_joint_names`` returns ALL joints,
+        including passive/mimic followers (e.g. gripper finger joints) that have
+        no driving actuator, and it omits tendon-driven actuators (a grasp tendon
+        is an actuator, not a joint). A policy must key its actions by what the
+        engine can actually drive, which is the set of ACTUATORS scoped to the
+        robot's namespace - exactly what :meth:`send_action` resolves against.
+
+        Returns the short (namespace-stripped) actuator names in model order.
+        Falls back to :meth:`robot_joint_names` only if the model is unavailable.
+        """
+        if self._world is None or self._world._model is None or robot_name not in self._world.robots:
+            return []
+        mj = self._mj
+        model = self._world._model
+        # namespace already carries its trailing "/" (e.g. "xarm7/"); normalize
+        # so we don't double it. Empty namespace -> unprefixed single-robot scene.
+        raw_ns = self._world.robots[robot_name].namespace or ""
+        namespace = raw_ns if raw_ns.endswith("/") or raw_ns == "" else raw_ns + "/"
+        keys: list[str] = []
+        for ai in range(model.nu):
+            aname = mj.mj_id2name(model, mj.mjtObj.mjOBJ_ACTUATOR, ai)
+            if not aname:
+                continue
+            if namespace and aname.startswith(namespace):
+                keys.append(aname[len(namespace):])
+            elif not namespace and "/" not in aname:
+                keys.append(aname)
+        # Fall back to joint names only if the robot has no namespaced actuators
+        # (degenerate single-robot scenes without a prefix).
+        return keys or list(self._world.robots[robot_name].joint_names)
+
     def bind_policy_sim_context(self, policy: Any, robot_name: str) -> None:
         """Hand the compiled MjModel + robot namespace to policies that opt in.
 

strands_robots/simulation/mujoco/spec_builder.py

@@ -125,11 +125,20 @@ def _target_quat(position: list[float], target: list[float]) -> list[float] | No
         return None
     fwd /= flen
 
+    # Choose an up vector not parallel to the view direction. World +Z is the
+    # natural up, but a vertical (top-down / bottom-up) camera has forward
+    # parallel to +Z, making cross(fwd, up) degenerate. In that case fall back
+    # to world +Y as up so a straight-down camera still gets a valid, sensible
+    # orientation (this is the common "overhead view" case, NOT an error).
     up = np.array([0.0, 0.0, 1.0])
     right = np.cross(fwd, up)
     rlen = float(np.linalg.norm(right))
     if rlen < 1e-9:
-        return None
+        up = np.array([0.0, 1.0, 0.0])
+        right = np.cross(fwd, up)
+        rlen = float(np.linalg.norm(right))
+        if rlen < 1e-9:
+            return None
     right /= rlen
     image_up = np.cross(right, fwd)
     image_up /= float(np.linalg.norm(image_up))

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",
 ]