diff --git a/src/openenv/core/rubrics/base.py b/src/openenv/core/rubrics/base.py index 0e2753047..96bf88019 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 6a62aceb5..db94ee10b 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 15f9cfb3c..7e662d70a 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 433f9ca06..8d3efed2d 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 38ea025e9..ae4ac3a8f 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 16822e9c2..4805c5e02 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()