def observation_step()

in mae_envs/modules/objects.py [0:0]

• 25 lines of code
• 6 McCabe index (conditional complexity)

    def observation_step(self, env, sim):
        qpos = sim.data.qpos.copy()
        qvel = sim.data.qvel.copy()

        box_inds = np.expand_dims(np.arange(self.curr_n_boxes), -1)
        box_geom_idxs = np.expand_dims(self.box_geom_idxs, -1)
        box_qpos = qpos[self.box_qpos_idxs]
        box_qvel = qvel[self.box_qvel_idxs]
        box_angle = normalize_angles(box_qpos[:, 3:])
        polar_angle = np.concatenate([np.cos(box_angle), np.sin(box_angle)], -1)
        if self.polar_obs:
            box_qpos = np.concatenate([box_qpos[:, :3], polar_angle], -1)
        box_obs = np.concatenate([box_qpos, box_qvel], -1)

        if self.boxid_obs:
            box_obs = np.concatenate([box_obs, box_inds], -1)
        if self.n_elongated_boxes[1] > 0 or self.boxsize_obs:
            box_obs = np.concatenate([box_obs, self.box_size_array], -1)

        obs = {'box_obs': box_obs,
               'box_angle': box_angle,
               'box_geom_idxs': box_geom_idxs,
               'box_pos': box_qpos[:, :3],
               'box_xpos': sim.data.geom_xpos[self.box_geom_idxs]}

        if self.mark_box_corners:
            obs.update({'box_corner_pos': sim.data.site_xpos[self.box_corner_idxs],
                        'box_corner_idxs': np.expand_dims(self.box_corner_idxs, -1)})

        return obs