Esempio n. 1
0
class SimManager(object):
    """Object to generate sequence of images with their transforms"""
    def __init__(self,
                 filepath,
                 random_params={},
                 gpu_render=False,
                 gui=False,
                 display_data=False):
        self.model = load_model_from_path(filepath)
        self.sim = MjSim(self.model)
        self.filepath = filepath
        self.gui = gui
        self.display_data = display_data
        # Take the default random params and update anything we need
        self.RANDOM_PARAMS = {}
        self.RANDOM_PARAMS.update(random_params)

        if gpu_render:
            self.viewer = MjViewer(self.sim)
        else:
            self.viewer = None

        # Get start state of params to slightly jitter later
        self.START_GEOM_POS = self.model.geom_pos.copy()
        self.START_GEOM_SIZE = self.model.geom_size.copy()
        self.START_GEOM_QUAT = self.model.geom_quat.copy()
        self.START_BODY_POS = self.model.body_pos.copy()
        self.START_BODY_QUAT = self.model.body_quat.copy()
        self.START_MATID = self.model.geom_matid.copy()
        #self.FLOOR_OFFSET = self.model.body_pos[self.model.body_name2id('floor')]

        self.tex_modder = TextureModder(self.sim)
        self.tex_modder.whiten_materials(
        )  # ensures materials won't impact colors
        self.cam_modder = CameraModder(self.sim)
        self.light_modder = LightModder(self.sim)
        self.start_obj_pose = self.sim.data.get_joint_qpos(
            'object:joint').copy()

    def get_data(self, num_images=10):
        """
        Returns camera intrinsics, and a sequence of images, pose transforms, and
        camera transforms
        """
        # randomize the scene
        self._rand_textures()
        self._rand_lights()
        self._rand_object()
        self._rand_walls()
        self._rand_distract()
        sequence = defaultdict(list)
        context = {}
        # object pose
        obj_pose, robot_pose = self._get_ground_truth()
        context["obj_world_pose"] = obj_pose
        context["robot_world_pose"] = robot_pose
        self._cam_step = 0
        self._cam_choices = np.array([[-1.75, 0, 2], [-1.75, 0, 1.62],
                                      [-1.3, 1.7, 1.62]])
        self._curr_cam_pos = self._cam_choices[0]
        for i in range(num_images):
            self._next_camera()
            self._forward()
            img = self._get_cam_frame()
            sequence["img"].append(img)
            cam_pos = self.cam_modder.get_pos("camera1")
            cam_quat = self.cam_modder.get_quat("camera1")
            cam_pose = np.concatenate([cam_pos, cam_quat]).astype(np.float32)
            sequence["cam_pose"].append(cam_pose)

        cam_id = self.cam_modder.get_camid("camera1")
        fovy = self.sim.model.cam_fovy[cam_id]
        width, height = 640, 480
        f = 0.5 * height / math.tan(fovy * math.pi / 360)
        camera_intrinsics = np.array(
            ((f, 0, width / 2), (0, f, height / 2), (0, 0, 1)))
        context['cam_matrix'] = camera_intrinsics
        return context, sequence

    def _forward(self):
        """Advances simulator a step (NECESSARY TO MAKE CAMERA AND LIGHT MODDING WORK)
        And add some visualization"""
        self.sim.forward()
        if self.viewer and self.gui:
            # Get angle of camera and display it
            quat = np.quaternion(*self.model.cam_quat[0])
            ypr = quaternion.as_euler_angles(quat) * 180 / np.pi
            cam_pos = self.model.cam_pos[0]
            cam_fovy = self.model.cam_fovy[0]
            #self.viewer.add_marker(pos=cam_pos, label="CAM: {}{}".format(cam_pos, ypr))
            #self.viewer.add_marker(pos=cam_pos, label="CAM: {}".format(ypr))
            #self.viewer.add_marker(pos=cam_pos, label="CAM: {}".format(cam_pos))
            self.viewer.add_marker(pos=cam_pos,
                                   label="FOVY: {}, CAM: {}".format(
                                       cam_fovy, cam_pos))
            self.viewer.render()

    def _get_ground_truth(self):
        """
        Return the  position to the robot, and quaternion to the robot quaternion
        7 dim total
        """
        robot_gid = self.sim.model.geom_name2id('base_link')
        obj_gid = self.sim.model.geom_name2id('object')
        # only x and y pos needed
        obj_world_pos = self.sim.data.geom_xpos[obj_gid]
        robot_world_pos = self.sim.data.geom_xpos[robot_gid]
        # obj_pos_in_robot_frame = (self.sim.data.geom_xpos[obj_gid] - self.sim.data.geom_xpos[robot_gid])[:2]

        # robot_quat = quaternion.as_quat_array(self.model.geom_quat[robot_gid].copy())
        obj_world_quat = self.model.geom_quat[obj_gid].copy()
        robot_world_quat = self.model.geom_quat[robot_gid].copy()
        # # want quat of obj relative to robot frame
        # # obj_q = robot_q * localrot
        # # robot_q.inv * obj_q = localrot
        # rel_quat = quaternion.as_float_array(robot_quat.inverse() * obj_quat)
        # pose = np.concatenate([obj_pos_in_robot_frame, rel_quat]).astype(np.float32)
        obj_pose = self.sim.data.get_joint_qpos('object:joint').copy()
        robot_pose = np.concatenate([robot_world_pos,
                                     robot_world_quat]).astype(np.float32)
        return obj_pose, robot_pose

    def _get_cam_frame(self, ground_truth=None):
        """Grab an image from the camera (224, 244, 3) to feed into CNN"""
        #IMAGE_NOISE_RVARIANCE = Range(0.0, 0.0001)
        cam_img = self.sim.render(
            640, 480, camera_name='camera1'
        )[::-1, :, :]  # Rendered images are upside-down.
        # make camera crop be more like kinect
        #cam_img = self.sim.render(854, 480, camera_name='camera1')[::-1, 107:-107, :] # Rendered images are upside-down.

        #image_noise_variance = sample(IMAGE_NOISE_RVARIANCE)
        #cam_img = (skimage.util.random_noise(cam_img, mode='gaussian', var=image_noise_variance) * 255).astype(np.uint8)

        if self.display_data:
            print(ground_truth)
            #label = str(ground_truth[3:6])
            display_image(cam_img, mode='preproc')  #, label)

        # cam_img = preproc_image(cam_img)
        return cam_img

    def _randomize(self):
        self._rand_textures()
        self._rand_camera()
        self._rand_lights()
        #self._rand_robot()
        self._rand_object()
        self._rand_walls()
        self._rand_distract()

    def _rand_textures(self):
        """Randomize all the textures in the scene, including the skybox"""
        bright = np.random.binomial(1, 0.5)
        for name in self.sim.model.geom_names + ('skybox', ):
            self.tex_modder.rand_all(name)
            if bright:
                self.tex_modder.brighten(name, np.random.randint(0, 150))

    def _rand_camera(self):
        """Randomize pos, orientation, and fov of camera
        real camera pos is -1.75, 0, 1.62
        FOVY:
        Kinect2 is 53.8
        ASUS is 45
        https://www.asus.com/us/3D-Sensor/Xtion_PRO_LIVE/specifications/
        http://smeenk.com/kinect-field-of-view-comparison/
        """
        # Params
        FOVY_R = Range(40, 50)
        #X = Range(-3, -1)
        #Y = Range(-1, 3)
        #Z = Range(1, 2)
        #C_R3D = Range3D(X, Y, Z)
        #cam_pos = sample_xyz(C_R3D)
        #L_R3D = rto3d([-0.1, 0.1])

        C_R3D = Range3D([-0.07, 0.07], [-0.07, 0.07], [-0.07, 0.07])
        ANG3 = Range3D([-3, 3], [-3, 3], [-3, 3])

        # Look approximately at the robot, but then randomize the orientation around that
        cam_choices = np.array([[-1.75, 0, 1.62], [-1.3, 1.7, 1.62],
                                [-1.75, 0, 2]])
        cam_pos = cam_choices[np.random.choice(len(cam_choices))]
        # cam_pos = get_real_cam_pos(FLAGS.real_data_path)
        target_id = self.model.body_name2id(FLAGS.look_at)

        cam_off = 0  #sample_xyz(L_R3D)
        target_off = 0  #sample_xyz(L_R3D)
        quat = look_at(cam_pos + cam_off,
                       self.sim.data.body_xpos[target_id] + target_off)
        quat = jitter_angle(quat, ANG3)
        #quat = jitter_quat(quat, 0.01)

        cam_pos += sample_xyz(C_R3D)

        self.cam_modder.set_quat('camera1', quat)
        self.cam_modder.set_pos('camera1', cam_pos)
        self.cam_modder.set_fovy('camera1', 60)  # hard code to wide fovy

    def _next_camera(self):
        """Randomize pos, orientation, and fov of camera
        real camera pos is -1.75, 0, 1.62
        FOVY:
        Kinect2 is 53.8
        ASUS is 45
        https://www.asus.com/us/3D-Sensor/Xtion_PRO_LIVE/specifications/
        http://smeenk.com/kinect-field-of-view-comparison/
        """
        # Params
        FOVY_R = Range(40, 50)
        #X = Range(-3, -1)
        #Y = Range(-1, 3)
        #Z = Range(1, 2)
        #C_R3D = Range3D(X, Y, Z)
        #cam_pos = sample_xyz(C_R3D)
        #L_R3D = rto3d([-0.1, 0.1])

        C_R3D = Range3D([-0.07, 0.07], [-0.07, 0.07], [-0.07, 0.07])
        ANG3 = Range3D([-3, 3], [-3, 3], [-3, 3])

        # Look approximately at the robot, but then randomize the orientation around that

        # linearly interpolate to the next camera every K steps
        K = 5
        goal_cam_pos = self._cam_choices[(self._cam_step // K) + 1]
        offset = goal_cam_pos - self._curr_cam_pos
        offset *= (self._cam_step % K) / K
        self._curr_cam_pos += offset
        cam_pos = self._curr_cam_pos
        self._cam_step += 1

        # cam_pos = cam_choices[np.random.choice(len(cam_choices))]
        # cam_pos = get_real_cam_pos(FLAGS.real_data_path)
        target_id = self.model.body_name2id(FLAGS.look_at)

        cam_off = 0  #sample_xyz(L_R3D)
        target_off = 0  #sample_xyz(L_R3D)
        quat = look_at(cam_pos + cam_off,
                       self.sim.data.body_xpos[target_id] + target_off)
        quat = jitter_angle(quat, ANG3)
        #quat = jitter_quat(quat, 0.01)

        cam_pos += sample_xyz(C_R3D)

        self.cam_modder.set_quat('camera1', quat)
        self.cam_modder.set_pos('camera1', cam_pos)
        self.cam_modder.set_fovy('camera1', 60)  # hard code to wide fovy

    def _rand_lights(self):
        """Randomize pos, direction, and lights"""
        # light stuff
        #X = Range(-1.5, 1.5)
        #Y = Range(-1.2, 1.2)
        #Z = Range(0, 2.8)
        X = Range(-1.5, -0.5)
        Y = Range(-0.6, 0.6)
        Z = Range(1.0, 1.5)
        LIGHT_R3D = Range3D(X, Y, Z)
        LIGHT_UNIF = Range3D(Range(0, 1), Range(0, 1), Range(0, 1))

        # TODO: also try not altering the light dirs and just keeping them at like -1, or [0, -0.15, -1.0]
        for i, name in enumerate(self.model.light_names):
            lid = self.model.light_name2id(name)
            # random sample 80% of any given light being on
            if lid != 0:
                self.light_modder.set_active(name, sample([0, 1]) < 0.8)
                self.light_modder.set_dir(name, sample_light_dir())

            self.light_modder.set_pos(name, sample_xyz(LIGHT_R3D))

            #self.light_modder.set_dir(name, sample_xyz(rto3d([-1,1])))

            #self.light_modder.set_specular(name, sample_xyz(LIGHT_UNIF))
            #self.light_modder.set_diffuse(name, sample_xyz(LIGHT_UNIF))
            #self.light_modder.set_ambient(name, sample_xyz(LIGHT_UNIF))

            spec = np.array([sample(Range(0.5, 1))] * 3)
            diffuse = np.array([sample(Range(0.5, 1))] * 3)
            ambient = np.array([sample(Range(0.5, 1))] * 3)

            self.light_modder.set_specular(name, spec)
            self.light_modder.set_diffuse(name, diffuse)
            self.light_modder.set_ambient(name, ambient)
            #self.model.light_directional[lid] = sample([0,1]) < 0.2
            self.model.light_castshadow[lid] = sample([0, 1]) < 0.5

    def _rand_robot(self):
        """Randomize joint angles and jitter orientation"""
        jnt_shape = self.sim.data.qpos.shape
        self.sim.data.qpos[:] = sample_joints(self.model.jnt_range, jnt_shape)

        robot_gid = self.model.geom_name2id('robot_table_link')
        self.model.geom_quat[robot_gid] = jitter_quat(
            self.START_GEOM_QUAT[robot_gid], 0.01)

    def _rand_object(self):
        obj_gid = self.sim.model.geom_name2id('object')
        obj_bid = self.sim.model.geom_name2id('object')
        table_gid = self.model.geom_name2id('object_table')
        table_bid = self.model.body_name2id('object_table')

        obj_pose = self.start_obj_pose.copy()
        xval = self.model.geom_size[table_gid][
            0]  #- self.model.geom_size[obj_gid][0]
        yval = self.model.geom_size[table_gid][
            1]  #- self.model.geom_size[obj_gid][1]

        O_X = Range(-xval, xval)
        O_Y = Range(-yval, yval)
        O_Z = Range(0, 0)
        O_R3D = Range3D(O_X, O_Y, O_Z)

        newpos = obj_pose[:3] + sample_xyz(O_R3D)
        newquat = jitter_quat(obj_pose[3:], 0.1)
        obj_pose[:3] = newpos
        obj_pose[3:] = newquat
        self.sim.data.set_joint_qpos('object:joint', obj_pose)

        #T_X = Range(-0.1, 0.1)
        #T_Y = Range(-0.1, 0.1)
        #T_Z = Range(-0.1, 0.1)
        #T_R3D = Range3D(T_X, T_Y, T_Z)
        #self.model.body_pos[table_bid] = self.START_BODY_POS[table_bid] + sample_xyz(T_R3D)
        ## randomize orientation a wee bit
        #self.model.geom_quat[table_gid] = jitter_quat(self.START_GEOM_QUAT[table_gid], 0.01)

    def _rand_walls(self):
        wall_bids = {
            name: self.model.body_name2id(name)
            for name in ['wall_' + dir for dir in 'nesw']
        }
        window_gid = self.model.geom_name2id('west_window')
        #floor_gid = self.model.geom_name2id('floor')

        WA_X = Range(-0.2, 0.2)
        WA_Y = Range(-0.2, 0.2)
        WA_Z = Range(-0.1, 0.1)
        WA_R3D = Range3D(WA_X, WA_Y, WA_Z)

        WI_X = Range(-0.1, 0.1)
        WI_Y = Range(0, 0)
        WI_Z = Range(-0.5, 0.5)
        WI_R3D = Range3D(WI_X, WI_Y, WI_Z)

        R = Range(0, 0)
        P = Range(-10, 10)
        Y = Range(0, 0)
        RPY_R = Range3D(R, P, Y)

        #self.model.geom_quat[floor_gid] = jitter_quat(self.START_GEOM_QUAT[floor_gid], 0.01)
        #self.model.geom_pos[floor_gid] = self.START_GEOM_POS[floor_gid] + [0,0,sample(-0.1,0.1)

        self.model.geom_quat[window_gid] = sample_quat(RPY_R)
        #self.model.geom_quat[window_gid] = jitter_quat(self.START_GEOM_QUAT[window_gid], 0.01)
        self.model.geom_pos[
            window_gid] = self.START_GEOM_POS[window_gid] + sample_xyz(WI_R3D)

        for name in wall_bids:
            gid = wall_bids[name]
            self.model.body_quat[gid] = jitter_quat(self.START_BODY_QUAT[gid],
                                                    0.01)
            self.model.body_pos[gid] = self.START_BODY_POS[gid] + sample_xyz(
                WA_R3D)

    def _rand_distract(self):
        PREFIX = 'distract'
        geom_names = [
            name for name in self.model.geom_names if name.startswith(PREFIX)
        ]

        # Size range
        SX = Range(0.01, 0.5)
        SY = Range(0.01, 0.9)
        SZ = Range(0.01, 0.5)
        S3D = Range3D(SX, SY, SZ)
        # Back range
        B_PX = Range(-0.5, 2)
        B_PY = Range(-1.5, 2)
        B_PZ = Range(0, 3)
        B_P3D = Range3D(B_PX, B_PY, B_PZ)
        # Front range
        F_PX = Range(-2, -0.5)
        F_PY = Range(-2, 1)
        F_PZ = Range(0, 0.5)
        F_P3D = Range3D(F_PX, F_PY, F_PZ)

        for name in geom_names:
            gid = self.model.geom_name2id(name)
            range = B_P3D if np.random.binomial(1, 0.5) else F_P3D

            self.model.geom_pos[gid] = sample_xyz(range)
            self.model.geom_quat[gid] = random_quat()
            self.model.geom_size[gid] = sample_xyz(S3D, mode='logspace')
            self.model.geom_type[gid] = sample_geom_type()
            self.model.geom_rgba[gid][-1] = np.random.binomial(1, 0.5)

    def _set_visible(self, prefix, range_top, visible):
        """Helper function to set visibility of several objects"""
        if not visible:
            if range_top == 0:
                name = prefix
                gid = self.model.geom_name2id(name)
                self.model.geom_rgba[gid][-1] = 0.0

            for i in range(range_top):
                name = "{}{}".format(prefix, i)
                gid = self.model.geom_name2id(name)
                self.model.geom_rgba[gid][-1] = 0.0
        else:
            if range_top == 0:
                name = prefix
                gid = self.model.geom_name2id(name)
                self.model.geom_rgba[gid][-1] = 1.0

            for i in range(range_top):
                name = "{}{}".format(prefix, i)
                gid = self.model.geom_name2id(name)
                self.model.geom_rgba[gid][-1] = 1.0
Esempio n. 2
0
class ArenaModder(BaseModder):
    """
    Object to handle randomization of all relevant properties of Mujoco sim

    """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Get start state of params to slightly jitter later
        self.start_geo_size = self.model.geom_size.copy()
        self.start_geom_quat = self.model.geom_quat.copy()
        self.start_body_pos = self.model.body_pos.copy()
        self.start_body_quat = self.model.body_quat.copy()
        self.start_matid = self.model.geom_matid.copy()
        self.floor_offset = self.model.body_pos[self.model.body_name2id(
            'floor')]
        self.rock_mod_cache = None

        self.tex_modder = TextureModder(self.sim)
        self.cam_modder = CameraModder(self.sim)
        self.light_modder = LightModder(self.sim)

        # Set these both externally to visualize
        self.visualize = False
        self.viewer = None

    def whiten_materials(self):
        self.tex_modder.whiten_materials()

    def randomize(self):
        self.mod_textures()
        self.mod_lights()
        self.mod_camera()
        self.mod_walls()
        self.mod_extras()
        self.mod_rocks()

    def mod_textures(self):
        """Randomize all the textures in the scene, including the skybox"""
        for name in self.sim.model.geom_names:
            if name != "billboard" and "light_disc" not in name:
                self.tex_modder.rand_all(name)
                ##texture_sizes = [32, 64, 128, 256, 512, 1024]
                ##gid = self.model.geom_name2id(name)
                ##mid = self.model.geom_matid[gid]
                ##tid = self.model.mat_texid[mid]
                ##self.model.tex_width[tid] = sample_from_list(texture_sizes)
                ##self.model.tex_height[tid] = 6 * self.model.tex_width[tid]
        self.tex_modder.rand_all('skybox')

    def mod_lights(self):
        """Randomize pos, direction, and lights"""
        # light stuff
        LIGHT_RX = Range(LEFTX, RIGHTX)
        LIGHT_RY = Range(BINY, DIGY)
        LIGHT_RZ = Range(AFZ, AFZ + ZHIGH)
        LIGHT_RANGE3D = Range3D(LIGHT_RX, LIGHT_RY, LIGHT_RZ)
        LIGHT_UNIF = Range3D(Range(0, 1), Range(0, 1), Range(0, 1))

        for i, name in enumerate(self.model.light_names):
            lid = self.model.light_name2id(name)
            # random sample 80% of any given light being on
            self.light_modder.set_active(name, sample([0, 1]) < 0.8)
            #self.light_modder.set_active(name, 0)
            dir_xyz = sample_light_dir()
            self.light_modder.set_pos(name, sample_xyz(LIGHT_RANGE3D))
            self.light_modder.set_dir(name, dir_xyz)
            self.light_modder.set_specular(name, sample_xyz(LIGHT_UNIF))
            #self.light_modder.set_diffuse(name, sample_xyz(LIGHT_UNIF))
            #self.light_modder.set_ambient(name, sample_xyz(LIGHT_UNIF))
            #self.model.light_directional[lid] = sample([0,1]) < 0.01

    def mod_camera(self):
        """Randomize pos, direction, and fov of camera"""
        # Params
        XOFF = 1.0
        CAM_RX = Range(ACX - XOFF, ACX + XOFF)  # center of arena +/- 0.5
        CAM_RY = Range(BINY + 0.2, SZ_ENDY)
        CAM_RZ = Range(AFZ + ZLOW, AFZ + ZHIGH)
        CAM_RANGE3D = Range3D(CAM_RX, CAM_RY, CAM_RZ)
        CAM_RYAW = Range(-95, -85)
        CAM_RPITCH = Range(65, 90)
        CAM_RROLL = Range(85, 95)  # this might actually be pitch?
        CAM_ANGLE3 = Range3D(CAM_RYAW, CAM_RPITCH, CAM_RROLL)

        # "The horizontal field of view is computed automatically given the
        # window size and the vertical field of view." - Mujoco
        # This range was calculated using: themetalmuncher.github.io/fov-calc/
        # ZED has 110° hfov --> 78° vfov, Logitech C920 has 78° hfov ---> 49° vfov
        # These were rounded all the way down to 40° and up to 80°. It starts
        # to look a bit bad in the upper range, but I think it will help
        # generalization.
        CAM_RFOVY = Range(40, 80)

        # Actual mods
        self.cam_modder.set_pos('camera1', sample_xyz(CAM_RANGE3D))
        self.cam_modder.set_quat('camera1', sample_quat(CAM_ANGLE3))

        fovy = sample(CAM_RFOVY)
        self.cam_modder.set_fovy('camera1', fovy)

    def _set_visible(self, prefix, range_top, visible):
        """Helper function to set visibility of several objects"""
        if not visible:
            if range_top == 0:
                name = prefix
                gid = self.model.geom_name2id(name)
                self.model.geom_rgba[gid][-1] = 0.0

            for i in range(range_top):
                name = "{}{}".format(prefix, i)
                gid = self.model.geom_name2id(name)
                self.model.geom_rgba[gid][-1] = 0.0
        else:
            if range_top == 0:
                name = prefix
                gid = self.model.geom_name2id(name)
                self.model.geom_rgba[gid][-1] = 1.0

            for i in range(range_top):
                name = "{}{}".format(prefix, i)
                gid = self.model.geom_name2id(name)
                self.model.geom_rgba[gid][-1] = 1.0

    def mod_extra_distractors(self, visible=True):
        """mod rocks and tools on the side of the arena"""
        # TODO: I might consider changing these to look like rocks instead of
        # just random shapes.  It just looks weird to me right now.  Ok for now,
        # but it seems a bit off.
        N = 20
        self._set_visible("side_obj", N, visible)
        if not visible:
            return

        Z_JITTER = 0.05
        OBJ_XRANGE = Range(0.01, 0.1)
        OBJ_YRANGE = Range(0.01, 0.1)
        OBJ_ZRANGE = Range(0.01, 0.1)
        OBJ_SIZE_RANGE = Range3D(OBJ_XRANGE, OBJ_YRANGE, OBJ_ZRANGE)

        floor_gid = self.model.geom_name2id("floor")

        left_body_id = self.model.body_name2id("left_wall")
        left_geom_id = self.model.geom_name2id("left_wall")
        right_body_id = self.model.body_name2id("right_wall")
        right_geom_id = self.model.geom_name2id("right_wall")

        left_center = self.model.body_pos[left_body_id]
        left_geo = self.model.geom_size[left_geom_id]
        left_height = left_center[2] + left_geo[2]
        left_xrange = Range(left_center[0] - left_geo[0],
                            left_center[0] + left_geo[0])
        left_yrange = Range(left_center[1] - left_geo[1],
                            left_center[1] + left_geo[1])
        left_zrange = 0.02 + Range(left_height - Z_JITTER,
                                   left_height + Z_JITTER)
        left_range = Range3D(left_xrange, left_yrange, left_zrange)

        right_center = self.model.body_pos[right_body_id]
        right_geo = self.model.geom_size[right_geom_id]
        right_height = right_center[2] + right_geo[2]
        right_xrange = Range(right_center[0] - right_geo[0],
                             right_center[0] + right_geo[0])
        right_yrange = Range(right_center[1] - right_geo[1],
                             right_center[1] + right_geo[1])
        right_zrange = 0.02 + Range(right_height - Z_JITTER,
                                    right_height + Z_JITTER)
        right_range = Range3D(right_xrange, right_yrange, right_zrange)

        for i in range(N):
            name = "side_obj{}".format(i)
            obj_bid = self.model.body_name2id(name)
            obj_gid = self.model.geom_name2id(name)
            self.model.geom_quat[obj_gid] = random_quat()
            self.model.geom_size[obj_gid] = sample_xyz(OBJ_SIZE_RANGE)
            self.model.geom_type[obj_gid] = sample_geom_type()

            # 50% chance of invisible
            if sample([0, 1]) < 0.5:
                self.model.geom_rgba[obj_gid][-1] = 0.0
            else:
                self.model.geom_rgba[obj_gid][-1] = 1.0
            ## 50% chance of same color as floor and rocks
            if sample([0, 1]) < 0.5:
                self.model.geom_matid[obj_gid] = self.model.geom_matid[
                    floor_gid]
            else:
                self.model.geom_matid[obj_gid] = self.start_matid[obj_gid]

            # 10 always on the left, 10 always on the right
            if i < 10:
                self.model.body_pos[obj_bid] = sample_xyz(left_range)
            else:
                self.model.body_pos[obj_bid] = sample_xyz(right_range)

    def mod_extra_judges(self, visible=True):
        """mod NASA judges around the perimeter of the arena"""
        # TODO: might want to add regions on the sides of the arena, but these
        # may be covered by the distractors already
        N = 5
        self._set_visible("judge", N, visible)
        if not visible:
            return

        JUDGE_XRANGE = Range(0.1, 0.2)
        JUDGE_YRANGE = Range(0.1, 0.2)
        JUDGE_ZRANGE = Range(0.75, 1.0)
        JUDGE_SIZE_RANGE = Range3D(JUDGE_XRANGE, JUDGE_YRANGE, JUDGE_ZRANGE)

        digwall_bid = self.model.body_name2id("dig_wall")
        digwall_gid = self.model.geom_name2id("dig_wall")

        digwall_center = self.model.body_pos[digwall_bid]
        digwall_geo = self.model.geom_size[digwall_gid]
        digwall_xrange = Range(-1.0 + digwall_center[0] - digwall_geo[0],
                               1.0 + digwall_center[0] + digwall_geo[0])
        digwall_yrange = Range(digwall_center[1] + 0.5,
                               digwall_center[1] + 1.5)
        digwall_zrange = JUDGE_ZRANGE - 0.75
        digwall_range = Range3D(digwall_xrange, digwall_yrange, digwall_zrange)

        for i in range(N):
            name = "judge{}".format(i)
            judge_bid = self.model.body_name2id(name)
            judge_gid = self.model.geom_name2id(name)

            #self.model.geom_quat[judge_gid] = jitter_quat(self.start_geom_quat[judge_gid], 0.05)
            self.model.geom_quat[judge_gid] = random_quat()
            self.model.geom_size[judge_gid] = sample_xyz(JUDGE_SIZE_RANGE)
            self.model.geom_type[judge_gid] = sample_geom_type()
            if self.model.geom_type[judge_gid] == 3 or self.model.geom_type[
                    judge_gid] == 5:
                self.model.geom_size[judge_gid][1] = self.model.geom_size[
                    judge_gid][2]

            self.model.body_pos[judge_bid] = sample_xyz(digwall_range)

            # 50% chance of invisible
            self.model.geom_rgba[judge_gid][-1] = sample([0, 1]) < 0.5

    def mod_extra_robot_parts(self, visible=True):
        """add distractor parts of robots in the lower area of the camera frame"""
        N = 3
        self._set_visible("robot_part", N, visible)
        if not visible:
            return
        # Project difference into camera coordinate frame
        cam_pos = self.model.cam_pos[0]
        cam_quat = np.quaternion(*self.model.cam_quat[0])
        lower_range = Range3D([0.0, 0.0], [-0.2, -0.3], [-0.2, -0.3])
        lower_size = Range3D([0.2, 0.6], [0.01, 0.15], [0.01, 0.15])
        lower_angle = Range3D([-85.0, -95.0], [-180, 180], [-85, -95])

        upper_range = Range3D([-0.6, 0.6], [-0.05, 0.05], [-0.05, 0.05])
        upper_size = Range3D([0.005, 0.05], [0.005, 0.05], [0.01, 0.3])
        upper_angle = Range3D([-85.0, -95.0], [-180, 180], [-85, -95])

        name = "robot_part0"
        lower_bid = self.model.body_name2id(name)
        lower_gid = self.model.geom_name2id(name)

        lower_pos = cam_pos + quaternion.rotate_vectors(
            cam_quat, sample_xyz(lower_range))
        self.model.body_pos[lower_bid] = lower_pos
        self.model.geom_size[lower_gid] = sample_xyz(lower_size)
        self.model.geom_quat[lower_gid] = sample_quat(lower_angle)
        self.model.geom_type[lower_gid] = sample_geom_type(reject=["capsule"])

        if self.model.geom_type[lower_gid] == 5:
            self.model.geom_size[lower_gid][0] = self.model.geom_size[
                lower_gid][2]

        for i in range(1, 10):
            name = "robot_part{}".format(i)
            upper_bid = self.model.body_name2id(name)
            upper_gid = self.model.geom_name2id(name)
            upper_pos = lower_pos + sample_xyz(upper_range)
            self.model.body_pos[upper_bid] = upper_pos
            self.model.geom_size[upper_gid] = sample_xyz(upper_size)
            self.model.geom_type[upper_gid] = sample_geom_type()

            # 50% of the time, choose random angle instead reasonable angle
            if sample([0, 1]) < 0.5:
                self.model.geom_quat[upper_gid] = sample_quat(upper_angle)
            else:
                self.model.geom_quat[upper_gid] = random_quat()

    def mod_extra_arena_structure(self, visible=True):
        """add randomized structure of the arena in the background with 
        pillars and a crossbar"""
        N = 16
        self._set_visible("arena_structure", N, visible)
        if not visible:
            return

        STARTY = DIGY + 3.0
        ENDY = DIGY + 5.0
        XBAR_SIZE = Range3D([10.0, 20.0], [0.05, 0.1], [0.2, 0.5])
        XBAR_RANGE = Range3D([0.0, 0.0], [STARTY, ENDY], [3.0, 5.0])
        STRUCTURE_SIZE = Range3D([0.05, 0.1], [0.05, 0.1], [3.0, 6.0])
        STRUCTURE_RANGE = Range3D([np.nan, np.nan], [STARTY, ENDY], [0.0, 0.0])
        x_range = np.linspace(ACX - 10.0, ACX + 10.0, 15)

        # crossbar
        name = "arena_structure{}".format(N - 1)
        xbar_bid = self.model.body_name2id(name)
        xbar_gid = self.model.geom_name2id(name)

        self.model.geom_quat[xbar_gid] = jitter_quat(
            self.start_geom_quat[xbar_gid], 0.01)
        self.model.geom_size[xbar_gid] = sample_xyz(XBAR_SIZE)
        self.model.body_pos[xbar_bid] = sample_xyz(XBAR_RANGE)

        ### 10% chance of invisible
        ##if sample([0,1]) < 0.1:
        ##    self.model.geom_rgba[xbar_gid][-1] = 0.0
        ##else:
        ##    self.model.geom_rgba[xbar_gid][-1] = 1.0

        for i in range(N - 1):
            name = "arena_structure{}".format(i)
            arena_structure_bid = self.model.body_name2id(name)
            arena_structure_gid = self.model.geom_name2id(name)
            STRUCTURE_RANGE[0] = Range(x_range[i] - 0.1, x_range[i] + 0.1)

            self.model.geom_quat[arena_structure_gid] = jitter_quat(
                self.start_geom_quat[arena_structure_gid], 0.01)
            self.model.geom_size[arena_structure_gid] = sample_xyz(
                STRUCTURE_SIZE)
            self.model.body_pos[arena_structure_bid] = sample_xyz(
                STRUCTURE_RANGE)

            self.model.geom_matid[arena_structure_gid] = self.model.geom_matid[
                xbar_gid]

            # 10% chance of invisible
            if sample([0, 1]) < 0.1:
                self.model.geom_rgba[arena_structure_gid][-1] = 0.0
            else:
                self.model.geom_rgba[arena_structure_gid][-1] = 1.0

    def mod_extra_arena_background(self, visible=True, N=10):
        """ """
        self._set_visible("background", N, visible)
        if not visible:
            return

        BACKGROUND_XRANGE = Range(0.1, 0.5)
        BACKGROUND_YRANGE = Range(0.1, 0.5)
        BACKGROUND_ZRANGE = Range(0.1, 1.0)
        BACKGROUND_SIZE_RANGE = Range3D(BACKGROUND_XRANGE, BACKGROUND_YRANGE,
                                        BACKGROUND_ZRANGE)

        digwall_bid = self.model.body_name2id("dig_wall")
        digwall_gid = self.model.geom_name2id("dig_wall")
        c = digwall_center = self.model.body_pos[digwall_bid]

        background_range = Range3D([c[0] - 4.0, c[0] + 4.0],
                                   [c[1] + 8.0, c[1] + 12.0], [0.0, 5.0])

        for i in range(N):
            name = "background{}".format(i)
            background_bid = self.model.body_name2id(name)
            background_gid = self.model.geom_name2id(name)

            self.model.geom_quat[background_gid] = random_quat()
            self.model.geom_size[background_gid] = sample_xyz(
                BACKGROUND_SIZE_RANGE)
            self.model.geom_type[background_gid] = sample_geom_type()

            self.model.body_pos[background_bid] = sample_xyz(background_range)

    def mod_extra_lights(self, visible=True):
        """"""
        #visible = False
        MODE_TARGETBODY = 3
        STARTY = DIGY + 3.0
        ENDY = DIGY + 5.0
        LIGHT_RANGE = Range3D([ACX - 2.0, ACX + 2.0], [STARTY, ENDY],
                              [3.0, 5.0])
        any_washed = False

        for i in range(3):
            name = "extra_light{}".format(i)
            lid = self.model.light_name2id(name)
            self.model.light_active[lid] = visible
            if not visible:
                return

            self.model.light_pos[lid] = sample_xyz(LIGHT_RANGE)
            self.model.light_dir[lid] = sample_light_dir()
            self.model.light_directional[lid] = 0

            ### 3% chance of directional light, washing out colors
            ### (P(at least 1 will be triggered) ~= 0.1) = 1 - 3root(p)
            washout = sample([0, 1]) < 0.03
            any_washed = any_washed or washout
            self.model.light_directional[lid] = washout

        if any_washed:
            self.mod_extra_light_discs(
                visible=visible and sample([0, 1]) < 0.9)
        else:
            self.mod_extra_light_discs(visible=False)

    def mod_extra_light_discs(self, visible=True):
        """"""
        N = 10
        self._set_visible("light_disc", N, visible)
        if not visible:
            return

        Z_JITTER = 0.05
        DISC_XRANGE = Range(0.1, 4.0)
        DISC_YRANGE = Range(0.1, 4.0)
        DISC_ZRANGE = Range(0.1, 4.0)
        DISC_SIZE_RANGE = Range3D(DISC_XRANGE, DISC_YRANGE, DISC_ZRANGE)
        OUTR = 20.0
        INR = 10.0

        floor_bid = self.model.body_name2id("floor")
        c = self.model.body_pos[floor_bid]
        disc_xrange = Range(c[0] - OUTR, c[0] + OUTR)
        disc_yrange = Range(c[1], c[1] + OUTR)
        disc_zrange = Range(-5.0, 10.0)
        disc_range = Range3D(disc_xrange, disc_yrange, disc_zrange)

        for i in range(N):
            name = "light_disc{}".format(i)
            disc_bid = self.model.body_name2id(name)
            disc_gid = self.model.geom_name2id(name)
            disc_mid = self.model.geom_matid[disc_gid]

            self.model.geom_quat[disc_gid] = random_quat()
            self.model.geom_size[disc_gid] = sample_xyz(DISC_SIZE_RANGE)
            self.model.geom_type[disc_gid] = sample_geom_type()

            # keep trying to place the disc until it lands outside of blocking stuff
            # (they will land in a u shape, because they are sampled with a
            # constrain to not land in the middle)
            while True:
                xyz = sample_xyz(disc_range)

                if ((xyz[0] > (c[0] - INR) and xyz[0] < (c[0] + INR))
                        and (xyz[1] > (c[1] - INR) and xyz[1] < (c[1] + INR))):
                    continue
                else:
                    self.model.geom_pos[disc_gid] = xyz
                    break

            # 50% chance of invisible
            if sample([0, 1]) < 0.5:
                self.model.geom_rgba[disc_gid][-1] = 0.0
            else:
                self.model.geom_rgba[disc_gid][-1] = 1.0

    def mod_extras(self, visible=True):
        """
        Randomize extra properties of the world such as the extra rocks on the side
        of the arena wal

        The motivation for these mods are that it seems likely that these distractor 
        objects, judges, etc, in the real images could degrade the performance 
        of the detector. 

        Artifacts:
        - Rocks and tools on edges of bin
        - NASA judges around the perimeter
        - Parts of the robot in the frame
        - Background arena structure and crowd 
        - Bright light around edges of arena
        """
        self.mod_extra_distractors(visible)
        self.mod_extra_robot_parts(visible)
        self.mod_extra_lights(visible)
        # 10% of the time, hide the other distractor pieces
        self.mod_extra_judges(visible=visible and sample([0, 1]) < 0.9)
        self.mod_extra_arena_structure(
            visible=visible and sample([0, 1]) < 0.9)
        self.mod_extra_arena_background(
            visible=visible and sample([0, 1]) < 0.9)

    def mod_walls(self):
        """
        Randomize the x, y, and orientation of the walls slights.
        Also drastically randomize the height of the walls. In many cases they won't
        be seen at all. This will allow the model to generalize to scenarios without
        walls, or where the walls and geometry is slightly different than the sim 
        model
        """

        wall_names = ["left_wall", "right_wall", "bin_wall", "dig_wall"]
        for name in wall_names:
            geom_id = self.model.geom_name2id(name)
            body_id = self.model.body_name2id(name)

            jitter_x = Range(-0.2, 0.2)
            jitter_y = Range(-0.2, 0.2)
            jitter_z = Range(-0.75, 0.0)
            jitter3D = Range3D(jitter_x, jitter_y, jitter_z)

            self.model.body_pos[
                body_id] = self.start_body_pos[body_id] + sample_xyz(jitter3D)
            self.model.body_quat[body_id] = jitter_quat(
                self.start_body_quat[body_id], 0.005)
            if sample([0, 1]) < 0.05:
                self.model.body_pos[body_id][2] = -2.0

    def mod_dirt(self):
        """Randomize position and rotation of dirt"""
        # TODO: 50% chance to flip the dirt pile upsidedown, then we will have
        # to account for this in calculations

        # dirt stuff
        DIRT_RX = Range(0.0, 0.3)
        DIRT_RY = Range(0.0, 0.3)
        DIRT_RZ = Range(-0.05, 0.03)
        DIRT_RANGE3D = Range3D(DIRT_RX, DIRT_RY, DIRT_RZ)
        DIRT_RYAW = Range(-180, 180)
        DIRT_RPITCH = Range(-90.5, -89.5)
        DIRT_RROLL = Range(-0.5, 0.5)
        DIRT_ANGLE3 = Range3D(DIRT_RYAW, DIRT_RPITCH, DIRT_RROLL)
        dirt_bid = self.model.body_name2id("dirt")
        dirt_gid = self.model.geom_name2id("dirt")
        dirt_mid = self.model.geom_dataid[dirt_gid]

        # randomize position and yaw of dirt
        self.model.body_pos[dirt_bid] = self.start_body_pos[
            dirt_bid] + sample_xyz(DIRT_RANGE3D)
        self.model.geom_quat[dirt_gid] = sample_quat(DIRT_ANGLE3)

        vert_adr = self.model.mesh_vertadr[dirt_mid]
        vert_num = self.model.mesh_vertnum[dirt_mid]
        mesh_verts = self.model.mesh_vert[vert_adr:vert_adr + vert_num]

        rot_quat = self.model.geom_quat[dirt_gid]
        rots = quaternion.rotate_vectors(
            np.quaternion(*rot_quat).normalized(), mesh_verts)

        mesh_abs_pos = self.floor_offset + self.model.body_pos[dirt_bid] + rots

        #xy_indexes = mesh_abs_pos[:, 0:2]
        #z_heights = mesh_abs_pos[:, 2]

        # Return a function that the user can call to get the approximate
        # height of an xy location
        def local_mean_height(xy):
            """
            Take an xy coordinate, and the approximate z height of the mesh at that
            location. It works decently.

            Uses a weighted average mean of all points within a threshold of xy.
            """
            # grab all mesh points within threshold euclidean distance of xy
            gt0_xyz = mesh_abs_pos[mesh_abs_pos[:, 2] > 0.01]
            eudists = np.sum(np.square(gt0_xyz[:, 0:2] - xy), axis=1)
            indices = eudists < 0.1
            close_xyz = gt0_xyz[indices]

            # if there are any nearby points above 0
            if np.count_nonzero(close_xyz[:, 2]) > 0:
                # weights for weighted sum. closer points to xy have higher weight
                weights = 1 / (eudists[indices])
                weights = np.expand_dims(weights / np.sum(weights), axis=1)
                pos = np.sum(close_xyz * weights, axis=0)

                # show an "o" and a marker where the height is
                if self.visualize:
                    self.viewer.add_marker(pos=pos,
                                           label="o",
                                           size=np.array([0.01, 0.01, 0.01]),
                                           rgba=np.array([0.0, 1.0, 0.0, 1.0]))

                # approximate z height of ground
                return pos[2]
            else:
                return 0

        def always_zero(xy):
            return 0

        dirt_height_xy = local_mean_height
        #dirt_height_xy = always_zero
        return dirt_height_xy

    def mod_rocks(self):
        """
        Randomize the rocks so that the model will generalize to competition rocks
        This randomizations currently being done are:
            - Positions (within guesses of competition regions)
            - Orientations
            - Shuffling the 3 rock meshes so that they can be on the left, middle, or right
            - Generating new random rock meshes every n runs (with Blender)
        """
        # Rock placement range parameters
        ROCK_LANEX = 0.4  # width parameters of x range
        OUTER_EXTRA = 0.5  # how much farther rocks should go out on the right and left lanes
        ROCK_BUFFX = 0.2  # distacne between rock lanes
        # How far into the obstacle zone the rocks should start.
        ROCK_START_OFFSET = 0.2
        MID_START_OFFSET = 0.4  # bit more for middle rock
        ROCK_RY = Range(OBS_SY + ROCK_START_OFFSET, OBS_ENDY)
        MID_RY = Range(OBS_SY + MID_START_OFFSET, OBS_ENDY)
        ROCK_RZ = Range(AFZ - 0.02, AFZ + 0.2)
        # Position dependent ranges
        LEFT_RX = Range(-3 * ROCK_LANEX - OUTER_EXTRA,
                        -ROCK_LANEX - ROCK_BUFFX)
        MID_RX = Range(-ROCK_LANEX, ROCK_LANEX)
        RIGHT_RX = Range(ROCK_BUFFX + ROCK_LANEX, 3 * ROCK_LANEX + OUTER_EXTRA)
        # Form full 3D sample range
        LEFT_ROCK_RANGE = Range3D(LEFT_RX, ROCK_RY, ROCK_RZ)
        MID_ROCK_RANGE = Range3D(MID_RX, MID_RY, ROCK_RZ)
        RIGHT_ROCK_RANGE = Range3D(RIGHT_RX, ROCK_RY, ROCK_RZ)
        ROCK_RANGES = [LEFT_ROCK_RANGE, MID_ROCK_RANGE, RIGHT_ROCK_RANGE]

        # actual mods
        rock_body_ids = {}
        rock_geom_ids = {}
        rock_mesh_ids = {}
        max_height_idxs = {}
        rot_cache = {}
        #max_height_xys = {}

        dirt_height_xy = self.mod_dirt()

        for name in self.model.geom_names:
            if name[:4] != "rock":
                continue

            geom_id = self.model.geom_name2id(name)
            body_id = self.model.body_name2id(name)
            mesh_id = self.model.geom_dataid[geom_id]
            rock_geom_ids[name] = geom_id
            rock_body_ids[name] = body_id
            rock_mesh_ids[name] = mesh_id

            # Rotate the rock and get the z value of the highest point in the
            # rotated rock mesh
            rot_quat = random_quat()
            vert_adr = self.model.mesh_vertadr[mesh_id]
            vert_num = self.model.mesh_vertnum[mesh_id]
            mesh_verts = self.model.mesh_vert[vert_adr:vert_adr + vert_num]
            rots = quaternion.rotate_vectors(
                np.quaternion(*rot_quat).normalized(), mesh_verts)
            self.model.geom_quat[geom_id] = rot_quat
            max_height_idx = np.argmax(rots[:, 2])
            max_height_idxs[name] = max_height_idx
            rot_cache[name] = rots

        # Shuffle the positions of the rocks (l or m or r)
        shuffle_names = list(rock_body_ids.keys())
        random.shuffle(shuffle_names)
        self.rock_mod_cache = []
        for i in range(len(shuffle_names)):
            name = shuffle_names[i]
            rots = rot_cache[name]
            self.model.body_pos[rock_body_ids[name]] = np.array(
                sample_xyz(ROCK_RANGES[i]))

            max_height_idx = max_height_idxs[name]
            xyz_for_max_z = rots[max_height_idx]

            # xyz coords in global frame
            global_xyz = self.floor_offset + xyz_for_max_z + self.model.body_pos[
                rock_body_ids[name]]
            gxy = global_xyz[0:2]
            max_height = global_xyz[2]

            if self.visualize:
                self.viewer.add_marker(pos=global_xyz,
                                       label="m",
                                       size=np.array([0.01, 0.01, 0.01]),
                                       rgba=np.array([0.0, 0.0, 1.0, 1.0]))

            dirt_z = dirt_height_xy(gxy)
            #dirt_z = 0
            #print(name, dirt_z)

            z_height = max_height - dirt_z
            self.rock_mod_cache.append((name, z_height))

    def get_ground_truth(self):
        """
        self.rock_mod_cache set from self.mod_rocks

        Return 1d numpy array of 9 elements for positions of all 3 rocks including:
            - rock x dist from cam
            - rock y dist from cam
            - rock z height from arena floor
        """
        cam_pos = self.model.cam_pos[0]

        #line_pos = self.floor_offset + np.array([0.0, 0.75, 0.0])
        #self.viewer.add_marker(pos=line_pos)

        ##r0_pos = self.floor_offset + self.model.body_pos[self.model.body_name2id('rock0')]
        ##r1_pos = self.floor_offset + self.model.body_pos[self.model.body_name2id('rock1')]
        ##r2_pos = self.floor_offset + self.model.body_pos[self.model.body_name2id('rock2')]
        ##r0_diff = r0_pos - cam_pos
        ##r1_diff = r1_pos - cam_pos
        ##r2_diff = r2_pos - cam_pos

        ground_truth = np.zeros(9, dtype=np.float32)
        for i, slot in enumerate(self.rock_mod_cache):
            name = slot[0]
            z_height = slot[1]

            pos = self.floor_offset + self.model.body_pos[
                self.model.body_name2id(name)]
            diff = pos - cam_pos

            # Project difference into camera coordinate frame
            cam_angle = quaternion.as_euler_angles(
                np.quaternion(*self.model.cam_quat[0]))[0]
            cam_angle += np.pi / 2
            in_cam_frame = np.zeros_like(diff)
            x = diff[1]
            y = -diff[0]
            in_cam_frame[0] = x * np.cos(cam_angle) + y * np.sin(cam_angle)
            in_cam_frame[1] = -x * np.sin(cam_angle) + y * np.cos(cam_angle)
            in_cam_frame[2] = z_height
            # simple check that change of frame is mathematically valid
            assert (np.isclose(np.sum(np.square(diff[:2])),
                               np.sum(np.square(in_cam_frame[:2]))))
            # swap positions to match ROS standard coordinates
            ground_truth[3 * i + 0] = in_cam_frame[0]
            ground_truth[3 * i + 1] = in_cam_frame[1]
            ground_truth[3 * i + 2] = in_cam_frame[2]
            text = "{0} x: {1:.2f} y: {2:.2f} height:{3:.2f}".format(
                name, ground_truth[3 * i + 0], ground_truth[3 * i + 1],
                z_height)
            ##text = "x: {0:.2f} y: {1:.2f} height:{2:.2f}".format(ground_truth[3*i+0], ground_truth[3*i+1], z_height)
            ##text = "height:{0:.2f}".format(z_height)
            if self.visualize:
                self.viewer.add_marker(pos=pos, label=text, rgba=np.zeros(4))

        return ground_truth
Esempio n. 3
0
class ObjectCollectionEnv(BaseEnv):
    """Custom muluti-object environment"""
    def __init__(self, FLAGS):
        super().__init__(FLAGS)

        if self.FLAGS['baxter']:
            self.robot_offset = lambda: self.model.body_pos[self.name2bid(
                'base_ground')]
        else:
            self.robot_offset = lambda: self.model.body_pos[self.name2bid(
                'robot_table_floor')]

        self.table_center = lambda: TABLES[self.FLAGS['default_table']][
            'pos'] + self.robot_offset(
            ) + [0, 0, TABLES[self.FLAGS['default_table']]['wood'][2]]

        self.LIGHT = {
            'ambient': self.model.light_ambient,
            'active': self.model.light_active,
            'castshadow': self.model.light_castshadow,
            'diffuse': self.model.light_diffuse,
            'dir': self.model.light_dir,
            'pos': self.model.light_pos,
            'specular': self.model.light_specular
        }
        self.START_BODY_POS = self.model.body_pos.copy()
        self.START_LIGHT = {}
        for key in self.LIGHT:
            self.START_LIGHT[key] = self.LIGHT[key].copy()

        self.tex_modder = TextureModder(self.sim)
        self.cam_modder = CameraModder(self.sim)
        self.light_modder = LightModder(self.sim)
        if self.FLAGS['domrand']:
            self.tex_modder.whiten_materials()

        self.goal = {}

        # mean camera position defined relative to table center
        if FLAGS['baxter']:
            self.CAMERA_POS_MEAN = np.array([-0.64, 0.025, 1.0])
        else:
            self.CAMERA_POS_MEAN = np.array([-0.65, 0.0, 0.68])

        # establish values for discretizing
        self.X = self.FLAGS['ACTS']['x']
        self.Y = self.FLAGS['ACTS']['y']
        self.YAW = self.FLAGS['ACTS']['yaw']
        self.DIST = self.FLAGS['ACTS']['dist']

        self.action_space = self.FLAGS['action_space']
        self.observation_space = self.FLAGS['observation_space']

        if self.FLAGS['goal_conditioned']:
            self.goal['array'] = np.zeros(self.FLAGS['state_shape'])
            if self.FLAGS['use_image']:
                self.goal['image'] = np.zeros(self.FLAGS['image_shape'])

        self.paddle_bid = self.name2bid('paddle')
        self.paddle_gid = self.name2gid('paddle')
        self.N = int(
            0.5 /
            self.dt)  # N should correspond to about 0.5 seconds of timestep

        self.consecutive_fences = 0
        self.escape_count = 0
        self._set_invisiwall()
        self._reset_sim()

        #self.viewer.add_marker(pos=cam_pos, label="CAM: {}".format(cam_pos))

    def _get_object_poses(self):
        objs = []
        for i in range(self.FLAGS['num_objects']):
            obj = 'object{}:joint'.format(i)
            pose = self.sim.data.get_joint_qpos(obj).copy()
            pose[:3] -= self.table_center()
            if not self.FLAGS['use_quat']:
                pose = pose[:3]
            objs.append(pose)
        objs = np.stack(objs)
        return objs

    def _get_obs(self):
        self.sim.step()

        array = self._get_object_poses()

        obs = {'array': array, 'single': np.zeros(4)}
        if self.FLAGS['use_image']:
            image = self._get_cam_frame().astype(np.float32)
            image = cv2.resize(image,
                               self.FLAGS['image_shape'][:2],
                               interpolation=cv2.INTER_AREA)
            image /= 255.0

            obs.update({'image': image})

        if self.FLAGS['goal_conditioned']:
            obs.update({'goal_array': self.goal['array']})
            if self.FLAGS['use_image']:
                obs.update({'goal_image': self.goal['image']})
        return obs

    def _geofence_escape(self):
        """check if any of the cubes have escaped the table.  this should not be possible
        but still happens due to some glitch with mujoco

        return True if any cube escapes
        
        """
        cubes = self._get_object_poses()

        cube_size = self.sim.model.geom_size[self.name2gid('object0')]

        table_range = 0.5 * TABLES[self.FLAGS['default_table']]['wood']
        table_range[
            2] = 0.1  # be forgiving on the height, since we already have object_fell to check for drops
        table_box = (1 + 0.01) * np.stack([-table_range, table_range], axis=-1)

        for c in cubes:
            cube_box = c[:3, None] + np.stack([-cube_size, cube_size], axis=-1)
            if not overlap3D(cube_box, table_box):
                return True

        return False

    def _object_fell(self):
        """check if any of the cubes have fallen off the table"""
        cubes = self._get_object_poses()
        for c in cubes:
            if c[2] < -0.01:
                return True
        return False

    def step(self, action):
        # convert from tuple to dictionary. this is pretty risky and bad practice, but should throw errors for shape mismatch in most cases *crosses fingers*
        if isinstance(action, tuple) or isinstance(action, np.ndarray):
            names = ['x', 'y', 'yaw']
            names += ['dist'] if self.FLAGS['use_dist'] else []
            sorted_names = list(sorted(names))
            action = {
                sorted_names[i]: action[i]
                for i in range(len(sorted_names))
            }
        action = copy.deepcopy(action)
        safe_action = self._set_action(action)

        reward = 0.0

        if self._geofence_escape() or self._object_fell():
            reward += -10.0
            done = True
            print('AHHH')
        else:
            done = False

        obs = self._get_obs()

        if self.FLAGS['penalize_invisiwall'] and safe_action != '':
            if 'wall' in safe_action:
                reward += -1.0

        info = {'seed': self.current_seed}
        out_obs = self.get_out_obs()

        return out_obs, reward, done, info

    def _set_action(self, action):
        """Apply the action to the env
        
        Return True is safe action, else false
        """
        self._randomize()
        contact_invisiwall = False
        if self.FLAGS['discrete']:
            x = self.X[action['x']]
            y = self.Y[action['y']]
            yaw = self.YAW[action['yaw']]
            if self.FLAGS['use_dist']:
                dist = self.DIST[1 + action['dist']]
            else:
                dist = self.FLAGS['max_dx']
        else:
            spaces = self.action_space.spaces

            def clip_act(key):
                """clip and then map back to semantic action scale"""
                clipped = np.clip(action[key], -1.0, 1.0)
                return unmap_continuous(key, clipped, self.FLAGS)

            x = clip_act('x')
            y = clip_act('y')
            yaw = clip_act('yaw')
            if self.FLAGS['use_dist']:
                dist = clip_act('dist')
            else:
                dist = self.FLAGS['max_dx']

        # Parse action
        s_xyz = np.array([x, y, 0.5])
        s_xyz += self.table_center()
        self.model.geom_pos[self.paddle_gid] = s_xyz
        self.sim.data.geom_xpos[self.paddle_gid] = s_xyz

        quat = quaternion.from_euler_angles(0, 0, yaw).normalized().components
        self.sim.model.geom_quat[self.paddle_gid] = quat

        xp_aid = self.model.actuator_name2id('xp')
        xv_aid = self.model.actuator_name2id('xv')
        yp_aid = self.model.actuator_name2id('yp')
        yv_aid = self.model.actuator_name2id('yv')
        zp_aid = self.model.actuator_name2id('zp')
        zv_aid = self.model.actuator_name2id('zv')

        self.sim.data.set_joint_qpos('paddle:slidex', np.array([0]))
        self.sim.data.set_joint_qpos('paddle:slidey', np.array([0]))
        self.sim.data.set_joint_qpos('paddle:slidez', np.array([0]))
        self.sim.data.set_joint_qpos('paddle:hingez', np.array([0]))

        self.sim.data.ctrl[:] = 0.0

        # Set goal positions of paddle (lower to targetz, then push to targetx and targety based on action)
        if self.FLAGS['baxter']:
            targetz = -0.925 + 0.83
        else:
            targetz = 0.69
        targetx, targety = s_xyz[0] + np.cos(yaw) * dist, s_xyz[1] + np.sin(
            yaw) * dist
        # reverse push at the end so that we are not in contact with block when we lift.
        btargetx, btargety = s_xyz[0] + np.cos(yaw) * (
            0.5 * dist), s_xyz[1] + np.sin(yaw) * (0.5 * dist)

        # Draw indicator of paddle target pos
        target = np.array([targetx, targety, targetz])
        site_id = self.name2sid('target0')
        self.sim.model.site_rgba[site_id] = np.zeros(4)
        if self.FLAGS['render']:
            self.sim.model.site_pos[site_id] = target
            self.sim.model.site_quat[site_id] = quat

        # Vertical PID controller to lower and raise
        P = 400.0
        D = 40.0
        v_pid = PID(
            P,
            D,
            curr_fn=lambda: self.sim.data.geom_xpos[self.paddle_gid][2],
            goal=targetz)

        # Horizontal (xy) PID controllers to push
        HP = 600.0
        HD = 60.0
        hx_pid = PID(
            HP,
            HD,
            curr_fn=lambda: self.sim.data.geom_xpos[self.paddle_gid][0],
            goal=targetx)
        hy_pid = PID(
            HP,
            HD,
            curr_fn=lambda: self.sim.data.geom_xpos[self.paddle_gid][1],
            goal=targety)

        # Lower the paddle
        for i in range(self.N):
            self.sim.data.ctrl[zp_aid], self.sim.data.ctrl[
                zv_aid] = v_pid.step()
            self.sim.step()
            contact_invisiwall = contact_invisiwall or self._collide_wall()
            if self.FLAGS['render']:
                self.viewer.render()

        # Push forward to the target position
        for i in range(int(1.5 * self.N)):
            #self.sim.data.ctrl[zp_aid], self.sim.data.ctrl[zv_aid]  = v_pid.step()
            self.sim.data.ctrl[xp_aid], self.sim.data.ctrl[
                xv_aid] = hx_pid.step()
            self.sim.data.ctrl[yp_aid], self.sim.data.ctrl[
                yv_aid] = hy_pid.step()
            contact_invisiwall = contact_invisiwall or self._collide_wall()
            self.sim.step()
            if self.FLAGS['render']:
                self.viewer.render()

        # Reverse a bit
        hx_pid.reset_goal(btargetx)
        hy_pid.reset_goal(btargety)
        for i in range(int(0.1 * self.N)):
            #self.sim.data.ctrl[zp_aid], self.sim.data.ctrl[zv_aid]  = v_pid.step()
            self.sim.data.ctrl[xp_aid], self.sim.data.ctrl[
                xv_aid] = hx_pid.step()
            self.sim.data.ctrl[yp_aid], self.sim.data.ctrl[
                yv_aid] = hy_pid.step()
            contact_invisiwall = contact_invisiwall or self._collide_wall()
            self.sim.step()
            if self.FLAGS['render']:
                self.viewer.render()

        # Raise the paddle back up after pushing
        v_pid.reset_goal(1.5)
        for i in range(self.N * 3):
            self.sim.data.ctrl[zp_aid], self.sim.data.ctrl[
                zv_aid] = v_pid.step()
            self.sim.data.ctrl[xp_aid], self.sim.data.ctrl[
                xv_aid] = hx_pid.step()
            self.sim.data.ctrl[yp_aid], self.sim.data.ctrl[
                yv_aid] = hy_pid.step()
            contact_invisiwall = contact_invisiwall or self._collide_wall()
            self.sim.step()
            if self.FLAGS['render']:
                self.viewer.render()

        if self.FLAGS['penalize_invisiwall'] and contact_invisiwall:
            safe = 'wall'
        else:
            safe = ''
        return safe

    def reset(self):
        self.goal = self._sample_goal()

        did_reset_sim = False
        while not did_reset_sim:
            did_reset_sim = self._reset_sim(
                reset_mode=self.np_random.choice(self.FLAGS['reset_mode']))

        curr = self._get_obs()['array']
        if self.FLAGS['render'] and self.FLAGS['goal_conditioned']:
            goal_array = self.goal['array'] * 0.5 * TABLES[
                self.FLAGS['default_table']]['wood'][:2]
            site_dests = np.c_[goal_array, curr[:, 2]] + self.table_center()
            for i in range(self.FLAGS['num_objects']):
                site = self.name2sid('object{}'.format(i))
                if self.FLAGS['render']:
                    self.model.site_pos[site] = site_dests[i] - [0.0, 0.0, 0.1]
                    self.model.site_rgba[site] = np.ones(4)
        return self.get_out_obs()

    def _collide_wall(self):
        wall_gids = [self.name2gid('invisiwall_' + str(i)) for i in range(4)]
        obj_gids = [
            self.name2gid('object' + str(i))
            for i in range(self.FLAGS['num_objects'])
        ]

        bad_contacts = [
            1 for cc in self.sim.data.contact
            if (cc.geom1 in wall_gids and cc.geom2 in obj_gids) or (
                cc.geom1 in obj_gids and cc.geom2 in wall_gids)
        ]
        return len(bad_contacts) > 0

    def _set_invisiwall(self):
        # NOTE: Generally changing the shape of objects does not do well for contacts
        X = self.FLAGS['ACTS']['x']
        Y = self.FLAGS['ACTS']['y']
        DEPTH, Z = self.model.geom_size[self.name2gid('invisiwall_0')][1:]
        Z -= (Z / 2)

        idx = 0
        for dir in [0, -1]:
            for xy in ['x', 'y']:
                gid = self.name2gid('invisiwall_' + str(idx))
                if self.FLAGS['view_invisiwall']:
                    pass
                else:
                    self.model.geom_rgba[gid] = np.zeros((4, ))
                if xy == 'x':
                    yval = 1.05 * np.sign(Y[dir]) * (np.abs(Y[dir]) + DEPTH)
                    self.model.geom_pos[gid] = self.table_center() + np.array(
                        [0, yval, Z])
                    #self.model.geom_size[gid] = np.array([X, DEPTH, Z])
                else:
                    xval = 1.05 * np.sign(X[dir]) * (np.abs(X[dir]) + DEPTH)
                    self.model.geom_pos[gid] = self.table_center() + np.array(
                        [xval, 0, Z])
                    #self.model.geom_size[gid] = np.array([DEPTH, Y, Z])

                idx += 1

        self.sim.forward()
        self.sim.step()

    def _sample_xyzs(self, reset_mode=None):
        reset_mode = reset_mode or 'cluster'  #self.FLAGS['reset_mode']
        object_xyzs = []  # for checking collisions against
        sample_range = R3D(self.FLAGS['obj_rangex'], self.FLAGS['obj_rangey'],
                           R(0, 0))

        if reset_mode == 'uniform':
            for i in range(self.FLAGS['num_objects']):
                while True:
                    object_xyz = sim_utils.sample_xyz(self.np_random,
                                                      sample_range)
                    collision = False
                    # TODO: check if this is necessary. could get speedup
                    for other in object_xyzs:
                        collision = collision or np.linalg.norm(
                            object_xyz[:2] - other[:2]) < 0.0127

                    if not collision and sim_utils.in_range(
                            object_xyz[0],
                            self.FLAGS['obj_rangex']) and sim_utils.in_range(
                                object_xyz[1], self.FLAGS['obj_rangey']):
                        object_xyzs.append(object_xyz)
                        break
        elif reset_mode == 'cluster':
            diameter = self.model.geom_size[self.name2gid('object0')][0]
            dirs = list(
                itertools.product([-diameter * 2, 0.0, diameter * 2],
                                  [-diameter * 2, 0.0, diameter * 2]))

            first_object_xyz = sim_utils.sample_xyz(self.np_random,
                                                    sample_range)
            object_xyzs.append(first_object_xyz)

            for i in range(1, self.FLAGS['num_objects']):
                while True:
                    reference = object_xyzs[self.np_random.randint(
                        len(object_xyzs))]
                    dir = np.array(dirs[self.np_random.randint(len(dirs))])
                    dir *= self.np_random.uniform(1.0, 2.0)
                    xyz = reference.copy()
                    xyz[:2] = reference[:2] + dir

                    collision = False
                    for other in object_xyzs:
                        collision = collision or np.linalg.norm(
                            xyz[:2] - other[:2]) < 0.0127

                    if not collision and sim_utils.in_range(
                            xyz[0],
                            self.FLAGS['obj_rangex']) and sim_utils.in_range(
                                xyz[1], self.FLAGS['obj_rangey']):
                        object_xyzs.append(xyz)
                        break

        return object_xyzs

    def _sample_object_poses(self, reset_mode=None):
        #rstate = np.random.RandomState(seed=0)  # use this to always reset to deterministic state
        xyz = self._sample_xyzs(reset_mode=reset_mode)
        for i in range(self.FLAGS['num_objects']):
            gid = self.name2gid('object{}'.format(i))
            #self.model.geom_size[gid] = self.FLAGS['cube_size']
            obj = 'object{}:joint'.format(i)
            object_qpos = self.sim.data.get_joint_qpos(obj)
            assert object_qpos.shape == (7, )
            object_xpos = self.table_center() + [0.0, 0.0, 0.2] + xyz[i]
            object_quat = sim_utils.sample_quat(
                self.np_random, R3D(R(0, 180), R(0, 0), R(0, 0)))
            object_qpos = np.concatenate([object_xpos, object_quat])
            self.sim.data.set_joint_qpos(obj, object_qpos)

    def _reset_sim(self, reset_mode=None):
        self.object_has_fallen = False
        self.object_fallen_stuff = None
        self.sim.set_state(self.initial_state)
        if self.FLAGS['baxter']:
            set_table(self.model, 'object_table', 'folding', self.FLAGS)
        else:
            set_table(self.model, 'robot_table', 'robot', self.FLAGS)
            set_table(self.model, 'object_table', 'small', self.FLAGS)

        self.model.site_size[self.name2sid('target0')] = self.model.geom_size[
            self.paddle_gid]
        self.cam_pos = self.model.cam_pos[0] = self.table_center() + np.array(
            [-0.75, 0.0, 0.75])
        target_pos = self.sim.data.body_xpos[self.name2bid('object_table')]
        self.model.cam_quat[0] = look_at(self.cam_pos, target_pos)

        # reset objects
        self._sample_object_poses(reset_mode=reset_mode)

        for i in range(5 * self.N):
            self.sim.step()
            if self.FLAGS['render']:
                self.viewer.render()

        return not self._object_fell() and not self._geofence_escape()

    def _sample_goal(self):
        """NOTE: must be called before reset, as this can mess with state"""
        goal = {}
        if self.FLAGS['goal_conditioned']:
            # Same mechanism for resetting sim, in order to get to a random state
            did_reset_sim = False
            while not did_reset_sim:
                did_reset_sim = self._reset_sim(reset_mode='cluster')

            obs = self._get_obs()
            obs['array'][:, :3] /= (
                0.5 * TABLES[self.FLAGS['default_table']]['wood'])
            if not self.FLAGS['use_quat']:
                obs['array'] = obs['array'][..., :2]
            goal['array'] = copy.deepcopy(obs['array'])
            if self.FLAGS['use_image']:
                goal['image'] = copy.deepcopy(obs['image'])
            return goal
        else:
            sample_range = R3D(self.FLAGS['obj_rangex'],
                               self.FLAGS['obj_rangey'], R(0, 0))
            goal['array'] = sim_utils.sample_xyz(self.np_random, sample_range)
            return goal

    def get_out_obs(self):
        """This matches the observation space.  The other one is just for internal use"""
        obs = self._get_obs()

        obs['array'][:, :3] /= (0.5 *
                                TABLES[self.FLAGS['default_table']]['wood'])
        if not self.FLAGS['use_quat']:
            obs['array'] = obs['array'][..., :2]

        obs_names = ['array', 'single']
        if self.FLAGS['goal_conditioned']:
            assert not np.all(np.equal(obs['goal_array'],
                                       0.0)), "goal was not set before out_obs"
            obs_names.append('goal_array')
            if self.FLAGS['use_image']:
                obs_names.append('goal_image')
        if self.FLAGS['use_image']:
            obs_names.append('image')
            if self.FLAGS['use_canonical']:
                obs_names.append('canonical')
                self._canonicalize()
                self.sim.step()
                canon_obs = self._get_obs()
                obs['canonical'] = canon_obs['image']

        for key in obs:
            obs[key] = obs[key].astype(np.float32)
        out_obs = copy.deepcopy(tuple([obs[key] for key in sorted(obs_names)]))
        #print(obs['array'])
        #print()
        if np.max(np.abs(obs['array'])) > 1.01:
            print('AHHH out obs', obs['array'])
        return out_obs

    def _canonicalize(self):
        # Canonicalize textures
        for name in self.model.geom_names + ('skybox', ):
            if re.match('object\d', name) is not None:
                self.tex_modder.set_rgb(
                    name, np.array([255, 0, 255], dtype=np.uint8))
            else:
                self.tex_modder.set_rgb(name,
                                        np.array([255] * 3, dtype=np.uint8))
        # light
        for key in self.LIGHT:
            self.LIGHT[key][:] = self.START_LIGHT[key].copy()

        # camera
        if self.FLAGS['mild_canonical']:
            self.cam_pos = self.CAMERA_POS_MEAN + self.table_center()
        else:
            self.cam_pos = self.CAMERA_POS_MEAN + self.table_center() + [
                0.639, -0.025, 0.125
            ]
        target_id = self.model.body_name2id('object_table')
        target_pos = self.sim.data.body_xpos[
            target_id]  #+ sim_utils.sample_xyz(self.np_random, R3D(R(-0.05, 0.05), R(-0.05, 0.05), R(0.1,0.1)))
        quat = look_at(self.cam_pos, target_pos)
        self.cam_modder.set_quat('camera1', quat)
        self.cam_modder.set_pos('camera1', self.cam_pos)
        self.cam_modder.set_fovy('camera1', 45)

        # robot
        if self.FLAGS['baxter']:
            for name in self.model.joint_names:
                if 'paddle' not in name and 'object' not in name:
                    id = self.sim.model.joint_name2id(name)
                    self.sim.data.set_joint_qpos(name, 0.0)
            #self.model.body_pos[self.name2bid('base')] = self.START_BODY_POS[self.name2bid('base')]
        else:
            for joint in ['lbr4_j{}'.format(i) for i in range(7)]:
                id = self.sim.model.joint_name2id(joint)
                self.sim.data.set_joint_qpos(joint, 0.0)

        PREFIX = 'distract'
        geom_names = [
            name for name in self.model.geom_names if name.startswith(PREFIX)
        ]
        for name in geom_names:
            gid = self.model.geom_name2id(name)
            self.model.geom_pos[gid] = np.array([0, 0, -2])
            self.model.geom_rgba[gid][-1] = 0

    def _randomize(self):
        if self.FLAGS['domrand']:
            self._rand_textures()
            self._rand_camera()
            self._rand_lights()
            #self._rand_robot()
            self._rand_distract()
        else:
            self._canonicalize()

    def _rand_textures(self):
        """Randomize all the textures in the scene, including the skybox"""
        bright = self.np_random.binomial(1, 0.8)
        for name in self.sim.model.geom_names + ('skybox', ):
            self.tex_modder.rand_all(name)
            if bright:
                if name == 'object_table':
                    self.tex_modder.brighten(name,
                                             self.np_random.randint(150, 255))
                else:
                    self.tex_modder.brighten(name,
                                             self.np_random.randint(0, 150))

    def _rand_camera(self):
        """Randomize pos, orientation, and fov of camera
        FOVY:
        Kinect2 is 53.8
        ASUS is 45 
        https://www.asus.com/us/3D-Sensor/Xtion_PRO_LIVE/specifications/
        http://smeenk.com/kinect-field-of-view-comparison/
        """
        dx = 0.05
        self.cam_pos = self.CAMERA_POS_MEAN + self.table_center()
        C_R3D = R3D(R(-dx, dx), R(-dx, dx), R(-2 * dx, 2 * dx))
        self.cam_pos += sim_utils.sample_xyz(self.np_random, C_R3D)
        self._rand_camera_angle()
        self.cam_modder.set_pos('camera1', self.cam_pos)
        self.cam_modder.set_fovy('camera1',
                                 sim_utils.sample(self.np_random, R(44, 46)))

    def _rand_camera_angle(self):
        # Look approximately at the robot
        target_id = self.model.body_name2id('object_table')
        target_pos = self.sim.data.body_xpos[target_id] + sim_utils.sample_xyz(
            self.np_random, R3D(R(-0.02, 0.02), R(-0.02, 0.02), R(0.0, 0.2)))
        quat = look_at(self.cam_pos, target_pos)
        self.cam_modder.set_quat('camera1', quat)

    def _rand_lights(self):
        """Randomize pos, direction, and lights"""
        # light stuff
        X = R(-1.0, 1.0)
        Y = R(-0.6, 0.6)
        Z = R(0.1, 1.5)
        LIGHT_R3D = self.table_center()[:, None] + R3D(X, Y, Z)
        LIGHT_UNIF = R3D(R(0, 1), R(0, 1), R(0, 1))

        for i, name in enumerate(self.model.light_names):
            lid = self.model.light_name2id(name)
            # random sample 80% of any given light being on
            if lid != 0:
                self.light_modder.set_active(
                    name,
                    sim_utils.sample(self.np_random, [0, 1]) < 0.8)
                self.light_modder.set_dir(
                    name, sim_utils.sample_light_dir(self.np_random))

            self.light_modder.set_pos(
                name, sim_utils.sample_xyz(self.np_random, LIGHT_R3D))

            spec = np.array([sim_utils.sample(self.np_random, R(0.5, 1))] * 3)
            diffuse = np.array([sim_utils.sample(self.np_random, R(0.5, 1))] *
                               3)
            ambient = np.array([sim_utils.sample(self.np_random, R(0.5, 1))] *
                               3)

            self.light_modder.set_specular(name, spec)
            self.light_modder.set_diffuse(name, diffuse)
            self.light_modder.set_ambient(name, ambient)
            self.model.light_castshadow[lid] = sim_utils.sample(
                self.np_random, [0, 1]) < 0.5

    def _rand_robot(self):
        """Randomize joint angles"""
        if self.FLAGS['baxter']:
            for name in self.model.joint_names:
                if 'paddle' not in name and 'object' not in name:
                    id = self.sim.model.joint_name2id(name)
                    self.sim.data.set_joint_qpos(
                        name,
                        sim_utils.sample(self.np_random,
                                         self.model.jnt_range[id]))

            self.model.body_pos[self.name2bid('base')] = self.START_BODY_POS[
                self.name2bid('base')] + sim_utils.sample_xyz(
                    self.np_random, R3D(R(0, 0), R(0, 0), R(-0.05, 0.05)))

        else:
            for name in ['lbr4_j{}'.format(i) for i in range(7)]:
                id = self.sim.model.joint_name2id(name)
                self.sim.data.set_joint_qpos(
                    name,
                    sim_utils.sample(self.np_random, self.model.jnt_range[id]))

    def _rand_distract(self):
        """Randomize the position and size of the distractor objects"""
        PREFIX = 'distract'
        geom_names = [
            name for name in self.model.geom_names if name.startswith(PREFIX)
        ]

        # Size range
        SX = R(0.01, 0.3)
        SY = R(0.01, 0.3)
        SZ = R(0.01, 0.3)
        S3D = R3D(SX, SY, SZ)

        # Back range
        B_PX = R(0.5, 1.0)
        B_PY = R(-2, 2)
        B_PZ = R(0.1, 0.5)
        B_P3D = R3D(B_PX, B_PY, B_PZ)

        # Front range
        F_PX = R(-0.5, 0.5)
        F_PY = R(-2, 2)
        F_PZ = R(-0.1, 0.3)
        F_P3D = R3D(F_PX, F_PY, F_PZ)

        for name in geom_names:
            gid = self.model.geom_name2id(name)
            range = B_P3D if np.random.binomial(1, 0.5) else F_P3D

            mid = self.table_center().copy()
            mid[2] = -0.925

            self.model.geom_pos[gid] = mid + sim_utils.sample_xyz(
                self.np_random, range)
            self.model.geom_quat[gid] = sim_utils.random_quat(self.np_random)
            self.model.geom_size[gid] = sim_utils.sample_xyz(
                self.np_random, S3D)
            self.model.geom_type[gid] = sim_utils.sample_geom_type(
                self.np_random)
            self.model.geom_rgba[gid][-1] = np.random.binomial(1, 0.5)