From 9cb55f4e7f2bcc7beca4140e284e21d37698e927 Mon Sep 17 00:00:00 2001 From: TheDeadcoder Date: Fri, 30 Oct 2026 16:07:13 +0600 Subject: [PATCH] rubrics: reset child rubrics when a rubric is reset Rubric.reset() now resets every registered child, so stateful rubrics nested in Sequential, Gate, WeightedSum, RubricList and RubricDict (such as TrajectoryRubric) no longer keep the previous episode. TrajectoryRubric.reset() calls super().reset() to keep the same contract. Signed-off-by: TheDeadcoder --- src/openenv/core/rubrics/base.py | 9 +++- src/openenv/core/rubrics/trajectory.py | 1 + .../test_async_environment_integration.py | 26 +++++++++- tests/core/test_rubrics/test_containers.py | 48 +++++++++++++++++++ .../test_environment_integration.py | 17 ++++++- .../test_rubrics/test_trajectory_rubric.py | 10 ++++ 6 files changed, 107 insertions(+), 4 deletions(-) diff --git a/src/openenv/core/rubrics/base.py b/src/openenv/core/rubrics/base.py index 0e2753047f..96bf880195 100644 --- a/src/openenv/core/rubrics/base.py +++ b/src/openenv/core/rubrics/base.py @@ -234,8 +234,13 @@ def get_rubric(self, path: str) -> "Rubric": return current def reset(self) -> None: - """Reset any internal state. Override in subclasses if needed.""" - pass + """Reset internal state, including every child rubric. + + Override to clear your own state and call `super().reset()` so child + rubrics are reset too. + """ + for child in self.children(): + child.reset() def state_dict(self) -> Dict[str, Any]: """Serialize rubric configuration for checkpointing.""" diff --git a/src/openenv/core/rubrics/trajectory.py b/src/openenv/core/rubrics/trajectory.py index 6a62aceb5a..db94ee10bb 100644 --- a/src/openenv/core/rubrics/trajectory.py +++ b/src/openenv/core/rubrics/trajectory.py @@ -118,6 +118,7 @@ def compute_step_rewards(self) -> List[float]: def reset(self) -> None: """Clear accumulated trajectory. Call on env.reset().""" + super().reset() self._trajectory = [] @property diff --git a/tests/core/test_rubrics/test_async_environment_integration.py b/tests/core/test_rubrics/test_async_environment_integration.py index 15f9cfb3c1..7e662d70a0 100644 --- a/tests/core/test_rubrics/test_async_environment_integration.py +++ b/tests/core/test_rubrics/test_async_environment_integration.py @@ -14,7 +14,7 @@ import pytest from openenv.core.env_server.interfaces import Environment from openenv.core.env_server.types import Action, Observation, State -from openenv.core.rubrics import Rubric +from openenv.core.rubrics import Rubric, Sequential, TrajectoryRubric # Test fixtures - using Pydantic models @@ -72,6 +72,16 @@ async def forward(self, action, observation): return (score1 + score2) / 2 +class RecordingTrajectoryRubric(TrajectoryRubric): + """Trajectory rubric that only records steps.""" + + def score_trajectory(self, trajectory): + return 1.0 + + def compute_step_rewards(self): + return [1.0] * len(self._trajectory) + + class AsyncEnvironment(Environment[MockAction, MockObservation, MockState]): """Async environment implementation for testing.""" @@ -214,6 +224,20 @@ async def reset_async(self): await env.reset_async() assert rubric.step_count == 0 # Reset clears state + @pytest.mark.asyncio + async def test_composed_rubric_resets_children_on_reset_async(self): + """reset_async clears trajectory rubrics nested in containers.""" + trajectory = RecordingTrajectoryRubric() + env = AsyncEnvironment(rubric=Sequential(AsyncRubric(1.0), trajectory)) + + await env.reset_async() + await env.step_async(MockAction(content="a")) + await env.reset_async() + assert trajectory.trajectory == [] + + await env.step_async(MockAction(content="b")) + assert len(trajectory.trajectory) == 1 + @pytest.mark.asyncio async def test_async_rubric_introspection(self): """Can introspect async rubric from environment.""" diff --git a/tests/core/test_rubrics/test_containers.py b/tests/core/test_rubrics/test_containers.py index 433f9ca06e..8d3efed2d8 100644 --- a/tests/core/test_rubrics/test_containers.py +++ b/tests/core/test_rubrics/test_containers.py @@ -39,6 +39,21 @@ def forward(self, action: Any, observation: Any) -> float: return self.score +class EpisodeRubric(Rubric): + """Rubric that records the actions of the current episode.""" + + def __init__(self): + super().__init__() + self.actions = [] + + def forward(self, action: Any, observation: Any) -> float: + self.actions.append(action) + return 1.0 + + def reset(self) -> None: + self.actions = [] + + class TestSequential: """Test Sequential container.""" @@ -408,3 +423,36 @@ def test_nested_named_rubrics(self): assert "task.rubric_1" in paths # FixedRubric # Gate's child assert "task.rubric_0.rubric" in paths + + +class TestContainerReset: + """reset() reaches the children of every container.""" + + @pytest.mark.parametrize( + "make_container", + [ + lambda child: Sequential(child), + lambda child: Gate(child), + lambda child: WeightedSum([child], weights=[1.0]), + lambda child: RubricList([child]), + lambda child: RubricDict({"child": child}), + ], + ids=["sequential", "gate", "weighted_sum", "rubric_list", "rubric_dict"], + ) + def test_reset_resets_children(self, make_container): + child = EpisodeRubric() + container = make_container(child) + child("action", "obs") + + container.reset() + + assert child.actions == [] + + def test_reset_reaches_nested_children(self): + leaf = EpisodeRubric() + rubric = RubricDict({"task": Sequential(Gate(leaf), FixedRubric(1.0))}) + rubric["task"]("action", "obs") + + rubric.reset() + + assert leaf.actions == [] diff --git a/tests/core/test_rubrics/test_environment_integration.py b/tests/core/test_rubrics/test_environment_integration.py index 38ea025e94..ae4ac3a8f2 100644 --- a/tests/core/test_rubrics/test_environment_integration.py +++ b/tests/core/test_rubrics/test_environment_integration.py @@ -6,7 +6,7 @@ from openenv.core.env_server.interfaces import Environment from openenv.core.env_server.types import Action, Observation, State -from openenv.core.rubrics import Rubric, TrajectoryRubric +from openenv.core.rubrics import Rubric, TrajectoryRubric, WeightedSum # Test fixtures - using Pydantic models (not dataclasses) @@ -231,6 +231,21 @@ def test_multiple_episodes(self): assert ep1_len == 2 assert ep2_len == 1 # Reset cleared previous episode + def test_composed_rubric_resets_children_between_episodes(self): + """Env reset clears trajectory rubrics nested in containers.""" + trajectory = MockTrajectoryRubric() + env = SimpleEnvironment( + rubric=WeightedSum([FixedRubric(1.0), trajectory], weights=[0.5, 0.5]) + ) + + env.reset() + env.step(MockAction(content="a")) + env.reset() + assert trajectory.trajectory == [] + + env.step(MockAction(content="b")) + assert len(trajectory.trajectory) == 1 + def test_rubric_hooks_work(self): """Rubric hooks work through environment.""" rubric = FixedRubric(0.9) diff --git a/tests/core/test_rubrics/test_trajectory_rubric.py b/tests/core/test_rubrics/test_trajectory_rubric.py index 16822e9c25..4805c5e029 100644 --- a/tests/core/test_rubrics/test_trajectory_rubric.py +++ b/tests/core/test_rubrics/test_trajectory_rubric.py @@ -119,6 +119,16 @@ def test_reset_clears_trajectory(self): rubric.reset() assert len(rubric._trajectory) == 0 + def test_reset_resets_child_rubrics(self): + """reset() also resets child rubrics.""" + rubric = EqualCreditRubric() + rubric.history = EqualCreditRubric() + rubric.history(MockAction(), MockObservation(done=False)) + + rubric.reset() + + assert rubric.history.trajectory == [] + def test_trajectory_property_returns_copy(self): """trajectory property returns a copy.""" rubric = EqualCreditRubric()