Python mujoco_py.MjRenderContextOffscreen() Examples

def _reset_internal(self):
        """Resets simulation internal configurations."""
        # instantiate simulation from MJCF model
        self._load_model()
        self.mjpy_model = self.model.get_model(mode="mujoco_py")
        self.sim = MjSim(self.mjpy_model)
        self.initialize_time(self.control_freq)

        # create visualization screen or renderer
        if self.has_renderer and self.viewer is None:
            self.viewer = MujocoPyRenderer(self.sim)
            self.viewer.viewer.vopt.geomgroup[0] = (
                1 if self.render_collision_mesh else 0
            )
            self.viewer.viewer.vopt.geomgroup[1] = 1 if self.render_visual_mesh else 0

            # hiding the overlay speeds up rendering significantly
            self.viewer.viewer._hide_overlay = True

        elif self.has_offscreen_renderer:
            if self.sim._render_context_offscreen is None:
                render_context = MjRenderContextOffscreen(self.sim)
                self.sim.add_render_context(render_context)
            self.sim._render_context_offscreen.vopt.geomgroup[0] = (
                1 if self.render_collision_mesh else 0
            )
            self.sim._render_context_offscreen.vopt.geomgroup[1] = (
                1 if self.render_visual_mesh else 0
            )

        # additional housekeeping
        self.sim_state_initial = self.sim.get_state()
        self._get_reference()
        self.cur_time = 0
        self.timestep = 0
        self.done = False 

def reset_from_xml_string(self, xml_string):
        """Reloads the environment from an XML description of the environment."""

        # if there is an active viewer window, destroy it
        self.close()

        # load model from xml
        self.mjpy_model = load_model_from_xml(xml_string)

        self.sim = MjSim(self.mjpy_model)
        self.initialize_time(self.control_freq)
        if self.has_renderer and self.viewer is None:
            self.viewer = MujocoPyRenderer(self.sim)
            self.viewer.viewer.vopt.geomgroup[0] = (
                1 if self.render_collision_mesh else 0
            )
            self.viewer.viewer.vopt.geomgroup[1] = 1 if self.render_visual_mesh else 0

            # hiding the overlay speeds up rendering significantly
            self.viewer.viewer._hide_overlay = True

        elif self.has_offscreen_renderer:
            render_context = MjRenderContextOffscreen(self.sim)
            render_context.vopt.geomgroup[0] = 1 if self.render_collision_mesh else 0
            render_context.vopt.geomgroup[1] = 1 if self.render_visual_mesh else 0
            self.sim.add_render_context(render_context)

        self.sim_state_initial = self.sim.get_state()
        self._get_reference()
        self.cur_time = 0
        self.timestep = 0
        self.done = False

        # necessary to refresh MjData
        self.sim.forward() 

def __init__(self,
                 wrapped_env,
                 imsize=32,
                 keep_prev=0,
                 init_camera=None,
                 camera_name=None,
                 transpose=False,
                 grayscale=False,
                 normalize=False,
    ):
        self.quick_init(locals())
        super().__init__(wrapped_env)

        self.imsize = imsize
        if grayscale:
            self.image_length = self.imsize * self.imsize
        else:
            self.image_length = 3 * self.imsize * self.imsize
        # This is torch format rather than PIL image
        self.image_shape = (self.imsize, self.imsize)
        # Flattened past image queue
        self.history_length = keep_prev + 1
        self.history = deque(maxlen=self.history_length)
        # init camera
        if init_camera is not None:
            sim = self._wrapped_env.sim
            viewer = mujoco_py.MjRenderContextOffscreen(sim, device_id=-1)
            init_camera(viewer.cam)
            sim.add_render_context(viewer)
        self.camera_name = camera_name # None means default camera
        self.transpose = transpose
        self.grayscale = grayscale
        self.normalize = normalize
        self._render_local = False

        self.observation_space = Box(low=0.0,
                                     high=1.0,
                                     shape=(self.image_length * self.history_length,)) 

def _setup_render_rgb(sim: mujoco_py.MjSim) -> mujoco_py.MjSim:
  # create copy of simulation to customize rendering context
  # flags defined in mjvisualize.h
  render_sim = mujoco_py.MjSim(sim.model)
  render_sim.set_state(sim.get_state())
  render_ctx = mujoco_py.MjRenderContextOffscreen(render_sim)
  render_ctx.scn.stereo = 2 # side-by-side rendering
  return render_sim 

def _setup_render_seg(sim: mujoco_py.MjSim) -> mujoco_py.MjSim:
  # create copy of simulation to customize rendering context
  # flags defined in mjvisualize.h
  render_sim = mujoco_py.MjSim(sim.model)
  render_sim.set_state(sim.get_state())
  render_ctx = mujoco_py.MjRenderContextOffscreen(render_sim)
  render_ctx.vopt.flags[1] = 0 # textures off
  render_ctx.vopt.flags[17] = 0 # static body off
  render_ctx.scn.flags[0] = 0 # shadow off
  render_ctx.scn.flags[2] = 0 # reflection off
  render_ctx.scn.flags[4] = 0 # skybox off
  render_ctx.scn.stereo = 2 # side-by-side rendering
  return render_sim 

def _setup_render_depth(sim: mujoco_py.MjSim) -> mujoco_py.MjSim:
  # create copy of simulation to customize rendering context
  # flags defined in mjvisualize.h
  render_sim = mujoco_py.MjSim(sim.model)
  render_sim.set_state(sim.get_state())
  render_ctx = mujoco_py.MjRenderContextOffscreen(render_sim)
  render_ctx.vopt.flags[1] = 0 # textures off
  render_ctx.scn.flags[0] = 0 # shadow off
  render_ctx.scn.flags[2] = 0 # reflection off
  render_ctx.scn.flags[4] = 0 # skybox off
  render_ctx.scn.stereo = 2 # side-by-side rendering
  return render_sim 

def _get_viewer(self, mode):
        self.viewer = self._viewers.get(mode)
        if self.viewer is None:
            if mode == 'human':
                self.viewer = mujoco_py.MjViewer(self.sim)
            elif mode == 'rgb_array' or mode == 'depth_array':
                self.viewer = mujoco_py.MjRenderContextOffscreen(self.sim, -1)
                
            self.viewer_setup()
            self._viewers[mode] = self.viewer
        return self.viewer 

def _get_viewer(self, mode):
        self.viewer = self._viewers.get(mode)
        if self.viewer is None:
            if mode == 'human':
                self.viewer = mujoco_py.MjViewer(self.sim)
            elif mode == 'rgb_array' or mode == 'depth_array':
                self.viewer = mujoco_py.MjRenderContextOffscreen(self.sim, -1)
                
            self.viewer_setup()
            self._viewers[mode] = self.viewer
        return self.viewer 

def get_current_image_obs(self):
        # print(self._get_viewer)
        # print(self._get_viewer())
        # self.viewer = mujoco_py.MjRenderContextOffscreen(self.sim, -1)
        # self.viewer = mujoco_py.MjViewer()
        # self.viewer.start()
        # self.viewer.set_model(self.model)
        # self.viewer_setup()
        # self._get_viewer()

        
        # image = self.viewer.get_image()
        # pil_image = Image.frombytes('RGB', (image[1], image[2]), image[0])
        # pil_image = pil_image.resize((125,125), Image.ANTIALIAS)
        # image = np.flipud(np.array(pil_image))

        # image = self.render(mode='rgb_array', width=125, height=125)
        # transpose to make it have correct ordering of dimensions for pytorch
        # image = image.transpose((2,0,1))
        # image = np.array(image).astype(np.float32)
        # image /= 255.0

        image = self.render(mode='rgb_array', width=500, height=500)
        image = Image.fromarray(image)
        image = image.resize((125,125), PIL.Image.LANCZOS)
        image = np.array(image)
        image = image.transpose((2,0,1))
        image = np.array(image).astype(np.float32)
        image /= 255.0

        X = np.concatenate([
            self.sim.data.qpos.flat[:7],
            self.sim.data.qvel.flat[:7],
            self.get_body_com('tips_arm'),
            self.get_body_com('goal'),
        ]).copy()
        X = (X - self.state_mean) / self.state_std

        return image, X 

def _get_viewer(self, mode):
        self.viewer = self._viewers.get(mode)
        if self.viewer is None:
            if mode == 'human':
                self.viewer = mujoco_py.MjViewer(self.sim)
            elif mode == 'rgb_array':
                self.viewer = mujoco_py.MjRenderContextOffscreen(self.sim, device_id=-1)
            self._viewer_setup()
            self._viewers[mode] = self.viewer
        return self.viewer

    # Extension methods
    # ---------------------------- 

def _get_viewer(self, mode):
        self.viewer = self._viewers.get(mode)
        if self.viewer is None:
            if mode == 'human':
                self.viewer = mujoco_py.MjViewer(self.sim)
            elif mode == 'rgb_array' or mode == 'depth_array':
                self.viewer = mujoco_py.MjRenderContextOffscreen(self.sim, -1)

            self.viewer_setup()
            self._viewers[mode] = self.viewer
        return self.viewer 

def initialize_camera(self):
        # set camera parameters for viewing
        sim = self.sim
        viewer = mujoco_py.MjRenderContextOffscreen(sim)
        camera = viewer.cam
        camera.type = 1
        camera.trackbodyid = 0
        camera.elevation = -20
        sim.add_render_context(viewer) 

def render(self, mode='human', close=False):
        if mode == 'human':
            if self.viewer is None:
                self.viewer = MjViewer(self.sim)
            self.viewer.render()
        elif mode == 'rgb_array':
            if self.rgb_viewer is None:
                self.rgb_viewer = MjRenderContextOffscreen(self.sim, 0)
            self.rgb_viewer.render(500, 500)
            # window size used for old mujoco-py:
            data = self.rgb_viewer.read_pixels(500, 500, depth=False)
            # original image is upside-down, so flip it
            return data[::-1, :, :] 

def __init__(self,
                 wrapped_env,
                 imsize=32,
                 keep_prev=0,
                 init_camera=None,
                 camera_name=None,
                 transpose=False,
                 grayscale=False,
                 normalize=False,
                 ):
        import mujoco_py
        super().__init__(wrapped_env)

        self.imsize = imsize
        if grayscale:
            self.image_length = self.imsize * self.imsize
        else:
            self.image_length = 3 * self.imsize * self.imsize
        # This is torch format rather than PIL image
        self.image_shape = (self.imsize, self.imsize)
        # Flattened past image queue
        self.history_length = keep_prev + 1
        self.history = deque(maxlen=self.history_length)
        # init camera
        if init_camera is not None:
            sim = self._wrapped_env.sim
            viewer = mujoco_py.MjRenderContextOffscreen(sim, device_id=-1)
            init_camera(viewer.cam)
            sim.add_render_context(viewer)
        self.camera_name = camera_name  # None means default camera
        self.transpose = transpose
        self.grayscale = grayscale
        self.normalize = normalize
        self._render_local = False

        self.observation_space = Box(low=0.0,
                                     high=1.0,
                                     shape=(
                                     self.image_length * self.history_length,))