diff --git a/packages/openpi-client/src/openpi_client/runtime/runtime.py b/packages/openpi-client/src/openpi_client/runtime/runtime.py index 9552be091a..4412f48070 100644 --- a/packages/openpi-client/src/openpi_client/runtime/runtime.py +++ b/packages/openpi-client/src/openpi_client/runtime/runtime.py @@ -62,7 +62,6 @@ def _run_episode(self) -> None: while self._in_episode: self._step() - self._episode_steps += 1 # Sleep to maintain the desired frame rate now = time.time() @@ -86,6 +85,7 @@ def _step(self) -> None: for subscriber in self._subscribers: subscriber.on_step(observation, action) + self._episode_steps += 1 if self._environment.is_episode_complete() or ( self._max_episode_steps > 0 and self._episode_steps >= self._max_episode_steps ): diff --git a/packages/openpi-client/src/openpi_client/runtime/runtime_test.py b/packages/openpi-client/src/openpi_client/runtime/runtime_test.py new file mode 100644 index 0000000000..653cf40a2b --- /dev/null +++ b/packages/openpi-client/src/openpi_client/runtime/runtime_test.py @@ -0,0 +1,51 @@ +from openpi_client.runtime import agent as _agent +from openpi_client.runtime import environment as _environment +from openpi_client.runtime.runtime import Runtime +import pytest + + +class _FakeEnvironment(_environment.Environment): + def __init__(self, complete_after: int | None = None) -> None: + self.applied_actions = 0 + self._complete_after = complete_after + + def reset(self) -> None: + pass + + def is_episode_complete(self) -> bool: + return self._complete_after is not None and self.applied_actions >= self._complete_after + + def get_observation(self) -> dict: + return {"step": self.applied_actions} + + def apply_action(self, action: dict) -> None: + self.applied_actions += 1 + + +class _FakeAgent(_agent.Agent): + def get_action(self, observation: dict) -> dict: + return {"step": observation["step"]} + + def reset(self) -> None: + pass + + +@pytest.mark.parametrize("max_episode_steps", [1, 2, 3]) +def test_max_episode_steps_counts_completed_actions(max_episode_steps: int) -> None: + environment = _FakeEnvironment() + runtime = Runtime(environment, _FakeAgent(), [], max_episode_steps=max_episode_steps) + + runtime.run() + + assert environment.applied_actions == max_episode_steps + assert runtime._episode_steps == max_episode_steps # noqa: SLF001 + + +def test_environment_can_complete_episode_before_step_limit() -> None: + environment = _FakeEnvironment(complete_after=2) + runtime = Runtime(environment, _FakeAgent(), [], max_episode_steps=10) + + runtime.run() + + assert environment.applied_actions == 2 + assert runtime._episode_steps == 2 # noqa: SLF001