| | import logging |
| | import threading |
| | import time |
| |
|
| | from openpi_client.runtime import agent as _agent |
| | from openpi_client.runtime import environment as _environment |
| | from openpi_client.runtime import subscriber as _subscriber |
| |
|
| |
|
| | class Runtime: |
| | """The core module orchestrating interactions between key components of the system.""" |
| |
|
| | def __init__( |
| | self, |
| | environment: _environment.Environment, |
| | 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._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.""" |
| | for _ in range(self._num_episodes): |
| | self._run_episode() |
| |
|
| | |
| | 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 mark_episode_complete(self) -> None: |
| | """Marks the end of an episode.""" |
| | self._in_episode = False |
| |
|
| | def _run_episode(self) -> None: |
| | """Runs a single episode.""" |
| | logging.info("Starting episode...") |
| | self._environment.reset() |
| | self._agent.reset() |
| | for subscriber in self._subscribers: |
| | subscriber.on_episode_start() |
| |
|
| | 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._in_episode: |
| | self._step() |
| | self._episode_steps += 1 |
| |
|
| | |
| | now = time.time() |
| | dt = now - last_step_time |
| | if dt < step_time: |
| | time.sleep(step_time - dt) |
| | last_step_time = time.time() |
| | else: |
| | last_step_time = now |
| |
|
| | logging.info("Episode completed.") |
| | for subscriber in self._subscribers: |
| | subscriber.on_episode_end() |
| |
|
| | def _step(self) -> None: |
| | """A single step of the runtime loop.""" |
| | observation = self._environment.get_observation() |
| | action = self._agent.get_action(observation) |
| | self._environment.apply_action(action) |
| |
|
| | for subscriber in self._subscribers: |
| | subscriber.on_step(observation, action) |
| |
|
| | if self._environment.is_episode_complete() or ( |
| | self._max_episode_steps > 0 and self._episode_steps >= self._max_episode_steps |
| | ): |
| | self.mark_episode_complete() |
| |
|
