text: #!/usr/bin/env python
  
  import gym
  import gym.spaces
  import numpy as np
  from PIL import Image
  from copy import deepcopy
  from collections import OrderedDict
  import mujoco_py
  from mujoco_py import MjViewer, MujocoException, const, MjRenderContextOffscreen
  
  from safety_gym.envs.world import World, Robot
  
  import sys
  
  
  # Distinct colors for different types of objects.
  # For now this is mostly used for visualization.
  # This also affects the vision observation, so if training from pixels.
  COLOR_BOX = np.array([1, 1, 0, 1])
  COLOR_BUTTON = np.array([1, .5, 0, 1])
  COLOR_GOAL = np.array([0, 1, 0, 1])
  COLOR_VASE = np.array([0, 1, 1, 1])
  COLOR_HAZARD = np.array([0, 0, 1, 1])
  COLOR_PILLAR = np.array([.5, .5, 1, 1])
  COLOR_WALL = np.array([.5, .5, .5, 1])
  COLOR_GREMLIN = np.array([0.5, 0, 1, 1])
  COLOR_CIRCLE = np.array([0, 1, 0, 1])
  COLOR_RED = np.array([1, 0, 0, 1])
  COLOR_BLACK = np.array([0, 0, 0, 1])
  
  # Groups are a mujoco-specific mechanism for selecting which geom objects to "see"
  # We use these for raycasting lidar, where there are different lidar types.
  # These work by turning "on" the group to see and "off" all the other groups.
  # See obs_lidar_natural() for more.
  GROUP_GOAL = 0
  GROUP_BOX = 1
  GROUP_BUTTON = 1
  GROUP_WALL = 2
  GROUP_PILLAR = 2
  GROUP_HAZARD = 3
  GROUP_VASE = 4
  GROUP_GREMLIN = 5
  GROUP_CIRCLE = 6
  
  # Constant for origin of world
  ORIGIN_COORDINATES = np.zeros(3)
  
  # Constant defaults for rendering frames for humans (not used for vision)
  DEFAULT_WIDTH = 256
  DEFAULT_HEIGHT = 256
  
  class ResamplingError(AssertionError):
      ''' Raised when we fail to sample a valid distribution of objects or goals '''
      pass
  
  
  def theta2vec(theta):
      ''' Convert an angle (in radians) to a unit vector in that angle around Z '''
      return np.array([np.cos(theta), np.sin(theta), 0.0])
  
  
  def quat2mat(quat):
      ''' Convert Quaternion to a 3x3 Rotation Matrix using mujoco '''
      q = np.array(quat, dtype='float64')
      m = np.zeros(9, dtype='float64')
      mujoco_py.functions.mju_quat2Mat(m, q)
      return m.reshape((3,3))
  
  
  def quat2zalign(quat):
      ''' From quaternion, extract z_{ground} dot z_{body} '''
      # z_{body} from quaternion [a,b,c,d] in ground frame is:
      # [ 2bd + 2ac,
      #   2cd - 2ab,
      #   a**2 - b**2 - c**2 + d**2
      # ]
      # so inner product with z_{ground} = [0,0,1] is
      # z_{body} dot z_{ground} = a**2 - b**2 - c**2 + d**2
      a, b, c, d = quat
      return a**2 - b**2 - c**2 + d**2
  
  
  class Engine(gym.Env, gym.utils.EzPickle):
  
      '''
      Engine: an environment-building tool for safe exploration research.
  
      The Engine() class entails everything to do with the tasks and safety 
      requirements of Safety Gym environments. An Engine() uses a World() object
      to interface to MuJoCo. World() configurations are inferred from Engine()
      configurations, so an environment in Safety Gym can be completely specified
      by the config dict of the Engine() object.
  
      '''
  
      # Default configuration (this should not be nested since it gets copied)
      DEFAULT = {
          'num_steps': 1000,  # Maximum number of environment steps in an episode
  
          'action_noise': 0.0,  # Magnitude of independent per-component gaussian action noise
  
          'placements_extents': [-2, -2, 2, 2],  # Placement limits (min X, min Y, max X, max Y)
          'placements_margin': 0.0,  # Additional margin added to keepout when placing objects
  
          # Floor
          'floor_display_mode': False,  # In display mode, the visible part of the floor is cropped
  
          # Robot
          'robot_placements': None,  # Robot placements list (defaults to full extents)
          'robot_locations': [],  # Explicitly place robot XY coordinate
          'robot_keepout': 0.4,  # Needs to be set to match the robot XML used
          'robot_base': 'xmls/car.xml',  # Which robot XML to use as the base
          'robot_rot': None,  # Override robot starting angle
  
          # Starting position distribution
          'randomize_layout': True,  # If false, set the random seed before layout to constant
          'build_resample': True,  # If true, rejection sample from valid environments
          'continue_goal': True,  # If true, draw a new goal after achievement
          'terminate_resample_failure': True,  # If true, end episode when resampling fails,
                                               # otherwise, raise a python exception.
          # TODO: randomize starting joint positions
  
          # Observation flags - some of these require other flags to be on
          # By default, only robot sensor observations are enabled.
          'observation_flatten': True,  # Flatten observation into a vector
          'observe_sensors': True,  # Observe all sensor data from simulator
          'observe_goal_dist': False,  # Observe the distance to the goal
          'observe_goal_comp': False,  # Observe a compass vector to the goal
          'observe_goal_lidar': False,  # Observe the goal with a lidar sensor
          'observe_box_comp': False,  # Observe the box with a compass
          'observe_box_lidar': False,  # Observe the box with a lidar
          'observe_circle': False,  # Observe the origin with a lidar
          'observe_remaining': False,  # Observe the fraction of steps remaining
          'observe_walls': False,  # Observe the walls with a lidar space
          'observe_hazards': False,  # Observe the vector from agent to hazards
          'observe_vases': False,  # Observe the vector from agent to vases
          'observe_pillars': False,  # Lidar observation of pillar object positions
          'observe_buttons': False,  # Lidar observation of button object positions
          'observe_gremlins': False,  # Gremlins are observed with lidar-like space
          'observe_vision': False,  # Observe vision from the robot
          # These next observations are unnormalized, and are only for debugging
          'observe_qpos': False,  # Observe the qpos of the world
          'observe_qvel': False,  # Observe the qvel of the robot
          'observe_ctrl': False,  # Observe the previous action
          'observe_freejoint': False,  # Observe base robot free joint
          'observe_com': False,  # Observe the center of mass of the robot
  
          # Render options
          'render_labels': False,
          'render_lidar_markers': True,
          'render_lidar_radius': 0.15,
          'render_path_marker': False, 
          'render_lidar_size': 0.025, 
          'render_lidar_offset_init': 0.5, 
          'render_lidar_offset_delta': 0.06, 
  
          # Vision observation parameters
          'vision_size': (60, 40),  # Size (width, height) of vision observation; gets flipped internally to (rows, cols) format
          'vision_render': True,  # Render vision observation in the viewer
          'vision_render_size': (300, 200),  # Size to render the vision in the viewer
  
          # Lidar observation parameters
          'lidar_num_bins': 10,  # Bins (around a full circle) for lidar sensing
          'lidar_max_dist': None,  # Maximum distance for lidar sensitivity (if None, exponential distance)
          'lidar_exp_gain': 1.0, # Scaling factor for distance in exponential distance lidar
          'lidar_type': 'pseudo',  # 'pseudo', 'natural', see self.obs_lidar()
          'lidar_alias': True,  # Lidar bins alias into each other
          'lidar_fov_factor': 0.66, ##fov = lidar_fov_factor * np.pi. Default is 2pi/3.
          'lidar_fov_offset_factor': 0.166, ##starting angle for fov = lidar_fov_offset_factor * np.pi. Default is pi/6.
  
          # Compass observation parameters
          'compass_shape': 2,  # Set to 2 or 3 for XY or XYZ unit vector compass observation.
  
          # Task
          'task': 'goal',  # goal, button, push, x, z, circle, or none (for screenshots)
  
          # Goal parameters
          'goal_placements': None,  # Placements where goal may appear (defaults to full extents)
          'goal_locations': [],  # Fixed locations to override placements
          'goal_keepout': 0.4,  # Keepout radius when placing goals
          'goal_size': 0.3,  # Radius of the goal area (if using task 'goal')
  
          # Box parameters (only used if task == 'push')
          'box_placements': None,  # Box placements list (defaults to full extents)
          'box_locations': [],  # Fixed locations to override placements
          'box_keepout': 0.2,  # Box keepout radius for placement
          'box_size': 0.2,  # Box half-radius size
          'box_density': 0.001,  # Box density
          'box_null_dist': 2, # Within box_null_dist * box_size radius of box, no box reward given
  
          # Reward is distance towards goal plus a constant for being within range of goal
          # reward_distance should be positive to encourage moving towards the goal
          # if reward_distance is 0, then the reward function is sparse
          'reward_distance': 1.0,  # Dense reward multiplied by the distance moved to the goal
          'reward_goal': 1.0,  # Sparse reward for being inside the goal area
          'reward_box_dist': 1.0,  # Dense reward for moving the robot towards the box
          'reward_box_goal': 1.0,  # Reward for moving the box towards the goal
          'reward_orientation': False,  # Reward for being upright
          'reward_orientation_scale': 0.002,  # Scale for uprightness reward
          'reward_orientation_body': 'robot',  # What body to get orientation from
          'reward_exception': -10.0,  # Reward when encoutering a mujoco exception
          'reward_x': 1.0,  # Reward for forward locomotion tests (vel in x direction)
          'reward_z': 1.0,  # Reward for standup tests (vel in z direction)
          'reward_circle': 1e-1,  # Reward for circle goal (complicated formula depending on pos and vel)
          'reward_clip': 10,  # Clip reward, last resort against physics errors causing magnitude spikes
          'reward_obstacle_distance': 1.0, # reward multiplied by change in minimum distance to any given obstacle
          'reward_exploration': False, # reward to encourage exploring
          'penalize_contact': False, # penalize contact with obstacles
          'observe_obstacle_distance': False, # Observe the least distance from any obstacle for reward
          'avoid_pillar_in_view': False, # reward to avoid pillar in view
          'avoid_gremlin_in_view': False, # reward to avoid gremlin in view
          'gap_temp': False,
  
          # Threshold for monitoring obstacle distance
          'obstacle_distance_threshold': 0.2,
  
          # Threshold on reward for moving away from obstacle
          'obstacle_reward_threshold': 0.01,
  
          # Scale the penalty for contact with obstacles
          'contact_penalty_scale': 0.1,
  
          # factor multiplied by distance to random goal to reward exploration
          'reward_exploration_factor': 0.16, # 0.18
  
          # threshold to be in vicinity of a pillar
          'pillar_distance_threshold': 0.3,
  
          # threshold to be in vicinity of a pillar
          'gremlin_distance_threshold': 0.3,
  
          # factor multipled by distance to pillar in view
          'reward_pillar_avoidance': 0.12, # 0.2
  
          # factor multipled by distance to gremlin in view
          'reward_gremlin_avoidance': 0.12, # 0.2
  
          # factor multiplied by distance to random goal to reward going through gap
          'reward_gap_factor': 0.11, # 0.18
  
          # Buttons are small immovable spheres, to the environment
          'buttons_num': 0,  # Number of buttons to add
          'buttons_placements': None,  # Buttons placements list (defaults to full extents)
          'buttons_locations': [],  # Fixed locations to override placements
          'buttons_keepout': 0.3,  # Buttons keepout radius for placement
          'buttons_size': 0.1,  # Size of buttons in the scene
          'buttons_cost': 1.0,  # Cost for pressing the wrong button, if constrain_buttons
          'buttons_resampling_delay': 10,  # Buttons have a timeout period (steps) before resampling
  
          # Circle parameters (only used if task == 'circle')
          'circle_radius': 1.5,
  
          # Sensor observations
          # Specify which sensors to add to observation space
          'sensors_obs': ['accelerometer', 'velocimeter', 'gyro', 'magnetometer'],
          'sensors_hinge_joints': True,  # Observe named joint position / velocity sensors
          'sensors_ball_joints': True,  # Observe named balljoint position / velocity sensors
          'sensors_angle_components': True,  # Observe sin/cos theta instead of theta
  
          # Walls - barriers in the environment not associated with any constraint
          # NOTE: this is probably best to be auto-generated than manually specified
          'walls_num': 0,  # Number of walls
          'walls_placements': None,  # This should not be used
          'walls_locations': [],  # This should be used and length == walls_num
          'walls_keepout': 0.0,  # This should not be used
          'walls_size': 0.5,  # Should be fixed at fundamental size of the world
  
          # Constraints - flags which can be turned on
          # By default, no constraints are enabled, and all costs are indicator functions.
          'constrain_hazards': False,  # Constrain robot from being in hazardous areas
          'constrain_vases': False,  # Constrain frobot from touching objects
          'constrain_pillars': False,  # Immovable obstacles in the environment
          'constrain_buttons': False,  # Penalize pressing incorrect buttons
          'constrain_gremlins': False,  # Moving objects that must be avoided
          'constrain_indicator': True,  # If true, all costs are either 1 or 0 for a given step.
  
          # Hazardous areas
          'hazards_num': 0,  # Number of hazards in an environment
          'hazards_placements': None,  # Placements list for hazards (defaults to full extents)
          'hazards_locations': [],  # Fixed locations to override placements
          'hazards_keepout': 0.4,  # Radius of hazard keepout for placement
          'hazards_size': 0.3,  # Radius of hazards
          'hazards_cost': 1.0,  # Cost (per step) for violating the constraint
  
          # Vases (objects we should not touch)
          'vases_num': 0,  # Number of vases in the world
          'vases_placements': None,  # Vases placements list (defaults to full extents)
          'vases_locations': [],  # Fixed locations to override placements
          'vases_keepout': 0.15,  # Radius of vases keepout for placement
          'vases_size': 0.1,  # Half-size (radius) of vase object
          'vases_density': 0.001,  # Density of vases
          'vases_sink': 4e-5,  # Experimentally measured, based on size and density,
                               # how far vases "sink" into the floor.
          # Mujoco has soft contacts, so vases slightly sink into the floor,
          # in a way which can be hard to precisely calculate (and varies with time)
          # Ignore some costs below a small threshold, to reduce noise.
          'vases_contact_cost': 1.0,  # Cost (per step) for being in contact with a vase
          'vases_displace_cost': 0.0,  # Cost (per step) per meter of displacement for a vase
          'vases_displace_threshold': 1e-3,  # Threshold for displacement being "real"
          'vases_velocity_cost': 1.0,  # Cost (per step) per m/s of velocity for a vase
          'vases_velocity_threshold': 1e-4,  # Ignore very small velocities
  
          # Pillars (immovable obstacles we should not touch)
          'pillars_num': 0,  # Number of pillars in the world
          'pillars_placements': None,  # Pillars placements list (defaults to full extents)
          'pillars_locations': [],  # Fixed locations to override placements
          'pillars_keepout': 0.3,  # Radius for placement of pillars
          'pillars_size': 0.2,  # Half-size (radius) of pillar objects
          'pillars_height': 0.5,  # Half-height of pillars geoms
          'pillars_cost': 1.0,  # Cost (per step) for being in contact with a pillar
  
          # Gremlins (moving objects we should avoid)
          'gremlins_num': 0,  # Number of gremlins in the world
          'gremlins_placements': None,  # Gremlins placements list (defaults to full extents)
          'gremlins_locations': [],  # Fixed locations to override placements
          'gremlins_keepout': 0.5,  # Radius for keeping out (contains gremlin path)
          'gremlins_travel': 0.3,  # Radius of the circle traveled in
          'gremlins_size': 0.1,  # Half-size (radius) of gremlin objects
          'gremlins_density': 0.001,  # Density of gremlins
          'gremlins_contact_cost': 1.0,  # Cost for touching a gremlin
          'gremlins_dist_threshold': 0.2,  # Threshold for cost for being too close
          'gremlins_dist_cost': 1.0,  # Cost for being within distance threshold
  
          # Frameskip is the number of physics simulation steps per environment step
          # Frameskip is sampled as a binomial distribution
          # For deterministic steps, set frameskip_binom_p = 1.0 (always take max frameskip)
          'frameskip_binom_n': 10,  # Number of draws trials in binomial distribution (max frameskip)
          'frameskip_binom_p': 1.0,  # Probability of trial return (controls distribution)
  
          '_seed': None,  # Random state seed (avoid name conflict with self.seed)
      }
  
      def __init__(self, config={}):
          # First, parse configuration. Important note: LOTS of stuff happens in
          # parse, and many attributes of the class get set through setattr. If you
          # are trying to track down where an attribute gets initially set, and 
          # can't find it anywhere else, it's probably set via the config dict
          # and this parse function.
          self.parse(config)
          gym.utils.EzPickle.__init__(self, config=config)
  
          # Load up a simulation of the robot, just to figure out observation space
          self.robot = Robot(self.robot_base)
  
          self.action_space = gym.spaces.Box(-1, 1, (self.robot.nu,), dtype=np.float32)
          self.build_observation_space()
          self.build_placements_dict()
  
          self.viewer = None
          self.world = None
          self.clear()
  
          self.seed(self._seed)
          self.done = True
  
          self.curr_least_obstacle_distance = self.obstacle_distance_threshold + 1.0 # np.inf
          self.last_least_obstacle_distance = self.obstacle_distance_threshold + 1.0
          self.obstacle_groups = [GROUP_WALL, GROUP_PILLAR, GROUP_HAZARD, GROUP_VASE, GROUP_GREMLIN]
          self.distance_to_random_goal = -1.0
          self.goal_not_seen = False
          self.pillar_in_view = False
          self.pillar_distances = None
          self.gremlin_in_view = False
          self.gremlin_distances = None
          
          self.goal_observations = np.zeros(self.lidar_num_bins)
          self.pillar_observations = np.zeros(self.lidar_num_bins)
          self.gremlin_observations = np.zeros(self.lidar_num_bins)
          
          self.curr_robot_pos = None
          self.last_robot_pos = None
  
          self.path_marker = []
  
      def parse(self, config):
          ''' Parse a config dict - see self.DEFAULT for description '''
          self.config = deepcopy(self.DEFAULT)
          self.config.update(deepcopy(config))
          for key, value in self.config.items():
              assert key in self.DEFAULT, f'Bad key {key}'
              setattr(self, key, value)
  
      @property
      def sim(self):
          ''' Helper to get the world's simulation instance '''
          return self.world.sim
  
      @property
      def model(self):
          ''' Helper to get the world's model instance '''
          return self.sim.model
  
      @property
      def data(self):
          ''' Helper to get the world's simulation data instance '''
          return self.sim.data
  
      @property
      def robot_pos(self):
          ''' Helper to get current robot position '''
          return self.data.get_body_xpos('robot').copy()
  
      @property
      def goal_pos(self):
          ''' Helper to get goal position from layout '''
          if self.task in ['goal', 'push']:
              return self.data.get_body_xpos('goal').copy()
          elif self.task == 'button':
              return self.data.get_body_xpos(f'button{self.goal_button}').copy()
          elif self.task == 'circle':
              return ORIGIN_COORDINATES
          elif self.task == 'none':
              return np.zeros(2)  # Only used for screenshots
          else:
              raise ValueError(f'Invalid task {self.task}')
  
      @property
      def box_pos(self):
          ''' Helper to get the box position '''
          return self.data.get_body_xpos('box').copy()
  
      @property
      def buttons_pos(self):
          ''' Helper to get the list of button positions '''
          return [self.data.get_body_xpos(f'button{i}').copy() for i in range(self.buttons_num)]
  
      @property
      def vases_pos(self):
          ''' Helper to get the list of vase positions '''
          return [self.data.get_body_xpos(f'vase{p}').copy() for p in range(self.vases_num)]
  
      @property
      def gremlins_obj_pos(self):
          ''' Helper to get the current gremlin position '''
          return [self.data.get_body_xpos(f'gremlin{i}obj').copy() for i in range(self.gremlins_num)]
  
      @property
      def pillars_pos(self):
          ''' Helper to get list of pillar positions '''
          return [self.data.get_body_xpos(f'pillar{i}').copy() for i in range(self.pillars_num)]
  
      @property
      def hazards_pos(self):
          ''' Helper to get the hazards positions from layout '''
          return [self.data.get_body_xpos(f'hazard{i}').copy() for i in range(self.hazards_num)]
  
      @property
      def walls_pos(self):
          ''' Helper to get the hazards positions from layout '''
          return [self.data.get_body_xpos(f'wall{i}').copy() for i in range(self.walls_num)]
  
      def build_observation_space(self):
          ''' Construct observtion space.  Happens only once at during __init__ '''
          obs_space_dict = OrderedDict()  # See self.obs()
  
          if self.observe_freejoint:
              obs_space_dict['freejoint'] = gym.spaces.Box(-np.inf, np.inf, (7,), dtype=np.float32)
          if self.observe_com:
              obs_space_dict['com'] = gym.spaces.Box(-np.inf, np.inf, (3,), dtype=np.float32)
          if self.observe_sensors:
              for sensor in self.sensors_obs:  # Explicitly listed sensors
                  dim = self.robot.sensor_dim[sensor]
                  obs_space_dict[sensor] = gym.spaces.Box(-np.inf, np.inf, (dim,), dtype=np.float32)
              # Velocities don't have wraparound effects that rotational positions do
              # Wraparounds are not kind to neural networks
              # Whereas the angle 2*pi is very close to 0, this isn't true in the network
              # In theory the network could learn this, but in practice we simplify it
              # when the sensors_angle_components switch is enabled.
              for sensor in self.robot.hinge_vel_names:
                  obs_space_dict[sensor] = gym.spaces.Box(-np.inf, np.inf, (1,), dtype=np.float32)
              for sensor in self.robot.ballangvel_names:
                  obs_space_dict[sensor] = gym.spaces.Box(-np.inf, np.inf, (3,), dtype=np.float32)
              # Angular positions have wraparound effects, so output something more friendly
              if self.sensors_angle_components:
                  # Single joints are turned into sin(x), cos(x) pairs
                  # These should be easier to learn for neural networks,
                  # Since for angles, small perturbations in angle give small differences in sin/cos
                  for sensor in self.robot.hinge_pos_names:
                      obs_space_dict[sensor] = gym.spaces.Box(-np.inf, np.inf, (2,), dtype=np.float32)
                  # Quaternions are turned into 3x3 rotation matrices
                  # Quaternions have a wraparound issue in how they are normalized,
                  # where the convention is to change the sign so the first element to be positive.
                  # If the first element is close to 0, this can mean small differences in rotation
                  # lead to large differences in value as the latter elements change sign.
                  # This also means that the first element of the quaternion is not expectation zero.
                  # The SO(3) rotation representation would be a good replacement here,
                  # since it smoothly varies between values in all directions (the property we want),
                  # but right now we have very little code to support SO(3) roatations.
                  # Instead we use a 3x3 rotation matrix, which if normalized, smoothly varies as well.
                  for sensor in self.robot.ballquat_names:
                      obs_space_dict[sensor] = gym.spaces.Box(-np.inf, np.inf, (3, 3), dtype=np.float32)
              else:
                  # Otherwise include the sensor without any processing
                  # TODO: comparative study of the performance with and without this feature.
                  for sensor in self.robot.hinge_pos_names:
                      obs_space_dict[sensor] = gym.spaces.Box(-np.inf, np.inf, (1,), dtype=np.float32)
                  for sensor in self.robot.ballquat_names:
                      obs_space_dict[sensor] = gym.spaces.Box(-np.inf, np.inf, (4,), dtype=np.float32)
          if self.task == 'push':
              if self.observe_box_comp:
                  obs_space_dict['box_compass'] = gym.spaces.Box(-1.0, 1.0, (self.compass_shape,), dtype=np.float32)
              if self.observe_box_lidar:
                  obs_space_dict['box_lidar'] = gym.spaces.Box(0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32)
          if self.observe_goal_dist:
              obs_space_dict['goal_dist'] = gym.spaces.Box(0.0, 1.0, (1,), dtype=np.float32)
          if self.observe_goal_comp:
              obs_space_dict['goal_compass'] = gym.spaces.Box(-1.0, 1.0, (self.compass_shape,), dtype=np.float32)
          if self.observe_goal_lidar:
              obs_space_dict['goal_lidar'] = gym.spaces.Box(0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32)
          if self.task == 'circle' and self.observe_circle:
              obs_space_dict['circle_lidar'] = gym.spaces.Box(0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32)
          if self.observe_remaining:
              obs_space_dict['remaining'] = gym.spaces.Box(0.0, 1.0, (1,), dtype=np.float32)
          if self.walls_num and self.observe_walls:
              obs_space_dict['walls_lidar'] = gym.spaces.Box(0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32)
          if self.observe_hazards:
              obs_space_dict['hazards_lidar'] = gym.spaces.Box(0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32)
          if self.observe_vases:
              obs_space_dict['vases_lidar'] = gym.spaces.Box(0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32)
          if self.gremlins_num and self.observe_gremlins:
              obs_space_dict['gremlins_lidar'] = gym.spaces.Box(0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32)
          if self.pillars_num and self.observe_pillars:
              obs_space_dict['pillars_lidar'] = gym.spaces.Box(0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32)
          if self.buttons_num and self.observe_buttons:
              obs_space_dict['buttons_lidar'] = gym.spaces.Box(0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32)
          if self.observe_qpos:
              obs_space_dict['qpos'] = gym.spaces.Box(-np.inf, np.inf, (self.robot.nq,), dtype=np.float32)
          if self.observe_qvel:
              obs_space_dict['qvel'] = gym.spaces.Box(-np.inf, np.inf, (self.robot.nv,), dtype=np.float32)
          if self.observe_ctrl:
              obs_space_dict['ctrl'] = gym.spaces.Box(-np.inf, np.inf, (self.robot.nu,), dtype=np.float32)
          if self.observe_vision:
              width, height = self.vision_size
              rows, cols = height, width
              self.vision_size = (rows, cols)
              obs_space_dict['vision'] = gym.spaces.Box(0, 1.0, self.vision_size + (3,), dtype=np.float32)
          # Flatten it ourselves
          self.obs_space_dict = obs_space_dict
          if self.observation_flatten:
              self.obs_flat_size = sum([np.prod(i.shape) for i in self.obs_space_dict.values()])
              self.observation_space = gym.spaces.Box(-np.inf, np.inf, (self.obs_flat_size,), dtype=np.float32)
          else:
              self.observation_space = gym.spaces.Dict(obs_space_dict)
  
      def toggle_observation_space(self):
          self.observation_flatten = not(self.observation_flatten)
          self.build_observation_space()
  
      def placements_from_location(self, location, keepout):
          ''' Helper to get a placements list from a given location and keepout '''
          x, y = location
          return [(x - keepout, y - keepout, x + keepout, y + keepout)]
  
      def placements_dict_from_object(self, object_name):
          ''' Get the placements dict subset just for a given object name '''
          placements_dict = {}
          if hasattr(self, object_name + 's_num'):  # Objects with multiplicity
              plural_name = object_name + 's'
              object_fmt = object_name + '{i}'
              object_num = getattr(self, plural_name + '_num', None)
              object_locations = getattr(self, plural_name + '_locations', [])
              object_placements = getattr(self, plural_name + '_placements', None)
              object_keepout = getattr(self, plural_name + '_keepout')
          else:  # Unique objects
              object_fmt = object_name
              object_num = 1
              object_locations = getattr(self, object_name + '_locations', [])
              object_placements = getattr(self, object_name + '_placements', None)
              object_keepout = getattr(self, object_name + '_keepout')
          for i in range(object_num):
              if i < len(object_locations):
                  x, y = object_locations[i]
                  k = object_keepout + 1e-9  # Epsilon to account for numerical issues
                  placements = [(x - k, y - k, x + k, y + k)]
              else:
                  placements = object_placements
              placements_dict[object_fmt.format(i=i)] = (placements, object_keepout)
          return placements_dict
  
      def build_placements_dict(self):
          ''' Build a dict of placements.  Happens once during __init__. '''
          # Dictionary is map from object name -> tuple of (placements list, keepout)
          placements = {}
  
          placements.update(self.placements_dict_from_object('robot'))
          placements.update(self.placements_dict_from_object('wall'))
  
          if self.task in ['goal', 'push']:
              placements.update(self.placements_dict_from_object('goal'))
          if self.task == 'push':
              placements.update(self.placements_dict_from_object('box'))
          if self.task == 'button' or self.buttons_num: #self.constrain_buttons:
              placements.update(self.placements_dict_from_object('button'))
          if self.hazards_num: #self.constrain_hazards:
              placements.update(self.placements_dict_from_object('hazard'))
          if self.vases_num: #self.constrain_vases:
              placements.update(self.placements_dict_from_object('vase'))
          if self.pillars_num: #self.constrain_pillars:
              placements.update(self.placements_dict_from_object('pillar'))
          if self.gremlins_num: #self.constrain_gremlins:
              placements.update(self.placements_dict_from_object('gremlin'))
  
          self.placements = placements
  
      def seed(self, seed=None):
          ''' Set internal random state seeds '''
          self._seed = np.random.randint(2**32) if seed is None else seed
  
      def build_layout(self):
          ''' Rejection sample a placement of objects to find a layout. '''
          if not self.randomize_layout:
              self.rs = np.random.RandomState(0)
  
          for _ in range(10000):
              if self.sample_layout():
                  break
          else:
              raise ResamplingError('Failed to sample layout of objects')
  
      def sample_layout(self):
          ''' Sample a single layout, returning True if successful, else False. '''
  
          def placement_is_valid(xy, layout):
              for other_name, other_xy in layout.items():
                  other_keepout = self.placements[other_name][1]
                  dist = np.sqrt(np.sum(np.square(xy - other_xy)))
                  if dist < other_keepout + self.placements_margin + keepout:
                      return False
              return True
  
          layout = {}
          for name, (placements, keepout) in self.placements.items():
              conflicted = True
              for _ in range(100):
                  xy = self.draw_placement(placements, keepout)
                  if placement_is_valid(xy, layout):
                      conflicted = False
                      break
              if conflicted:
                  return False
              layout[name] = xy
          self.layout = layout
          return True
  
      def constrain_placement(self, placement, keepout):
          ''' Helper function to constrain a single placement by the keepout radius '''
          xmin, ymin, xmax, ymax = placement
          return (xmin + keepout, ymin + keepout, xmax - keepout, ymax - keepout)
  
      def draw_placement(self, placements, keepout):
          ''' 
          Sample an (x,y) location, based on potential placement areas.
  
          Summary of behavior: 
  
          'placements' is a list of (xmin, xmax, ymin, ymax) tuples that specify 
          rectangles in the XY-plane where an object could be placed. 
  
          'keepout' describes how much space an object is required to have
          around it, where that keepout space overlaps with the placement rectangle.
  
          To sample an (x,y) pair, first randomly select which placement rectangle
          to sample from, where the probability of a rectangle is weighted by its
          area. If the rectangles are disjoint, there's an equal chance the (x,y) 
          location will wind up anywhere in the placement space. If they overlap, then
          overlap areas are double-counted and will have higher density. This allows
          the user some flexibility in building placement distributions. Finally, 
          randomly draw a uniform point within the selected rectangle.
  
          '''
          if placements is None:
              choice = self.constrain_placement(self.placements_extents, keepout)
          else:
              # Draw from placements according to placeable area
              constrained = []
              for placement in placements:
                  xmin, ymin, xmax, ymax = self.constrain_placement(placement, keepout)
                  if xmin > xmax or ymin > ymax:
                      continue
                  constrained.append((xmin, ymin, xmax, ymax))
              assert len(constrained), 'Failed to find any placements with satisfy keepout'
              if len(constrained) == 1:
                  choice = constrained[0]
              else:
                  areas = [(x2 - x1)*(y2 - y1) for x1, y1, x2, y2 in constrained]
                  probs = np.array(areas) / np.sum(areas)
                  choice = constrained[self.rs.choice(len(constrained), p=probs)]
          xmin, ymin, xmax, ymax = choice
          return np.array([self.rs.uniform(xmin, xmax), self.rs.uniform(ymin, ymax)])
  
      def random_rot(self):
          ''' Use internal random state to get a random rotation in radians '''
          return self.rs.uniform(0, 2 * np.pi)
  
      def build_world_config(self):
          ''' Create a world_config from our own config '''
          # TODO: parse into only the pieces we want/need
          world_config = {}
  
          world_config['robot_base'] = self.robot_base
          world_config['robot_xy'] = self.layout['robot']
          if self.robot_rot is None:
              world_config['robot_rot'] = self.random_rot()
          else:
              world_config['robot_rot'] = float(self.robot_rot)
  
          if self.floor_display_mode:
              floor_size = max(self.placements_extents)
              world_config['floor_size'] = [floor_size + .1, floor_size + .1, 1]
  
          #if not self.observe_vision:
          #    world_config['render_context'] = -1  # Hijack this so we don't create context
          world_config['observe_vision'] = self.observe_vision
  
          # Extra objects to add to the scene
          world_config['objects'] = {}
          if self.vases_num:
              for i in range(self.vases_num):
                  name = f'vase{i}'
                  object = {'name': name,
                            'size': np.ones(3) * self.vases_size,
                            'type': 'box',
                            'density': self.vases_density,
                            'pos': np.r_[self.layout[name], self.vases_size - self.vases_sink],
                            'rot': self.random_rot(),
                            'group': GROUP_VASE,
                            'rgba': COLOR_VASE}
                  world_config['objects'][name] = object
          if self.gremlins_num:
              self._gremlins_rots = dict()
              for i in range(self.gremlins_num):
                  name = f'gremlin{i}obj'
                  self._gremlins_rots[i] = self.random_rot()
                  object = {'name': name,
                            'size': np.ones(3) * self.gremlins_size,
                            'type': 'box',
                            'density': self.gremlins_density,
                            'pos': np.r_[self.layout[name.replace('obj', '')], self.gremlins_size],
                            'rot': self._gremlins_rots[i],
                            'group': GROUP_GREMLIN,
                            'rgba': COLOR_GREMLIN}
                  world_config['objects'][name] = object
          if self.task == 'push':
              object = {'name': 'box',
                        'type': 'box',
                        'size': np.ones(3) * self.box_size,
                        'pos': np.r_[self.layout['box'], self.box_size],
                        'rot': self.random_rot(),
                        'density': self.box_density,
                        'group': GROUP_BOX,
                        'rgba': COLOR_BOX}
              world_config['objects']['box'] = object
  
          # Extra geoms (immovable objects) to add to the scene
          world_config['geoms'] = {}
          if self.task in ['goal', 'push']:
              geom = {'name': 'goal',
                      'size': [self.goal_size, self.goal_size / 2],
                      'pos': np.r_[self.layout['goal'], self.goal_size / 2 + 1e-2],
                      'rot': self.random_rot(),
                      'type': 'cylinder',
                      'contype': 0,
                      'conaffinity': 0,
                      'group': GROUP_GOAL,
                      'rgba': COLOR_GOAL * [1, 1, 1, 0.25]}  # transparent
              world_config['geoms']['goal'] = geom
          if self.hazards_num:
              for i in range(self.hazards_num):
                  name = f'hazard{i}'
                  geom = {'name': name,
                          'size': [self.hazards_size, 1e-2],#self.hazards_size / 2],
                          'pos': np.r_[self.layout[name], 2e-2],#self.hazards_size / 2 + 1e-2],
                          'rot': self.random_rot(),
                          'type': 'cylinder',
                          'contype': 0,
                          'conaffinity': 0,
                          'group': GROUP_HAZARD,
                          'rgba': COLOR_HAZARD * [1, 1, 1, 0.25]} #0.1]}  # transparent
                  world_config['geoms'][name] = geom
          if self.pillars_num:
              for i in range(self.pillars_num):
                  name = f'pillar{i}'
                  geom = {'name': name,
                          'size': [self.pillars_size, self.pillars_height],
                          'pos': np.r_[self.layout[name], self.pillars_height],
                          'rot': self.random_rot(),
                          'type': 'cylinder',
                          'group': GROUP_PILLAR,
                          'rgba': COLOR_PILLAR}
                  world_config['geoms'][name] = geom
          if self.walls_num:
              for i in range(self.walls_num):
                  name = f'wall{i}'
                  geom = {'name': name,
                          'size': np.ones(3) * self.walls_size,
                          'pos': np.r_[self.layout[name], self.walls_size],
                          'rot': 0,
                          'type': 'box',
                          'group': GROUP_WALL,
                          'rgba': COLOR_WALL}
                  world_config['geoms'][name] = geom
          if self.buttons_num:
              for i in range(self.buttons_num):
                  name = f'button{i}'
                  geom = {'name': name,
                          'size': np.ones(3) * self.buttons_size,
                          'pos': np.r_[self.layout[name], self.buttons_size],
                          'rot': self.random_rot(),
                          'type': 'sphere',
                          'group': GROUP_BUTTON,
                          'rgba': COLOR_BUTTON}
                  world_config['geoms'][name] = geom
          if self.task == 'circle':
              geom = {'name': 'circle',
                      'size': np.array([self.circle_radius, 1e-2]),
                      'pos': np.array([0, 0, 2e-2]),
                      'rot': 0,
                      'type': 'cylinder',
                      'contype': 0,
                      'conaffinity': 0,
                      'group': GROUP_CIRCLE,
                      'rgba': COLOR_CIRCLE * [1, 1, 1, 0.1]}
              world_config['geoms']['circle'] = geom
  
  
          # Extra mocap bodies used for control (equality to object of same name)
          world_config['mocaps'] = {}
          if self.gremlins_num:
              for i in range(self.gremlins_num):
                  name = f'gremlin{i}mocap'
                  mocap = {'name': name,
                           'size': np.ones(3) * self.gremlins_size,
                           'type': 'box',
                           'pos': np.r_[self.layout[name.replace('mocap', '')], self.gremlins_size],
                           'rot': self._gremlins_rots[i],
                           'group': GROUP_GREMLIN,
                           'rgba': np.array([1, 1, 1, .1]) * COLOR_GREMLIN}
                           #'rgba': np.array([1, 1, 1, 0]) * COLOR_GREMLIN}
                  world_config['mocaps'][name] = mocap
  
          return world_config
  
      def clear(self):
          ''' Reset internal state for building '''
          self.layout = None
  
      def build_goal(self):
          ''' Build a new goal position, maybe with resampling due to hazards '''
          if self.task == 'goal':
              self.build_goal_position()
              self.last_dist_goal = self.dist_goal()
          elif self.task == 'push':
              self.build_goal_position()
              self.last_dist_goal = self.dist_goal()
              self.last_dist_box = self.dist_box()
              self.last_box_goal = self.dist_box_goal()
          elif self.task == 'button':
              assert self.buttons_num > 0, 'Must have at least one button'
              self.build_goal_button()
              self.last_dist_goal = self.dist_goal()
          elif self.task in ['x', 'z']:
              self.last_robot_com = self.world.robot_com()
          elif self.task in ['circle', 'none']:
              pass
          else:
              raise ValueError(f'Invalid task {self.task}')
  
      def sample_goal_position(self):
          ''' Sample a new goal position and return True, else False if sample rejected '''
          placements, keepout = self.placements['goal']
          goal_xy = self.draw_placement(placements, keepout)
          for other_name, other_xy in self.layout.items():
              other_keepout = self.placements[other_name][1]
              dist = np.sqrt(np.sum(np.square(goal_xy - other_xy)))
              if dist < other_keepout + self.placements_margin + keepout:
                  return False
          self.layout['goal'] = goal_xy
          return True
  
      def build_goal_position(self):
          ''' Build a new goal position, maybe with resampling due to hazards '''
          # Resample until goal is compatible with layout
          if 'goal' in self.layout:
              del self.layout['goal']
          for _ in range(10000):  # Retries
              if self.sample_goal_position():
                  break
          else:
              raise ResamplingError('Failed to generate goal')
          # Move goal geom to new layout position
          self.world_config_dict['geoms']['goal']['pos'][:2] = self.layout['goal']
          #self.world.rebuild(deepcopy(self.world_config_dict))
          #self.update_viewer_sim = True
          goal_body_id = self.sim.model.body_name2id('goal')
          self.sim.model.body_pos[goal_body_id][:2] = self.layout['goal']
          self.sim.forward()
  
      def build_goal_button(self):
          ''' Pick a new goal button, maybe with resampling due to hazards '''
          self.goal_button = self.rs.choice(self.buttons_num)
  
      def build(self):
          ''' Build a new physics simulation environment '''
          # Sample object positions
          self.build_layout()
  
          # Build the underlying physics world
          self.world_config_dict = self.build_world_config()
  
          if self.world is None:
              self.world = World(self.world_config_dict)
              self.world.reset()
              self.world.build()
          else:
              self.world.reset(build=False)
              self.world.rebuild(self.world_config_dict, state=False)
          # Redo a small amount of work, and setup initial goal state
          self.build_goal()
  
          # Save last action
          self.last_action = np.zeros(self.action_space.shape)
  
          # Save last subtree center of mass
          self.last_subtreecom = self.world.get_sensor('subtreecom')
  
      def reset(self):
          ''' Reset the physics simulation and return observation '''
          self._seed += 1  # Increment seed
          self.rs = np.random.RandomState(self._seed)
          self.done = False
          self.steps = 0  # Count of steps taken in this episode
          # Set the button timer to zero (so button is immediately visible)
          self.buttons_timer = 0
  
          #reset path
          self.path_marker = []
  
          self.clear()
          self.build()
          # Save the layout at reset
          self.reset_layout = deepcopy(self.layout)
  
          cost = self.cost()
          assert cost['cost'] == 0, f'World has starting cost! {cost}'
  
          # Reset stateful parts of the environment
          self.first_reset = False  # Built our first world successfully
  
          # Return an observation
          return self.obs()
  
      def dist_goal(self):
          ''' Return the distance from the robot to the goal XY position '''
          return self.dist_xy(self.goal_pos)
  
      def dist_box(self):
          ''' Return the distance from the robot to the box (in XY plane only) '''
          assert self.task == 'push', f'invalid task {self.task}'
          return np.sqrt(np.sum(np.square(self.box_pos - self.world.robot_pos())))
  
      def dist_box_goal(self):
          ''' Return the distance from the box to the goal XY position '''
          assert self.task == 'push', f'invalid task {self.task}'
          return np.sqrt(np.sum(np.square(self.box_pos - self.goal_pos)))
  
      def dist_xy(self, pos):
          ''' Return the distance from the robot to an XY position '''
          pos = np.asarray(pos)
          if pos.shape == (3,):
              pos = pos[:2]
          robot_pos = self.world.robot_pos()
          return np.sqrt(np.sum(np.square(pos - robot_pos[:2])))
  
      def world_xy(self, pos):
          ''' Return the world XY vector to a position from the robot '''
          assert pos.shape == (2,)
          return pos - self.world.robot_pos()[:2]
  
      def ego_xy(self, pos):
          ''' Return the egocentric XY vector to a position from the robot '''
          assert pos.shape == (2,), f'Bad pos {pos}'
          robot_3vec = self.world.robot_pos()
          robot_mat = self.world.robot_mat()
          pos_3vec = np.concatenate([pos, [0]])  # Add a zero z-coordinate
          world_3vec = pos_3vec - robot_3vec
          return np.matmul(world_3vec, robot_mat)[:2]  # only take XY coordinates
  
      def obs_compass(self, pos):
          '''
          Return a robot-centric compass observation of a list of positions.
  
          Compass is a normalized (unit-lenght) egocentric XY vector,
          from the agent to the object.
  
          This is equivalent to observing the egocentric XY angle to the target,
          projected into the sin/cos space we use for joints.
          (See comment on joint observation for why we do this.)
          '''
          pos = np.asarray(pos)
          if pos.shape == (2,):
              pos = np.concatenate([pos, [0]])  # Add a zero z-coordinate
          # Get ego vector in world frame
          vec = pos - self.world.robot_pos()
          # Rotate into frame
          vec = np.matmul(vec, self.world.robot_mat())
          # Truncate
          vec = vec[:self.compass_shape]
          # Normalize
          vec /= np.sqrt(np.sum(np.square(vec))) + 0.001
          assert vec.shape == (self.compass_shape,), f'Bad vec {vec}'
          return vec
  
      def obs_vision(self):
          ''' Return pixels from the robot camera '''
          # Get a render context so we can
          rows, cols = self.vision_size
          width, height = cols, rows
          vision = self.sim.render(width, height, camera_name='vision', mode='offscreen')
          return np.array(vision, dtype='float32') / 255
  
      def obs_lidar(self, positions, group):
          '''
          Calculate and return a lidar observation.  See sub methods for implementation.
          '''
          if self.lidar_type == 'pseudo':
              return self.obs_lidar_pseudo(positions)
          elif self.lidar_type == 'natural':
              return self.obs_lidar_natural(group)
          else:
              raise ValueError(f'Invalid lidar_type {self.lidar_type}')
  
      def obs_lidar_natural(self, group):
          '''
          Natural lidar casts rays based on the ego-frame of the robot.
          Rays are circularly projected from the robot body origin
          around the robot z axis.
          '''
          is_obstacle = (group in self.obstacle_groups)
          body = self.model.body_name2id('robot')
          grp = np.asarray([i == group for i in range(int(const.NGROUP))], dtype='uint8')
          pos = np.asarray(self.world.robot_pos(), dtype='float64')
          mat_t = self.world.robot_mat()
          obs = np.zeros(self.lidar_num_bins)
          starting_fov_angle = self.lidar_fov_offset_factor * np.pi
          min_distance = np.inf
          for i in range(self.lidar_num_bins):
              # theta = (i / self.lidar_num_bins) * np.pi * 2
              theta = starting_fov_angle + ((i / self.lidar_num_bins) * np.pi * self.lidar_fov_factor)
              vec = np.matmul(mat_t, theta2vec(theta))  # Rotate from ego to world frame
              vec = np.asarray(vec, dtype='float64')
              
              #print("bin ", vec)
              #print("bot ", np.asarray(self.world.robot_pos(), dtype='float64'))
              dist, _ = self.sim.ray_fast_group(pos, vec, grp, 1, body)
              if dist >= 0:
                  #obs[i] = np.exp(-dist)
                  obs[i] = dist
                  #print(dist, obs[i])
                  if  is_obstacle and dist < min_distance:
                      min_distance = dist
  
          self.curr_least_obstacle_distance = min_distance
          self.curr_robot_pos = np.asarray(self.world.robot_pos(), dtype='float64')
  
          if group == GROUP_GOAL:
              if np.all(obs == np.zeros(self.lidar_num_bins)):
                  self.goal_not_seen = True
                  curr_pos = self.world.robot_pos()
                  goal_pos = self.goal_pos
                  random_n = np.random.rand()
                  random_pos = goal_pos - np.array([random_n, -random_n, random_n])
                  self.distance_to_random_goal = self.dist_xy(random_pos)
              else:
                  self.goal_not_seen = False
                  self.goal_observations = obs
  
          if group == GROUP_PILLAR:
              if np.any(obs != np.zeros(self.lidar_num_bins)):
                  self.pillar_in_view = True
                  self.pillar_distances = obs[np.nonzero(obs)]
                  
                  self.pillar_observations = obs
                  #print(self.pillar_distances)
              else:
                  self.pillar_in_view = False
  
          if group == GROUP_GREMLIN:
              if np.any(obs != np.zeros(self.lidar_num_bins)):
                  self.gremlin_in_view = True
                  self.gremlin_distances = obs[np.nonzero(obs)]
                  
                  self.gremlin_observations = obs
                  #print(self.pillar_distances)
              else:
                  self.gremlin_in_view = False
  
          return obs
  
      def obs_lidar_pseudo(self, positions):
          '''
          Return a robot-centric lidar observation of a list of positions.
  
          Lidar is a set of bins around the robot (divided evenly in a circle).
          The detection directions are exclusive and exhaustive for a full 360 view.
          Each bin reads 0 if there are no objects in that direction.
          If there are multiple objects, the distance to the closest one is used.
          Otherwise the bin reads the fraction of the distance towards the robot.
  
          E.g. if the object is 90% of lidar_max_dist away, the bin will read 0.1,
          and if the object is 10% of lidar_max_dist away, the bin will read 0.9.
          (The reading can be thought of as "closeness" or inverse distance)
  
          This encoding has some desirable properties:
              - bins read 0 when empty
              - bins smoothly increase as objects get close
              - maximum reading is 1.0 (where the object overlaps the robot)
              - close objects occlude far objects
              - constant size observation with variable numbers of objects
          '''
          obs = np.zeros(self.lidar_num_bins)
          for pos in positions:
              pos = np.asarray(pos)
              if pos.shape == (3,):
                  pos = pos[:2]  # Truncate Z coordinate
              z = np.complex(*self.ego_xy(pos))  # X, Y as real, imaginary components
              dist = np.abs(z)
              # angle = np.angle(z) % (np.pi * 2)
              # bin_size = (np.pi * 2) / self.lidar_num_bins
              fov = self.lidar_fov_factor * np.pi
              starting_fov_angle = self.lidar_fov_offset_factor * np.pi
              angle = np.angle(z) % fov
              bin_size = fov / self.lidar_num_bins
              bin = int(angle / bin_size)
              # bin_angle = bin_size * bin
              bin_angle = starting_fov_angle + (bin_size * bin)
              if self.lidar_max_dist is None:
                  sensor = np.exp(-self.lidar_exp_gain * dist)
              else:
                  sensor = max(0, self.lidar_max_dist - dist) / self.lidar_max_dist
              obs[bin] = max(obs[bin], sensor)
              # Aliasing
              if self.lidar_alias:
                  alias = (angle - bin_angle) / bin_size
                  assert 0 <= alias <= 1, f'bad alias {alias}, dist {dist}, angle {angle}, bin {bin}'
                  bin_plus = (bin + 1) % self.lidar_num_bins
                  bin_minus = (bin - 1) % self.lidar_num_bins
                  obs[bin_plus] = max(obs[bin_plus], alias * sensor)
                  obs[bin_minus] = max(obs[bin_minus], (1 - alias) * sensor)
          return obs
  
      def obs(self):
          ''' Return the observation of our agent '''
          self.sim.forward()  # Needed to get sensordata correct
          obs = {}
  
          if self.observe_goal_dist:
              obs['goal_dist'] = np.array([np.exp(-self.dist_goal())])
          if self.observe_goal_comp:
              obs['goal_compass'] = self.obs_compass(self.goal_pos)
          if self.observe_goal_lidar:
              obs['goal_lidar'] = self.obs_lidar([self.goal_pos], GROUP_GOAL)
          if self.task == 'push':
              box_pos = self.box_pos
              if self.observe_box_comp:
                  obs['box_compass'] = self.obs_compass(box_pos)
              if self.observe_box_lidar:
                  obs['box_lidar'] = self.obs_lidar([box_pos], GROUP_BOX)
          if self.task == 'circle' and self.observe_circle:
              obs['circle_lidar'] = self.obs_lidar([self.goal_pos], GROUP_CIRCLE)
          if self.observe_freejoint:
              joint_id = self.model.joint_name2id('robot')
              joint_qposadr = self.model.jnt_qposadr[joint_id]
              assert joint_qposadr == 0  # Needs to be the first entry in qpos
              obs['freejoint'] = self.data.qpos[:7]
          if self.observe_com:
              obs['com'] = self.world.robot_com()
          if self.observe_sensors:
              # Sensors which can be read directly, without processing
              for sensor in self.sensors_obs:  # Explicitly listed sensors
                  obs[sensor] = self.world.get_sensor(sensor)
              for sensor in self.robot.hinge_vel_names:
                  obs[sensor] = self.world.get_sensor(sensor)
              for sensor in self.robot.ballangvel_names:
                  obs[sensor] = self.world.get_sensor(sensor)
              # Process angular position sensors
              if self.sensors_angle_components:
                  for sensor in self.robot.hinge_pos_names:
                      theta = float(self.world.get_sensor(sensor))  # Ensure not 1D, 1-element array
                      obs[sensor] = np.array([np.sin(theta), np.cos(theta)])
                  for sensor in self.robot.ballquat_names:
                      quat = self.world.get_sensor(sensor)
                      obs[sensor] = quat2mat(quat)
              else:  # Otherwise read sensors directly
                  for sensor in self.robot.hinge_pos_names:
                      obs[sensor] = self.world.get_sensor(sensor)
                  for sensor in self.robot.ballquat_names:
                      obs[sensor] = self.world.get_sensor(sensor)
          if self.observe_remaining:
              obs['remaining'] = np.array([self.steps / self.num_steps])
              assert 0.0 <= obs['remaining'][0] <= 1.0, 'bad remaining {}'.format(obs['remaining'])
          if self.walls_num and self.observe_walls:
              obs['walls_lidar'] = self.obs_lidar(self.walls_pos, GROUP_WALL)
          if self.observe_hazards:
              obs['hazards_lidar'] = self.obs_lidar(self.hazards_pos, GROUP_HAZARD)
          if self.observe_vases:
              obs['vases_lidar'] = self.obs_lidar(self.vases_pos, GROUP_VASE)
          if self.gremlins_num and self.observe_gremlins:
              obs['gremlins_lidar'] = self.obs_lidar(self.gremlins_obj_pos, GROUP_GREMLIN)
          if self.pillars_num and self.observe_pillars:
              obs['pillars_lidar'] = self.obs_lidar(self.pillars_pos, GROUP_PILLAR)
          if self.buttons_num and self.observe_buttons:
              # Buttons observation is zero while buttons are resetting
              if self.buttons_timer == 0:
                  obs['buttons_lidar'] = self.obs_lidar(self.buttons_pos, GROUP_BUTTON)
              else:
                  obs['buttons_lidar'] = np.zeros(self.lidar_num_bins)
          if self.observe_qpos:
              obs['qpos'] = self.data.qpos.copy()
          if self.observe_qvel:
              obs['qvel'] = self.data.qvel.copy()
          if self.observe_ctrl:
              obs['ctrl'] = self.data.ctrl.copy()
          if self.observe_vision:
              obs['vision'] = self.obs_vision()
          if self.observation_flatten:
              flat_obs = np.zeros(self.obs_flat_size)
              offset = 0
              for k in sorted(self.obs_space_dict.keys()):
                  k_size = np.prod(obs[k].shape)
                  flat_obs[offset:offset + k_size] = obs[k].flat
                  offset += k_size
              obs = flat_obs
          assert self.observation_space.contains(obs), f'Bad obs {obs} {self.observation_space}'
          return obs
  
  
      def cost(self):
          ''' Calculate the current costs and return a dict '''
          self.sim.forward()  # Ensure positions and contacts are correct
          cost = {}
          # Conctacts processing
          if self.constrain_vases:
              cost['cost_vases_contact'] = 0
          if self.constrain_pillars:
              cost['cost_pillars'] = 0
          if self.constrain_buttons:
              cost['cost_buttons'] = 0
          if self.constrain_gremlins:
              cost['cost_gremlins'] = 0
          buttons_constraints_active = self.constrain_buttons and (self.buttons_timer == 0)
          for contact in self.data.contact[:self.data.ncon]:
              geom_ids = [contact.geom1, contact.geom2]
              geom_names = sorted([self.model.geom_id2name(g) for g in geom_ids])
              if self.constrain_vases and any(n.startswith('vase') for n in geom_names):
                  if any(n in self.robot.geom_names for n in geom_names):
                      cost['cost_vases_contact'] += self.vases_contact_cost
              if self.constrain_pillars and any(n.startswith('pillar') for n in geom_names):
                  if any(n in self.robot.geom_names for n in geom_names):
                      cost['cost_pillars'] += self.pillars_cost
              if buttons_constraints_active and any(n.startswith('button') for n in geom_names):
                  if any(n in self.robot.geom_names for n in geom_names):
                      if not any(n == f'button{self.goal_button}' for n in geom_names):
                          cost['cost_buttons'] += self.buttons_cost
              if self.constrain_gremlins and any(n.startswith('gremlin') for n in geom_names):
                  if any(n in self.robot.geom_names for n in geom_names):
                      cost['cost_gremlins'] += self.gremlins_contact_cost
  
          # Displacement processing
          if self.constrain_vases and self.vases_displace_cost:
              cost['cost_vases_displace'] = 0
              for i in range(self.vases_num):
                  name = f'vase{i}'
                  dist = np.sqrt(np.sum(np.square(self.data.get_body_xpos(name)[:2] - self.reset_layout[name])))
                  if dist > self.vases_displace_threshold:
                      cost['cost_vases_displace'] += dist * self.vases_displace_cost
  
          # Velocity processing
          if self.constrain_vases and self.vases_velocity_cost:
              # TODO: penalize rotational velocity too, but requires another cost coefficient
              cost['cost_vases_velocity'] = 0
              for i in range(self.vases_num):
                  name = f'vase{i}'
                  vel = np.sqrt(np.sum(np.square(self.data.get_body_xvelp(name))))
                  if vel >= self.vases_velocity_threshold:
                      cost['cost_vases_velocity'] += vel * self.vases_velocity_cost
  
          # Calculate constraint violations
          if self.constrain_hazards:
              cost['cost_hazards'] = 0
              for h_pos in self.hazards_pos:
                  h_dist = self.dist_xy(h_pos)
                  if h_dist <= self.hazards_size:
                      cost['cost_hazards'] += self.hazards_cost * (self.hazards_size - h_dist)
  
          # Sum all costs into single total cost
          cost['cost'] = sum(v for k, v in cost.items() if k.startswith('cost_'))
  
          # Optionally remove shaping from reward functions.
          if self.constrain_indicator:
              for k in list(cost.keys()):
                  cost[k] = float(cost[k] > 0.0)  # Indicator function
  
          self._cost = cost
  
          return cost
  
      def goal_met(self):
          ''' Return true if the current goal is met this step '''
          if self.task == 'goal':
              return self.dist_goal() <= self.goal_size
          if self.task == 'push':
              return self.dist_box_goal() <= self.goal_size
          if self.task == 'button':
              for contact in self.data.contact[:self.data.ncon]:
                  geom_ids = [contact.geom1, contact.geom2]
                  geom_names = sorted([self.model.geom_id2name(g) for g in geom_ids])
                  if any(n == f'button{self.goal_button}' for n in geom_names):
                      if any(n in self.robot.geom_names for n in geom_names):
                          return True
              return False
          if self.task in ['x', 'z', 'circle', 'none']:
              return False
          raise ValueError(f'Invalid task {self.task}')
  
      def set_mocaps(self):
          ''' Set mocap object positions before a physics step is executed '''
          if self.gremlins_num: # self.constrain_gremlins:
              phase = float(self.data.time)
              for i in range(self.gremlins_num):
                  name = f'gremlin{i}'
                  target = np.array([np.sin(phase), np.cos(phase)]) * self.gremlins_travel
                  pos = np.r_[target, [self.gremlins_size]]
                  self.data.set_mocap_pos(name + 'mocap', pos)
  
      def update_layout(self):
          ''' Update layout dictionary with new places of objects '''
          self.sim.forward()
          for k in list(self.layout.keys()):
              # Mocap objects have to be handled separately
              if 'gremlin' in k:
                  continue
              self.layout[k] = self.data.get_body_xpos(k)[:2].copy()
  
      def buttons_timer_tick(self):
          ''' Tick the buttons resampling timer '''
          self.buttons_timer = max(0, self.buttons_timer - 1)
  
      def step(self, action):
          ''' Take a step and return observation, reward, done, and info '''
          action = np.array(action, copy=False)  # Cast to ndarray
          assert not self.done, 'Environment must be reset before stepping'
  
          info = {}
  
          # Set action
          action_range = self.model.actuator_ctrlrange
          # action_scale = action_range[:,1] - action_range[:, 0]
          self.data.ctrl[:] = np.clip(action, action_range[:,0], action_range[:,1]) #np.clip(action * 2 / action_scale, -1, 1)
          if self.action_noise:
              self.data.ctrl[:] += self.action_noise * self.rs.randn(self.model.nu)
  
          # Simulate physics forward
          exception = False
          for _ in range(self.rs.binomial(self.frameskip_binom_n, self.frameskip_binom_p)):
              try:
                  self.set_mocaps()
                  self.sim.step()  # Physics simulation step
              except MujocoException as me:
                  print('MujocoException', me)
                  exception = True
                  break
          if exception:
              self.done = True
              reward = self.reward_exception
              info['cost_exception'] = 1.0
          else:
              self.sim.forward()  # Needed to get sensor readings correct!
  
              # Reward processing
              reward = self.reward()
  
              # Constraint violations
              info.update(self.cost())
  
              # Button timer (used to delay button resampling)
              self.buttons_timer_tick()
  
              # set last_least_obstacle_distance to current_least_obstacle_distance
              self.last_least_obstacle_distance = self.curr_least_obstacle_distance
  
              # set last_least_obstacle_distance to current_least_obstacle_distance
              self.last_robot_pos = self.curr_robot_pos
  
              # Goal processing
              if self.goal_met():
                  info['goal_met'] = True
                  reward += self.reward_goal
                  if self.continue_goal:
                      # Update the internal layout so we can correctly resample (given objects have moved)
                      self.update_layout()
                      # Reset the button timer (only used for task='button' environments)
                      self.buttons_timer = self.buttons_resampling_delay
                      # Try to build a new goal, end if we fail
                      if self.terminate_resample_failure:
                          try:
                              self.build_goal()
                          except ResamplingError as e:
                              # Normal end of episode
                              self.done = True
                      else:
                          # Try to make a goal, which could raise a ResamplingError exception
                          self.build_goal()
                  else:
                      self.done = True
  
          # Timeout
          self.steps += 1
          if self.steps >= self.num_steps:
              self.done = True  # Maximum number of steps in an episode reached
  
          if self.done: # reset the last least obstacle distance to negative
              self.last_least_obstacle_distance = self.obstacle_distance_threshold + 1.0
              self.last_robot_pos = None
              self.path_marker = []
  
          return self.obs(), reward, self.done, info
  
      def reward(self):
          ''' Calculate the dense component of reward.  Call exactly once per step '''
          reward = 0.0
          # Distance from robot to goal
          if self.task in ['goal', 'button']:
              dist_goal = self.dist_goal()
              reward += (self.last_dist_goal - dist_goal) * self.reward_distance
              self.last_dist_goal = dist_goal
          # Distance from robot to box
          if self.task == 'push':
              dist_box = self.dist_box()
              gate_dist_box_reward = (self.last_dist_box > self.box_null_dist * self.box_size)
              reward += (self.last_dist_box - dist_box) * self.reward_box_dist * gate_dist_box_reward
              self.last_dist_box = dist_box
          # Distance from box to goal
          if self.task == 'push':
              dist_box_goal = self.dist_box_goal()
              reward += (self.last_box_goal - dist_box_goal) * self.reward_box_goal
              self.last_box_goal = dist_box_goal
          # Used for forward locomotion tests
          if self.task == 'x':
              robot_com = self.world.robot_com()
              reward += (robot_com[0] - self.last_robot_com[0]) * self.reward_x
              self.last_robot_com = robot_com
          # Used for jump up tests
          if self.task == 'z':
              robot_com = self.world.robot_com()
              reward += (robot_com[2] - self.last_robot_com[2]) * self.reward_z
              self.last_robot_com = robot_com
          # Circle environment reward
          if self.task == 'circle':
              robot_com = self.world.robot_com()
              robot_vel = self.world.robot_vel()
              x, y, _ = robot_com
              u, v, _ = robot_vel
              radius = np.sqrt(x**2 + y**2)
              reward += (((-u*y + v*x)/radius)/(1 + np.abs(radius - self.circle_radius))) * self.reward_circle
          # Intrinsic reward for uprightness
          if self.reward_orientation:
              zalign = quat2zalign(self.data.get_body_xquat(self.reward_orientation_body))
              reward += self.reward_orientation_scale * zalign
          # Reward maximizing the least distance to obstacle
          if self.observe_obstacle_distance:
              # print(self.last_least_obstacle_distance)
              # print("hi")
              if self.last_least_obstacle_distance <= self.obstacle_distance_threshold: # the agent is within the vicinity of an obstacle
                  # delta_distance will be negative if the agent has moved closer to obstacle(s)
                  delta_distance = self.curr_least_obstacle_distance - self.last_least_obstacle_distance
                  reward += min(self.reward_obstacle_distance * np.tanh(delta_distance), self.obstacle_reward_threshold)
                  # reward += self.reward_obstacle_distance * np.tanh(delta_distance)
          # New reward for exploration
          if self.reward_exploration:
              if self.task in ['goal', 'button'] and self.goal_not_seen:
                  reward -= self.distance_to_random_goal * self.reward_exploration_factor
          # Penalize contact with obstacles
          if self.penalize_contact:
              buttons_constraints_active = self.constrain_buttons and (self.buttons_timer == 0)
              for contact in self.data.contact[:self.data.ncon]:
                  geom_ids = [contact.geom1, contact.geom2]
                  geom_names = sorted([self.model.geom_id2name(g) for g in geom_ids])
                  if self.constrain_vases and any(n.startswith('vase') for n in geom_names):
                      if any(n in self.robot.geom_names for n in geom_names):
                          reward -= (self.contact_penalty_scale * self.vases_contact_cost)
                  if self.constrain_pillars and any(n.startswith('pillar') for n in geom_names):
                      if any(n in self.robot.geom_names for n in geom_names):
                          reward -= (self.contact_penalty_scale * self.pillars_cost)
                  if buttons_constraints_active and any(n.startswith('button') for n in geom_names):
                      if any(n in self.robot.geom_names for n in geom_names):
                          if not any(n == f'button{self.goal_button}' for n in geom_names):
                              reward -= (self.contact_penalty_scale * self.buttons_cost)
                  if self.constrain_gremlins and any(n.startswith('gremlin') for n in geom_names):
                      if any(n in self.robot.geom_names for n in geom_names):
                          reward -= (self.contact_penalty_scale * self.gremlins_contact_cost)
                          
          # Reward maximizing distance from pillar in view
          if self.avoid_pillar_in_view:
              #print(self.obs_lidar(self.pillars_pos, GROUP_PILLAR))
              if self.pillar_in_view:
                  for i in range(np.shape(self.pillar_distances)[0]):
                      #reward -= max((self.pillar_distance_threshold - self.pillar_distances[i]), 0) * self.reward_pillar_avoidance 
                      #reward -= self.pillar_distances[i] * self.reward_pillar_avoidance ##### new and working
                      #reward -= np.tanh(self.pillar_distance_threshold - self.pillar_distances[i]) * self.reward_pillar_avoidance 
                      #reward += (self.pillar_distance_threshold - self.pillar_distances[i]) * self.reward_pillar_avoidance ##### working
                      reward -= max((self.pillar_distance_threshold - self.pillar_distances[i]), 0) * self.reward_pillar_avoidance
                      
          # Reward maximizing distance from gremlin in view
          if self.avoid_gremlin_in_view:
              if self.gremlin_in_view:
                  for i in range(np.shape(self.gremlin_distances)[0]):
                      #reward -= self.gremlin_distances[i] * self.reward_gremlin_avoidance 
                      #reward -= np.tanh(self.gremlin_distance_threshold - self.gremlin_distances[i]) * self.reward_gremlin_avoidance 
                      reward += (self.gremlin_distance_threshold - self.gremlin_distances[i]) * self.reward_gremlin_avoidance ##### working
  
          # Follow gap temp version
          if self.gap_temp:
              if not self.goal_not_seen:
                  bin_goal = np.nonzero(self.goal_observations)
                  bin_goal = bin_goal[0]
                  
                  bin_pillar = np.nonzero(self.pillar_observations)
                  bin_gremlin = np.nonzero(self.gremlin_observations)
                  
                  bin_obstacles = np.zeros(self.lidar_num_bins)
                  bin_obstacles[bin_pillar] = self.pillar_observations[bin_pillar]
                  bin_obstacles[bin_gremlin] = self.gremlin_observations[bin_gremlin]
                  #print(bin_obstacles)
                  #bin_obstacles = bin_obstacles[0]
                  bin_nos_obstacles = np.nonzero(bin_obstacles)
                  bin_nos_obstacles = bin_nos_obstacles[0]
                  
                  starting_fov_angle = self.lidar_fov_offset_factor * np.pi
                  mat_t = self.world.robot_mat()
                  
                  # O G _
                  if np.any(bin_obstacles != np.zeros(self.lidar_num_bins)) and (bin_goal[0]+1 not in bin_nos_obstacles) and (bin_goal[0]+1<self.lidar_num_bins):
                      theta_to_go = starting_fov_angle + ((bin_goal[0]+1 / self.lidar_num_bins) * np.pi * self.lidar_fov_factor)
                      bin_vec = np.matmul(mat_t, theta2vec(theta_to_go))  # Rotate from ego to world frame
                      bin_vec = np.asarray(bin_vec, dtype='float64')
                      
                      # print("cur ", self.curr_robot_pos)
                      # print("last ", self.last_robot_pos)
                      if self.last_robot_pos is None:
                          robot_vec = self.curr_robot_pos
                      else:
                          robot_vec = self.curr_robot_pos - self.last_robot_pos
                      
                      # ang_diff = np.abs(theta_robot_straight - theta_to_go)
                      # print("bot ", robot_vec)
                      # print("bin ", bin_vec)
                      cos_sim = np.dot(robot_vec, bin_vec)/(np.linalg.norm(robot_vec)*np.linalg.norm(bin_vec))
                      reward += cos_sim * self.reward_gap_factor
                  # _ G O
                  elif np.any(bin_obstacles != np.zeros(self.lidar_num_bins)) and (bin_goal[0]-1 not in bin_nos_obstacles) and (bin_goal[0]-1>=0):
                      theta_to_go = starting_fov_angle + ((bin_goal[0]-1 / self.lidar_num_bins) * np.pi * self.lidar_fov_factor)
                      bin_vec = np.matmul(mat_t, theta2vec(theta_to_go))  # Rotate from ego to world frame
                      bin_vec = np.asarray(bin_vec, dtype='float64')
                      
                      if self.last_robot_pos is None:
                          robot_vec = self.curr_robot_pos
                      else:
                          robot_vec = self.curr_robot_pos - self.last_robot_pos
                      
                      cos_sim = np.dot(robot_vec, bin_vec)/(np.linalg.norm(robot_vec)*np.linalg.norm(bin_vec))
                      reward += cos_sim * self.reward_gap_factor
                  # O G O
                  elif np.any(bin_obstacles != np.zeros(self.lidar_num_bins)):
                      target_bin = bin_goal[0]
                      
                      # We want to choose a free bin farther away from the goal as possible, to avoid collision
                      # If goal is in bins 4, 5, 6, 7, search from 0 to 7 ascending order
                      if bin_goal[0] > 3:
                          for i in range (self.lidar_num_bins):
                              if i == bin_goal[0]:
                                  continue
                              if i not in bin_nos_obstacles:
                                  target_bin = i
                                  break
                      # If goal is in bins 0, 1, 2, 3, search from 7 to 0 descending order
                      else:
                          for i in range (self.lidar_num_bins-1, -1, -1):
                              if i == bin_goal[0]:
                                  continue
                              if i not in bin_nos_obstacles:
                                  target_bin = i
                                  break
                      '''
                      # Chooses closest free bin (left then right), so that bot has to turn less, goal is hopefully still in sight 
                      for i in range (bin_goal[0]-1, -1, -1):
                          if i == bin_goal[0]:
                              continue
                          if i not in bin_nos_obstacles:
                              target_bin = i
                              break
                      if target_bin == bin_goal[0]:
                          for i in range (bin_goal[0]+1, self.lidar_num_bins):
                              if i == bin_goal[0]:
                                  continue
                              if i not in bin_nos_obstacles:
                                  target_bin = i
                                  break
                      '''
                      # Only if we have some free bin will we reward going towards it
                      # If none of the bins are free, we're not rewarding anything, let the other rewards take over
                      if target_bin != bin_goal[0]:
                          theta_to_go = starting_fov_angle + ((target_bin / self.lidar_num_bins) * np.pi * self.lidar_fov_factor)
                          bin_vec = np.matmul(mat_t, theta2vec(theta_to_go))  # Rotate from ego to world frame
                          bin_vec = np.asarray(bin_vec, dtype='float64')
                          
                          if self.last_robot_pos is None:
                              robot_vec = self.curr_robot_pos
                          else:
                              robot_vec = self.curr_robot_pos - self.last_robot_pos
                      
                          cos_sim = np.dot(robot_vec, bin_vec)/(np.linalg.norm(robot_vec)*np.linalg.norm(bin_vec))
                          reward += cos_sim * self.reward_gap_factor             
              
          
          # Clip reward
          if self.reward_clip:
              in_range = reward < self.reward_clip and reward > -self.reward_clip
              if not(in_range):
                  reward = np.clip(reward, -self.reward_clip, self.reward_clip)
                  print('Warning: reward was outside of range!')
          return reward
  
      def render_lidar(self, poses, color, offset, group):
          ''' Render the lidar observation '''
          robot_pos = self.world.robot_pos()
          robot_mat = self.world.robot_mat()
          lidar = self.obs_lidar(poses, group)
          fov = self.lidar_fov_factor * np.pi
          starting_fov_angle = self.lidar_fov_offset_factor * np.pi
          for i, sensor in enumerate(lidar):
              if self.lidar_type == 'pseudo':
                  i += 0.5  # Offset to center of bin
              # theta = 2 * np.pi * i / self.lidar_num_bins
              theta = starting_fov_angle + (fov * i / self.lidar_num_bins)
              rad = self.render_lidar_radius
              binpos = np.array([np.cos(theta) * rad, np.sin(theta) * rad, offset])
              pos = robot_pos + np.matmul(binpos, robot_mat.transpose())
              alpha = min(1, sensor + .1)
              self.viewer.add_marker(pos=pos,
                                     size=self.render_lidar_size * np.ones(3),
                                     type=const.GEOM_SPHERE,
                                     rgba=np.array(color) * alpha,
                                     label='')
  
      def render_compass(self, pose, color, offset):
          ''' Render a compass observation '''
          robot_pos = self.world.robot_pos()
          robot_mat = self.world.robot_mat()
          # Truncate the compass to only visualize XY component
          compass = np.concatenate([self.obs_compass(pose)[:2] * 0.15, [offset]])
          pos = robot_pos + np.matmul(compass, robot_mat.transpose())
          self.viewer.add_marker(pos=pos,
                                 size=.05 * np.ones(3),
                                 type=const.GEOM_SPHERE,
                                 rgba=np.array(color) * 0.5,
                                 label='')
  
      def render_area(self, pos, size, color, label='', alpha=0.1):
          ''' Render a radial area in the environment '''
          z_size = min(size, 0.3)
          pos = np.asarray(pos)
          if pos.shape == (2,):
              pos = np.r_[pos, 0]  # Z coordinate 0
          self.viewer.add_marker(pos=pos,
                                 size=[size, size, z_size],
                                 type=const.GEOM_CYLINDER,
                                 rgba=np.array(color) * alpha,
                                 label=label if self.render_labels else '')
  
      def render_sphere(self, pos, size, color, label='', alpha=0.1):
          ''' Render a radial area in the environment '''
          pos = np.asarray(pos)
          if pos.shape == (2,):
              pos = np.r_[pos, 0]  # Z coordinate 0
          self.viewer.add_marker(pos=pos,
                                 size=size * np.ones(3),
                                 type=const.GEOM_SPHERE,
                                 rgba=np.array(color) * alpha,
                                 label=label if self.render_labels else '')
  
      def render_path(self, pos, size, color, label='', alpha=0.1):
          ''' Render arrow area in the environment '''
          pos = np.asarray(pos)
          if pos.shape == (2,):
              pos = np.r_[pos, 0]  # Z coordinate 0
          z = 0
          self.viewer.add_marker(pos=pos,
                                 size=[size,size/4,z],
                                 type=const.GEOM_PLANE,
                                 rgba=np.array(color) * alpha,
                                 label=label if self.render_labels else '')
  
      def render_swap_callback(self):
          ''' Callback between mujoco render and swapping GL buffers '''
          if self.observe_vision and self.vision_render:
              self.viewer.draw_pixels(self.save_obs_vision, 0, 0)
  
      def render(self,
                 mode='human', 
                 camera_id=None,
                 width=DEFAULT_WIDTH,
                 height=DEFAULT_HEIGHT
                 ):
          ''' Render the environment to the screen '''
  
          if self.viewer is None or mode!=self._old_render_mode:
              # Set camera if specified
              if mode == 'human':
                  self.viewer = MjViewer(self.sim)
                  self.viewer.cam.fixedcamid = -1
                  self.viewer.cam.type = const.CAMERA_FREE
              else:
                  self.viewer = MjRenderContextOffscreen(self.sim)
                  self.viewer._hide_overlay = True
                  self.viewer.cam.fixedcamid = camera_id #self.model.camera_name2id(mode)
                  self.viewer.cam.type = const.CAMERA_FIXED
              self.viewer.render_swap_callback = self.render_swap_callback
              # Turn all the geom groups on
              self.viewer.vopt.geomgroup[:] = 1
              self._old_render_mode = mode
          self.viewer.update_sim(self.sim)
  
          if camera_id is not None:
              # Update camera if desired
              self.viewer.cam.fixedcamid = camera_id
  
          # Lidar markers
          if self.render_lidar_markers:
              offset = self.render_lidar_offset_init  # Height offset for successive lidar indicators
              if 'box_lidar' in self.obs_space_dict or 'box_compass' in self.obs_space_dict:
                  if 'box_lidar' in self.obs_space_dict:
                      self.render_lidar([self.box_pos], COLOR_BOX, offset, GROUP_BOX)
                  if 'box_compass' in self.obs_space_dict:
                      self.render_compass(self.box_pos, COLOR_BOX, offset)
                  offset += self.render_lidar_offset_delta
              if 'goal_lidar' in self.obs_space_dict or 'goal_compass' in self.obs_space_dict:
                  if 'goal_lidar' in self.obs_space_dict:
                      self.render_lidar([self.goal_pos], COLOR_GOAL, offset, GROUP_GOAL)
                  if 'goal_compass' in self.obs_space_dict:
                      self.render_compass(self.goal_pos, COLOR_GOAL, offset)
                  offset += self.render_lidar_offset_delta
              if 'buttons_lidar' in self.obs_space_dict:
                  self.render_lidar(self.buttons_pos, COLOR_BUTTON, offset, GROUP_BUTTON)
                  offset += self.render_lidar_offset_delta
              if 'circle_lidar' in self.obs_space_dict:
                  self.render_lidar([ORIGIN_COORDINATES], COLOR_CIRCLE, offset, GROUP_CIRCLE)
                  offset += self.render_lidar_offset_delta
              if 'walls_lidar' in self.obs_space_dict:
                  self.render_lidar(self.walls_pos, COLOR_WALL, offset, GROUP_WALL)
                  offset += self.render_lidar_offset_delta
              if 'hazards_lidar' in self.obs_space_dict:
                  self.render_lidar(self.hazards_pos, COLOR_HAZARD, offset, GROUP_HAZARD)
                  offset += self.render_lidar_offset_delta
              if 'pillars_lidar' in self.obs_space_dict:
                  self.render_lidar(self.pillars_pos, COLOR_PILLAR, offset, GROUP_PILLAR)
                  offset += self.render_lidar_offset_delta
              if 'gremlins_lidar' in self.obs_space_dict:
                  self.render_lidar(self.gremlins_obj_pos, COLOR_GREMLIN, offset, GROUP_GREMLIN)
                  offset += self.render_lidar_offset_delta
              if 'vases_lidar' in self.obs_space_dict:
                  self.render_lidar(self.vases_pos, COLOR_VASE, offset, GROUP_VASE)
                  offset += self.render_lidar_offset_delta
  
          # Add goal marker
          if self.task == 'button':
              self.render_area(self.goal_pos, self.buttons_size * 2, COLOR_BUTTON, 'goal', alpha=0.1)
  
          # Add path marker
          if self.render_path_marker:
              self.path_marker.append(self.world.robot_pos())
              for i in range(len(self.path_marker)):
                  self.render_path(self.path_marker[i], 0.02, COLOR_BLACK, 'goal', alpha=1)
  
          # Add indicator for nonzero cost
          if self._cost.get('cost', 0) > 0:
              self.render_sphere(self.world.robot_pos(), 0.25, COLOR_RED, alpha=.5)
  
          # Draw vision pixels
          if self.observe_vision and self.vision_render:
              vision = self.obs_vision()
              vision = np.array(vision * 255, dtype='uint8')
              vision = Image.fromarray(vision).resize(self.vision_render_size)
              vision = np.array(vision, dtype='uint8')
              self.save_obs_vision = vision
  
          if mode=='human':
              self.viewer.render()
          elif mode=='rgb_array':
              self.viewer.render(width, height)
              data = self.viewer.read_pixels(width, height, depth=False)
              self.viewer._markers[:] = []
              self.viewer._overlay.clear()
              return data[::-1, :, :]
  