text: import numpy as np
  
  from env.sawyer.sawyer import SawyerEnv
  
  
  class SawyerAssemblyObstacleEnv(SawyerEnv):
      def __init__(self, **kwargs):
          super().__init__("sawyer_assembly_obstacle.xml", **kwargs)
          self._get_reference()
  
      def _get_reference(self):
          super()._get_reference()
  
      @property
      def dof(self):
          return 7
  
      @property
      def init_qpos(self):
          return np.array([0.427, 0.13, 0.0557, 0.114, -0.0622, 0.0276, 0.00356])
  
      def _reset(self):
          init_qpos = (
              self.init_qpos + self.np_random.randn(self.init_qpos.shape[0]) * 0.02
          )
          self.sim.data.qpos[self.ref_joint_pos_indexes] = init_qpos
          self.sim.data.qvel[self.ref_joint_vel_indexes] = 0.0
          self.sim.forward()
  
          return self._get_obs()
  
      def compute_reward(self, action):
          info = {}
          reward = 0
          reward_type = self._kwargs["reward_type"]
          pegHeadPos = self.sim.data.get_site_xpos("pegHead")
          hole = self.sim.data.get_site_xpos("hole")
          dist = np.linalg.norm(pegHeadPos - hole)
          hole_bottom = self.sim.data.get_site_xpos("hole_bottom")
          dist_to_hole_bottom = np.linalg.norm(pegHeadPos - hole_bottom)
          dist_to_hole = np.linalg.norm(pegHeadPos - hole)
          reward_reach = 0
          if dist < 0.3:
              reward_reach += 0.4 * (1 - np.tanh(15 * dist_to_hole))
          reward += reward_reach
          if dist_to_hole_bottom < 0.025:
              reward += self._kwargs["success_reward"]
              self._success = True
              self._terminal = True
  
          return reward, info
  
      def _get_obs(self):
          di = super()._get_obs()
          di["hole"] = self.sim.data.get_site_xpos("hole")
          di["pegHead"] = self.sim.data.get_site_xpos("pegHead")
          di["pegEnd"] = self.sim.data.get_site_xpos("pegEnd")
          di["peg_quat"] = self._get_quat("peg")
          return di
  
      @property
      def static_bodies(self):
          return ["table"]
  
      @property
      def manipulation_bodies(self):
          return ["furniture", "0_part0", "1_part1", "4_part4", "2_part2"]
  
      @property
      def manipulation_geom_ids(self):
          body_ids = []
          for body_name in self.manipulation_bodies:
              body_ids.append(self.sim.model.body_name2id(body_name))
  
          geom_ids = []
          for geom_id, body_id in enumerate(self.sim.model.geom_bodyid):
              if body_id in body_ids:
                  geom_ids.append(geom_id)
          return geom_ids
  
      @property
      def static_geoms(self):
          return []
  
      @property
      def static_geom_ids(self):
          body_ids = []
          for body_name in self.static_bodies:
              body_ids.append(self.sim.model.body_name2id(body_name))
  
          geom_ids = []
          for geom_id, body_id in enumerate(self.sim.model.geom_bodyid):
              if body_id in body_ids:
                  geom_ids.append(geom_id)
          return geom_ids
  
      def _step(self, action, is_planner=False):
          """
          (Optional) does gripper visualization after actions.
          """
          assert len(action) == self.dof, "environment got invalid action dimension"
  
          if not is_planner or self._prev_state is None:
              self._prev_state = self.sim.data.qpos[self.ref_joint_pos_indexes].copy()
  
          if self._i_term is None:
              self._i_term = np.zeros_like(self.mujoco_robot.dof)
  
          if is_planner:
              rescaled_ac = np.clip(
                  action[: self.robot_dof], -self._ac_scale, self._ac_scale
              )
          else:
              rescaled_ac = np.clip(
                  action[: self.robot_dof] * self._ac_scale,
                  -self._ac_scale,
                  self._ac_scale,
              )
          desired_state = self._prev_state + rescaled_ac
  
          n_inner_loop = int(self._frame_dt / self.dt)
          for _ in range(n_inner_loop):
              self.sim.data.qfrc_applied[
                  self.ref_joint_vel_indexes
              ] = self.sim.data.qfrc_bias[self.ref_joint_vel_indexes].copy()
  
              if self.use_robot_indicator:
                  self.sim.data.qfrc_applied[
                      self.ref_indicator_joint_pos_indexes
                  ] = self.sim.data.qfrc_bias[self.ref_indicator_joint_pos_indexes].copy()
  
              if self.use_target_robot_indicator:
                  self.sim.data.qfrc_applied[
                      self.ref_target_indicator_joint_pos_indexes
                  ] = self.sim.data.qfrc_bias[
                      self.ref_target_indicator_joint_pos_indexes
                  ].copy()
              self._do_simulation(desired_state)
  
          self._prev_state = np.copy(desired_state)
          reward, info = self.compute_reward(action)
  
          return self._get_obs(), reward, self._terminal, info
  