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()
idea is to use GlfwWindow.update to show
add
to dm_control/viewer/application.py
modify dm_control/viewer/runtime.py
add
to dm_control/viewer/init.py
modify gym_aloha/env.py
modify packages/openpi-client/src/openpi_client/runtime/runtime.py
modify examples/aloha_sim/main.py