@@ -554,6 +554,161 @@ def term(_sim: SimEngine) -> float:
554554 return term
555555
556556
557+ # Stateful reward terms (declarative phase machine)
558+ #
559+ # A plain RewardTerm is stateless: ``(SimEngine) -> float``. Some rewards need
560+ # memory across steps - a pick-place curriculum advances Reach -> Grasp ->
561+ # Transport -> Place, awards a one-time bonus on each transition, and only ever
562+ # moves forward. Rather than hardcode any specific task, we expose ONE
563+ # generic primitive, ``staged_reward``, that composes EXISTING registry
564+ # predicates into a phase machine. The task itself is then authored as data
565+ # (a spec dict / YAML) by a human or LLM - never as shipped code, and never via
566+ # ``eval`` (sub-predicates are compiled through :func:`make_predicate`, the same
567+ # closed-registry path as every other DSL call).
568+
569+
570+ class StatefulRewardTerm :
571+ """A reward term that carries per-episode state and must be ``reset()``.
572+
573+ Duck-typed by consumers: anything with ``__call__(sim) -> float`` AND a
574+ zero-arg ``reset()`` is treated as episode-stateful. ``SimEnv.reset`` and
575+ ``DeclarativeBenchmark.on_episode_start`` call ``reset()`` on any reward
576+ term that has it, so stateless plain-function terms are unaffected.
577+ """
578+
579+ def __call__ (self , sim : SimEngine ) -> float : # pragma: no cover - interface
580+ raise NotImplementedError
581+
582+ def reset (self ) -> None : # pragma: no cover - interface
583+ raise NotImplementedError
584+
585+
586+ class _StagedReward (StatefulRewardTerm ):
587+ """Monotonic multi-stage (phase-machine) reward built from sub-predicates.
588+
589+ Each stage declares:
590+ - ``reward``: a float-valued registry predicate giving the dense
591+ shaping signal while the machine is IN that stage.
592+ - ``advance_when``: a bool-valued registry predicate; the FIRST step it
593+ returns True the machine awards ``bonus`` once and advances to the
594+ next stage. Phases only ever move forward (no regression), matching
595+ curriculum semantics and giving a stable, non-oscillating signal.
596+ - ``bonus``: a one-time scalar added on the transition out of the stage
597+ (default 0.0).
598+
599+ The last stage has no ``advance_when`` gate (the task is "done" there for
600+ reward purposes; episode termination is a separate ``success`` predicate).
601+ Per step the emitted reward is ``current_stage.reward(sim) +
602+ (bonus if this step advanced else 0.0)``.
603+ """
604+
605+ def __init__ (
606+ self ,
607+ stages : list [tuple [RewardTerm , BoolPredicate | None , float ]],
608+ ) -> None :
609+ self ._stages = stages
610+ self ._phase = 0
611+
612+ def reset (self ) -> None :
613+ self ._phase = 0
614+
615+ @property
616+ def phase (self ) -> int :
617+ """Current stage index (0-based). Exposed for logging / tests."""
618+ return self ._phase
619+
620+ def __call__ (self , sim : SimEngine ) -> float :
621+ if not self ._stages :
622+ return 0.0
623+ phase = min (self ._phase , len (self ._stages ) - 1 )
624+ reward_fn , advance_fn , bonus = self ._stages [phase ]
625+ r = float (reward_fn (sim ))
626+ # Advance (and award the one-time bonus) only if there IS a next stage
627+ # and this stage declares a gate that now fires.
628+ if self ._phase < len (self ._stages ) - 1 and advance_fn is not None and bool (advance_fn (sim )):
629+ self ._phase += 1
630+ return r + float (bonus )
631+ return r
632+
633+
634+ def _staged_reward (stages : list [Any ]) -> RewardTerm :
635+ """Factory: compile a declared stage list into a :class:`_StagedReward`.
636+
637+ This is the single new primitive that turns the stateless DSL into a
638+ declarative phase machine. It recursively compiles each stage's ``reward``
639+ and ``advance_when`` through :func:`make_predicate`, so the whole thing
640+ stays inside the closed-registry / no-``eval`` safety contract: a spec can
641+ only ever reference predicates that already exist in the registry.
642+
643+ Args:
644+ stages: Ordered list of stage dicts. Each stage::
645+
646+ {
647+ "reward": {"predicate": <float-term name>, **kwargs},
648+ "advance_when": {"predicate": <bool-pred name>, **kwargs}, # omit on last stage
649+ "bonus": <float>, # optional, default 0.0
650+ }
651+
652+ Returns:
653+ A callable+resettable :class:`_StagedReward`.
654+
655+ Raises:
656+ ValueError: stages is not a non-empty list, a stage is malformed, a
657+ non-final stage omits ``advance_when``, or ``bonus`` is non-numeric.
658+ TypeError: surfaced from :func:`make_predicate` for bad sub-kwargs.
659+ """
660+ if not isinstance (stages , list ) or not stages :
661+ raise ValueError ("staged_reward: 'stages' must be a non-empty list of stage dicts" )
662+
663+ compiled : list [tuple [RewardTerm , BoolPredicate | None , float ]] = []
664+ n = len (stages )
665+ for i , stage in enumerate (stages ):
666+ if not isinstance (stage , dict ):
667+ raise ValueError (f"staged_reward: stage[{ i } { type (stage ).__name__ } )
668+ unknown = set (stage .keys ()) - {"reward" , "advance_when" , "bonus" }
669+ if unknown :
670+ raise ValueError (
671+ f"staged_reward: stage[{ i } { sorted (unknown )}
672+ )
673+
674+ reward_call = stage .get ("reward" )
675+ if not isinstance (reward_call , dict ) or "predicate" not in reward_call :
676+ raise ValueError (
677+ f"staged_reward: stage[{ i }
678+ "like {predicate: distance_neg, body_a: ..., body_b: ...}"
679+ )
680+ reward_name = reward_call ["predicate" ]
681+ reward_kwargs = {k : v for k , v in reward_call .items () if k != "predicate" }
682+ reward_fn = make_predicate (reward_name , ** reward_kwargs )
683+
684+ advance_call = stage .get ("advance_when" )
685+ advance_fn : BoolPredicate | None
686+ if advance_call is None :
687+ if i != n - 1 :
688+ raise ValueError (
689+ f"staged_reward: stage[{ i }
690+ "'advance_when' (a bool predicate gating the transition to the next stage)"
691+ )
692+ advance_fn = None
693+ else :
694+ if not isinstance (advance_call , dict ) or "predicate" not in advance_call :
695+ raise ValueError (
696+ f"staged_reward: stage[{ i }
697+ "like {predicate: distance_less_than, body_a: ..., body_b: ..., threshold: ...}"
698+ )
699+ advance_name = advance_call ["predicate" ]
700+ advance_kwargs = {k : v for k , v in advance_call .items () if k != "predicate" }
701+ advance_fn = make_predicate (advance_name , ** advance_kwargs )
702+
703+ bonus_raw = stage .get ("bonus" , 0.0 )
704+ if isinstance (bonus_raw , bool ) or not isinstance (bonus_raw , (int , float )):
705+ raise ValueError (f"staged_reward: stage[{ i } { bonus_raw !r} )
706+
707+ compiled .append ((reward_fn , advance_fn , float (bonus_raw )))
708+
709+ return _StagedReward (compiled )
710+
711+
557712# Registry
558713
559714PREDICATE_REGISTRY : dict [str , PredicateFactory ] = {
@@ -574,6 +729,8 @@ def term(_sim: SimEngine) -> float:
574729 "distance_neg" : _distance_neg ,
575730 "joint_progress" : _joint_progress ,
576731 "constant" : _constant ,
732+ # stateful (phase machine)
733+ "staged_reward" : _staged_reward ,
577734}
578735
579736
@@ -634,6 +791,7 @@ def make_predicate(name: str, **kwargs: Any) -> Callable[[SimEngine], Any]:
634791 "BoolPredicate" ,
635792 "PredicateFactory" ,
636793 "RewardTerm" ,
794+ "StatefulRewardTerm" ,
637795 "make_predicate" ,
638796 "register_predicate" ,
639797]
0 commit comments