diff --git a/robohive/envs/env_base.py b/robohive/envs/env_base.py index 6e9afac2..a0d36988 100644 --- a/robohive/envs/env_base.py +++ b/robohive/envs/env_base.py @@ -534,13 +534,33 @@ def reset(self, reset_qpos=None, reset_qvel=None, seed=None, **kwargs): # return self.step(a) - def close(self): + def close(self, close_hardware=True): """ Clean up the environment - """ - if self.sim: + close_hardware: True (default) closes the robot's persistent hardware connection. + False leaves it open so a later env.make(is_hardware=True) can + reuse the same connection instead of reconnecting. + """ + self._explicitly_closed = True + if getattr(self, 'robot', None): + self.robot.close(close_hardware=close_hardware) + if getattr(self, 'sim_obsd', None) and self.sim_obsd is not self.sim: + self.sim_obsd.close() + if getattr(self, 'sim', None): self.sim.close() + # Warn (don't auto-cleanup) if a hardware-backed env is garbage collected without + # close() ever being called. Mirrors Robot.__del__: cleanup must stay an explicit, + # deliberate action (esp. for close_hardware=False reuse), not something GC timing decides. + def __del__(self): + if getattr(self, 'robot', None) and getattr(self.robot, 'is_hardware', False) \ + and not getattr(self, '_explicitly_closed', False): + raise RuntimeWarning( + f"RoboHive:> {type(self).__name__} is being garbage collected without close() ever being called. " + "If hardware is still connected, this leaves the persistent connection dangling. " + "Call env.close() (or env.close(close_hardware=False) to intentionally keep the " + "hardware connection alive) before letting the env go out of scope.") + @property def dt(self): return self.sim.model.opt.timestep * self.frame_skip @@ -696,6 +716,16 @@ def mj_render(self): self.sim.renderer.render_to_window() + @property + def viewer_exit_requested(self): + """ + True once the onscreen viewer has been asked to close (Escape / window close). + Always False for offscreen/none rendering, which has no interactive window; + use a KeyboardInterrupt (Ctrl+C) to break out of those rollouts instead. + """ + return bool(self.sim) and self.sim.renderer.exit_requested + + def viewer_setup(self, distance=2.5, azimuth=90, elevation=-30, lookat=None, render_actuator=None, render_tendon=None): """ Setup the default camera @@ -741,8 +771,17 @@ def examine_policy(self, self.mujoco_render_frames = False # start rollouts + # stop_requested becomes True either when the onscreen viewer is closed + # (Escape / window close, only possible with render=='onscreen') or when the + # user hits Ctrl+C (the only way to interrupt offscreen/none rollouts, which + # have no interactive window). Either way we stop early but still finalize + # and return whatever paths were completed, instead of crashing or looping + # invisibly to the end of `horizon`/`num_episodes`. paths = [] + stop_requested = False for ep in range(num_episodes): + if stop_requested: + break ep_t0 = timer.time() observations=[] actions=[] @@ -755,27 +794,35 @@ def examine_policy(self, done = False t = 0 ep_rwd = 0.0 - while t < horizon and done is False: - a = policy.get_action(o)[0] if mode == 'exploration' else policy.get_action(o)[1]['evaluation'] - next_o, rwd, done, *_, env_info = self.step(a) - ep_rwd += rwd - # render offscreen visuals - if render =='offscreen': - curr_frame = self.sim.renderer.render_offscreen( - width=frame_size[0], - height=frame_size[1], - camera_id=camera_name, - device_id=device_id) - - frames[t,:,:,:] = curr_frame - prompt(t, end=', ', flush=True, type=Prompt.INFO) - observations.append(o) - actions.append(a) - rewards.append(rwd) - # agent_infos.append(agent_info) - env_infos.append(env_info) - o = next_o - t = t+1 + try: + while t < horizon and done is False and not self.viewer_exit_requested: + a = policy.get_action(o)[0] if mode == 'exploration' else policy.get_action(o)[1]['evaluation'] + next_o, rwd, done, *_, env_info = self.step(a) + ep_rwd += rwd + # render offscreen visuals + if render =='offscreen': + curr_frame = self.sim.renderer.render_offscreen( + width=frame_size[0], + height=frame_size[1], + camera_id=camera_name, + device_id=device_id) + + frames[t,:,:,:] = curr_frame + prompt(t, end=', ', flush=True, type=Prompt.INFO) + observations.append(o) + actions.append(a) + rewards.append(rwd) + # agent_infos.append(agent_info) + env_infos.append(env_info) + o = next_o + t = t+1 + except KeyboardInterrupt: + prompt("\nRollout interrupted by user (Ctrl+C)", type=Prompt.WARN) + stop_requested = True + + if self.viewer_exit_requested: + prompt("Rollout stopped: viewer was closed", type=Prompt.WARN) + stop_requested = True prompt("Total reward = %3.3f, Total time = %2.3f" % (ep_rwd, timer.time()-ep_t0), type=Prompt.INFO) path = dict( @@ -788,14 +835,14 @@ def examine_policy(self, ) paths.append(path) - # save offscreen buffers as video + # save offscreen buffers as video (only the frames actually rendered so far) if render =='offscreen': file_name = output_dir + filename + str(ep) + ".mp4" # check if the platform is OS -- make it compatible with quicktime if platform == "darwin": - skvideo.io.vwrite(file_name, np.asarray(frames),outputdict={"-pix_fmt": "yuv420p"}) + skvideo.io.vwrite(file_name, np.asarray(frames[:t]),outputdict={"-pix_fmt": "yuv420p"}) else: - skvideo.io.vwrite(file_name, np.asarray(frames)) + skvideo.io.vwrite(file_name, np.asarray(frames[:t])) prompt("saved", file_name, type=Prompt.INFO) self.mujoco_render_frames = False @@ -837,7 +884,16 @@ def examine_policy_new(self, self.mujoco_render_frames = False # start rollouts + # stop_requested becomes True either when the onscreen viewer is closed + # (Escape / window close, only possible with render=='onscreen') or when the + # user hits Ctrl+C (the only way to interrupt offscreen/none rollouts, which + # have no interactive window). Either way we stop early but still finalize + # and return whatever trace was completed, instead of crashing or looping + # invisibly to the end of `horizon`/`num_episodes`. + stop_requested = False for ep in range(num_episodes): + if stop_requested: + break # initialize ----------------------------- ep_t0 = timer.time() @@ -850,39 +906,47 @@ def examine_policy_new(self, # Rollout -------------------------------- obs, rwd, done, *_, env_info = self.forward(update_exteroception=True) # t=0 - while t < horizon and done is False: - - # print(t, t*self.dt, self.time, t*self.dt-self.time) - # Get step's actions ---------------------- - act = policy.get_action(obs)[0] if mode == 'exploration' else policy.get_action(obs)[1]['evaluation'] - - # render offscreen visuals ---------------------- - if render =='offscreen': - curr_frame = self.sim.renderer.render_offscreen( - width=frame_size[0], - height=frame_size[1], - camera_id=camera_name, - device_id=device_id) - - frames[t,:,:,:] = curr_frame - prompt(str(t), end=', ', flush=True, type=Prompt.INFO) - - # log values at time=t ---------------------------------- - datum_dict = dict( - time=self.time, - observations=obs, - actions=act.copy(), - rewards=rwd, - env_infos=env_info, - done=done, - ) - trace.append_datums(group_key=group_key, dataset_key_val=datum_dict) - - - # step env using actions from t=>t+1 ---------------------- - obs, rwd, done, *_, env_info = self.step(act, update_exteroception=True) - t = t+1 - ep_rwd += rwd + try: + while t < horizon and done is False and not self.viewer_exit_requested: + + # print(t, t*self.dt, self.time, t*self.dt-self.time) + # Get step's actions ---------------------- + act = policy.get_action(obs)[0] if mode == 'exploration' else policy.get_action(obs)[1]['evaluation'] + + # render offscreen visuals ---------------------- + if render =='offscreen': + curr_frame = self.sim.renderer.render_offscreen( + width=frame_size[0], + height=frame_size[1], + camera_id=camera_name, + device_id=device_id) + + frames[t,:,:,:] = curr_frame + prompt(str(t), end=', ', flush=True, type=Prompt.INFO) + + # log values at time=t ---------------------------------- + datum_dict = dict( + time=self.time, + observations=obs, + actions=act.copy(), + rewards=rwd, + env_infos=env_info, + done=done, + ) + trace.append_datums(group_key=group_key, dataset_key_val=datum_dict) + + + # step env using actions from t=>t+1 ---------------------- + obs, rwd, done, *_, env_info = self.step(act, update_exteroception=True) + t = t+1 + ep_rwd += rwd + except KeyboardInterrupt: + prompt("\nRollout interrupted by user (Ctrl+C)", type=Prompt.WARN) + stop_requested = True + + if self.viewer_exit_requested: + prompt("Rollout stopped: viewer was closed", type=Prompt.WARN) + stop_requested = True # record last step and finalize the rollout -------------------------------- act = np.nan*np.ones(self.action_space.shape) @@ -902,9 +966,9 @@ def examine_policy_new(self, file_name = output_dir + filename + str(ep) + ".mp4" # check if the platform is OS -- make it compatible with quicktime if platform == "darwin": - skvideo.io.vwrite(file_name, np.asarray(frames),outputdict={"-pix_fmt": "yuv420p"}) + skvideo.io.vwrite(file_name, np.asarray(frames[:t]),outputdict={"-pix_fmt": "yuv420p"}) else: - skvideo.io.vwrite(file_name, np.asarray(frames)) + skvideo.io.vwrite(file_name, np.asarray(frames[:t])) prompt("saved: "+file_name, type=Prompt.ALWAYS) self.mujoco_render_frames = False @@ -931,4 +995,4 @@ def get_reward_dict(self, obs_dict): Compute rewards dictionary Implement this in each subclass. """ - raise NotImplementedError + raise NotImplementedError \ No newline at end of file diff --git a/robohive/physics/mj_sim_scene.py b/robohive/physics/mj_sim_scene.py index 4ddb1b98..4b3d97c1 100644 --- a/robohive/physics/mj_sim_scene.py +++ b/robohive/physics/mj_sim_scene.py @@ -20,6 +20,7 @@ import_utils.dm_control_isavailable() import_utils.mujoco_isavailable() import dm_control.mujoco as dm_mujoco +from dm_control.rl.control import PhysicsError from robohive.physics.sim_scene import SimScene from robohive.renderer.mj_renderer import MJRenderer @@ -232,8 +233,8 @@ def advance(self, substeps: int = 1, render:bool = True): # Step the simulation substeps (frame_skip) times. try: self.sim.step(substeps) - except: - prompt("Simulation couldn't be stepped as intended. Issuing a reset", type=Prompt.WARN) + except PhysicsError as e: + prompt(f"Simulation couldn't be stepped as intended ({e}). Issuing a reset", type=Prompt.WARN) self.sim.reset() if render: diff --git a/robohive/renderer/mj_renderer.py b/robohive/renderer/mj_renderer.py index e979cdf7..cccfe7e1 100644 --- a/robohive/renderer/mj_renderer.py +++ b/robohive/renderer/mj_renderer.py @@ -31,7 +31,6 @@ def __init__(self, sim): self._window = None self._renderer = None self._paused = False - self._user_exit = False self._closing = False @@ -42,7 +41,7 @@ def key_callback(self, keycode): # Escape if keycode == 256: - self._user_exit = True + self._exit_requested = True def setup_renderer(self, model, height, width): @@ -58,7 +57,7 @@ def render_to_window(self): This function is a no-op if the window was already created. """ - if not self._window and not self._user_exit: + if not self._window and not self._exit_requested: self._window = viewer.launch_passive(self._sim.model.ptr, self._sim.data.ptr, key_callback=self.key_callback) self._update_camera_properties(self._window.cam) self._update_viewer_settings(self._window.opt) @@ -71,14 +70,20 @@ def refresh_window(self): """Refreshes the rendered window if one is present.""" if self._window is None or self._closing: return + + # Native window-close ('X' button) doesn't go through key_callback, so it never + # sets _exit_requested on its own; catch it here via is_running() too. + if not self._window.is_running(): + self._exit_requested = True + self._window.sync() # Keep checking to unpause if paused - while self._paused and not self._user_exit: + while self._paused and not self._exit_requested: # print("paused") time.sleep(.2) - if self._user_exit: + if self._exit_requested: self.close() diff --git a/robohive/renderer/renderer.py b/robohive/renderer/renderer.py index efcaccbc..8a570926 100644 --- a/robohive/renderer/renderer.py +++ b/robohive/renderer/renderer.py @@ -40,6 +40,15 @@ def __init__(self, sim): self._sim = sim self._camera_settings = {} self._viewer_settings = {} + self._exit_requested = False + + @property + def exit_requested(self) -> bool: + """True once the user has asked to stop rendering (e.g. Escape in an onscreen + viewer). Backends without an interactive window (offscreen/none) never set this; + callers driving a rollout loop should check it after every render_to_window()/ + refresh_window() call and break out if it's True.""" + return self._exit_requested @abc.abstractmethod def render_to_window(self): diff --git a/robohive/robot/robot.py b/robohive/robot/robot.py index 26cf5e66..89668241 100644 --- a/robohive/robot/robot.py +++ b/robohive/robot/robot.py @@ -87,17 +87,30 @@ def __init__(self, self.sim = mj_sim # Configure the robot - if self.robot_config is None: + if self.is_hardware and type(self).robot_config is not None: + # Session persists on the class (not the instance). + # A hardware connection is shared across separate Robot() objects, + # e.g. across an env.close()/make() cycle, instead of reconnecting. + prompt("Reusing a previours session of {}".format(self.name), 'white', 'on_grey') + else: prompt("Configuring a new session for {}".format(self.name), 'white', 'on_grey') - self.robot_config = self.configure_robot(self.sim, config_path) + robot_config = self.configure_robot(self.sim, config_path) if _ROBOT_VIZ: - self.configure_robot_viz(self.robot_config) + self.configure_robot_viz(robot_config) # start the robot if self.is_hardware is True: prompt("Initializing robot: %s"%(self.name), 'white', 'on_grey') - self.robot_config = self.hardware_init(self.robot_config) - else: - prompt("Reusing a previours session of {}".format(self.name), 'white', 'on_grey') + robot_config = self.hardware_init(robot_config) + # Only hardware sessions are persisted across instances + type(self).robot_config = robot_config + else: + # Sim sessions are NOT persisted across instances + self.robot_config = robot_config + + # Tracks whether close() has already been invoked on this instance (with either + # close_hardware value), to avoid a spurious __del__ warning when the hardware + # connection was intentionally left open via close(close_hardware=False). + self._explicitly_closed = False # check robot health if self.is_hardware is True: @@ -745,20 +758,30 @@ def reset(self, # Clear the robot class. Note that it doesn't close the persistent connection def __del__(self): - if self.robot_config is not None and self.is_hardware: + if self.robot_config is not None and self.is_hardware and not getattr(self, '_explicitly_closed', False): raise RuntimeWarning("RoboHive:> Robot class is being cleared from the workspace. This is expected if we still need to maintain the active connection to the hardware. A persistent connection to robot is still maintained and will be used next time a robot class is created. Ensure that a robot.close() is called to terminate the persistent connection before exiting the program.") - # Close the persistnent connection to the robot. This should be called only once at the end when persistent connection is no longer needed. - def close(self): - if self.robot_config is not None: - status = self.hardware_close() if self.is_hardware else True - if status: - prompt(f"Closed {self.name} (Status: {status})", 'white', 'on_grey', flush=True) - self.robot_config = None - else: - prompt(f"Error closing {self.name} (Status: {status})", 'red', 'on_grey', flush=True, type=Prompt.ERROR) - else: + # Close the connection to the robot. + # close_hardware=True (default): permanently closes the persistent hardware connection. + # close_hardware=False: leaves the hardware connection open/untouched so a subsequent + # Robot(is_hardware=True) (e.g. via env.make()) can reuse it instead of reconnecting. + def close(self, close_hardware=True): + self._explicitly_closed = True + if self.robot_config is None: prompt(f"Trying to close a non-existent robot", flush=True, type=Prompt.WARN) + return + if not close_hardware: + prompt(f"Leaving {self.name}'s hardware connection open (persistent)", 'white', 'on_grey', flush=True) + return + status = self.hardware_close() if self.is_hardware else True + if status: + prompt(f"Closed {self.name} (Status: {status})", 'white', 'on_grey', flush=True) + # Clear both: hardware sessions live on the class, sim sessions + # shadow it as an instance attr — null out whichever is set. + type(self).robot_config = None + self.__dict__.pop('robot_config', None) + else: + prompt(f"Error closing {self.name} (Status: {status})", 'red', 'on_grey', flush=True, type=Prompt.ERROR) def demo_robot(): diff --git a/robohive/tests/test_envs.py b/robohive/tests/test_envs.py index 21bec889..60c6945b 100644 --- a/robohive/tests/test_envs.py +++ b/robohive/tests/test_envs.py @@ -103,8 +103,8 @@ def check_env(self, environment_id, input_seed): # reset env2.reset() - del(env1) - del(env2) + env1.close() + env2.close() @implement_for("gym", None, "0.26") @@ -149,7 +149,7 @@ def check_old_envs(self, module_name, env_names, lite=False, seed=1234): # test step (everything together) observation, _reward, done, *_, _info = env.step(np.zeros(env.sim.model.nu)) - del(envw) + env.close() if __name__ == '__main__': diff --git a/robohive/utils/examine_env.py b/robohive/utils/examine_env.py index 513c0c11..649565a3 100644 --- a/robohive/utils/examine_env.py +++ b/robohive/utils/examine_env.py @@ -122,5 +122,8 @@ def main(env_name, policy_path, mode, seed, num_episodes, render, camera_name, o render_keys = ['env_infos/visual_dict/'+ key for key in env.visual_keys] paths.render(output_dir=output_dir, output_format="mp4", groups=["Trial0",], datasets=render_keys, input_fps=1/env.dt) + # cleanup + env.close() + if __name__ == '__main__': main()