Support multi-episode runs and max episode steps in runtime (#167)

Author: jimmyt857

examples/aloha_real/env.py

@@ -21,7 +21,7 @@ def reset(self) -> None:
         self._ts = self._env.reset()
 
     @override
-    def done(self) -> bool:
+    def is_episode_complete(self) -> bool:
         return False
 
     @override

examples/aloha_real/main.py

@@ -17,6 +17,9 @@ class Args:
 
     action_horizon: int = 25
 
+    num_episodes: int = 1
+    max_episode_steps: int = 1000
+
 
 def main(args: Args) -> None:
     runtime = _runtime.Runtime(
@@ -32,6 +35,8 @@ def main(args: Args) -> None:
         ),
         subscribers=[],
         max_hz=50,
+        num_episodes=args.num_episodes,
+        max_episode_steps=args.max_episode_steps,
     )
 
     runtime.run()

examples/aloha_sim/env.py

@@ -27,7 +27,7 @@ def reset(self) -> None:
         self._episode_reward = 0.0
 
     @override
-    def done(self) -> bool:
+    def is_episode_complete(self) -> bool:
         return self._done
 
     @override

packages/openpi-client/src/openpi_client/runtime/environment.py

@@ -16,11 +16,11 @@ def reset(self) -> None:
         """
 
     @abc.abstractmethod
-    def done(self) -> bool:
-        """Allow the environment to signal that the task is done.
+    def is_episode_complete(self) -> bool:
+        """Allow the environment to signal that the episode is complete.
 
-        This will be called after each step. It should return `True` if the task is
-        done (either successfully or unsuccessfully), and `False` otherwise.
+        This will be called after each step. It should return `True` if the episode is
+        complete (either successfully or unsuccessfully), and `False` otherwise.
         """
 
     @abc.abstractmethod

packages/openpi-client/src/openpi_client/runtime/runtime.py

@@ -16,41 +16,52 @@ def __init__(
         agent: _agent.Agent,
         subscribers: list[_subscriber.Subscriber],
         max_hz: float = 0,
+        num_episodes: int = 1,
+        max_episode_steps: int = 0,
     ) -> None:
         self._environment = environment
         self._agent = agent
         self._subscribers = subscribers
         self._max_hz = max_hz
+        self._num_episodes = num_episodes
+        self._max_episode_steps = max_episode_steps
 
-        self._running = False
+        self._in_episode = False
+        self._episode_steps = 0
 
     def run(self) -> None:
         """Runs the runtime loop continuously until stop() is called or the environment is done."""
-        self._loop()
+        for _ in range(self._num_episodes):
+            self._run_episode()
+
+        # Final reset, this is important for real environments to move the robot to its home position.
+        self._environment.reset()
 
     def run_in_new_thread(self) -> threading.Thread:
         """Runs the runtime loop in a new thread."""
         thread = threading.Thread(target=self.run)
         thread.start()
         return thread
 
-    def stop(self) -> None:
-        """Stops the runtime loop."""
-        self._running = False
+    def mark_episode_complete(self) -> None:
+        """Marks the end of an episode."""
+        self._in_episode = False
 
-    def _loop(self) -> None:
-        """The runtime loop."""
+    def _run_episode(self) -> None:
+        """Runs a single episode."""
         logging.info("Starting episode...")
         self._environment.reset()
         for subscriber in self._subscribers:
             subscriber.on_episode_start()
 
-        self._running = True
+        self._in_episode = True
+        self._episode_steps = 0
         step_time = 1 / self._max_hz if self._max_hz > 0 else 0
         last_step_time = time.time()
 
-        while self._running:
+        while self._in_episode:
             self._step()
+            self._episode_steps += 1
 
             # Sleep to maintain the desired frame rate
             now = time.time()
@@ -74,5 +85,7 @@ def _step(self) -> None:
         for subscriber in self._subscribers:
             subscriber.on_step(observation, action)
 
-        if self._environment.done():
-            self.stop()
+        if self._environment.is_episode_complete() or (
+            self._max_episode_steps > 0 and self._episode_steps >= self._max_episode_steps
+        ):
+            self.mark_episode_complete()