success to run human interaction gui with aloha_sim

[csjiyw]

idea is to use GlfwWindow.update to show

add

  def launch_passive(self, environment_loader, policy=None):
      if environment_loader is None:
        raise ValueError('"environment_loader" argument is required.')
      if callable(environment_loader):
        self._environment_loader = environment_loader
      else:
        self._environment_loader = lambda: environment_loader
      self._policy = policy
      self._load_environment(zoom_to_scene=True)
      def tick():
        self._viewport.set_size(*self._window.shape)
        self._tick()
        return self._renderer.pixels
      self.tick_f=tick  

to dm_control/viewer/application.py

modify dm_control/viewer/runtime.py

  def _step(self):
   
    finished = True
    with self._error_logger:
      if self._policy:
        action = self._policy(self._time_step)
        self._time_step = self._env.step(action)
        self._last_action = action
        finished = self._time_step.last()
      else:
        action = self._default_action
        # finished=False
        # no action we use outer action
      
    return finished or self._error_logger.errors_found

add

def launch_passive(environment_loader, policy=None, title='Explorer', width=1024,
           height=768):

  app = application.Application(title=title, width=width, height=height)
  app.launch_passive(environment_loader=environment_loader, policy=policy)
  return app

to dm_control/viewer/init.py

modify gym_aloha/env.py

class AlohaEnv(gym.Env):
    metadata = {"render_modes": ["rgb_array","human"], "render_fps": 15}

    def __init__(
        self,
        task,
        obs_type="pixels",
        render_mode="human",
        observation_width=640,
        observation_height=480,
        visualization_width=640,
        visualization_height=480,
    ):
        super().__init__()
        self.task = task
        self.obs_type = obs_type
        self.render_mode = render_mode
        self.observation_width = observation_width
        self.observation_height = observation_height
        self.visualization_width = visualization_width
        self.visualization_height = visualization_height

        self._env = self._make_env_task(self.task)

        if self.obs_type == "state":
            raise NotImplementedError()
            self.observation_space = spaces.Box(
                low=np.array([0] * len(JOINTS)),  # ???
                high=np.array([255] * len(JOINTS)),  # ???
                dtype=np.float64,
            )
        elif self.obs_type == "pixels":
            self.observation_space = spaces.Dict(
                {
                    "top": spaces.Box(
                        low=0,
                        high=255,
                        shape=(self.observation_height, self.observation_width, 3),
                        dtype=np.uint8,
                    )
                }
            )
        elif self.obs_type == "pixels_agent_pos":
            self.observation_space = spaces.Dict(
                {
                    "pixels": spaces.Dict(
                        {
                            "top": spaces.Box(
                                low=0,
                                high=255,
                                shape=(self.observation_height, self.observation_width, 3),
                                dtype=np.uint8,
                            )
                        }
                    ),
                    "agent_pos": spaces.Box(
                        low=-1000.0,
                        high=1000.0,
                        shape=(len(JOINTS),),
                        dtype=np.float64,
                    ),
                }
            )

        self.action_space = spaces.Box(low=-1, high=1, shape=(len(ACTIONS),), dtype=np.float32)
        self.app=None
        

    def render(self):
        if self.render_mode=="human":
            if self.app==None:
                self.app=viewer.launch_passive(self._env)
            else:
                self.app._window.update(self.app.tick_f)
        else:
            return self._render(visualize=True)

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

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:
    

    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()

        # 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 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._environment._gym.render()
        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

            # Sleep to maintain the desired frame rate
            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)
        self._environment._gym.render()



        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()

modify examples/aloha_sim/main.py

def main(args: Args) -> None:
    runtime = _runtime.Runtime(
        environment=_env.AlohaSimEnvironment(
            task=args.task,
            seed=args.seed,
        ),
        agent=_policy_agent.PolicyAgent(
            policy=action_chunk_broker.ActionChunkBroker(
                policy=_websocket_client_policy.WebsocketClientPolicy(
                    host=args.host,
                    port=args.port,
                ),
                action_horizon=args.action_horizon,
            )
        ),
        # subscribers=[
        #     _saver.VideoSaver(args.out_dir),
        # ],
        max_hz=50,
    )

    runtime.run()

[svlevine]

Thanks for the example!