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
def in_cone2d(origin_pts, origin_angles, cone_angle, target_pts): ''' Computes whether 2D points target_pts are in the cones originating from origin_pts at angle origin_angles with cone spread angle cone_angle. Args: origin_pts (np.ndarray): array with shape (n_points, 2) of origin points origin_angles (np.ndarray): array with shape (n_points,) of origin angles cone_angle (float): cone angle width target_pts (np.ndarray): target points to check whether in cones Returns: np.ndarray of bools. Each row corresponds to origin cone, and columns to target points ''' assert isinstance(origin_pts, np.ndarray) assert isinstance(origin_angles, np.ndarray) assert isinstance(cone_angle, float) assert isinstance(target_pts, np.ndarray) assert origin_pts.shape[0] == origin_angles.shape[0] assert len(origin_angles.shape) == 1, "Angles should only have 1 dimension" np.seterr(divide='ignore', invalid='ignore') cone_vec = np.array([np.cos(origin_angles), np.sin(origin_angles)]).T # Compute normed vectors between all pairs of agents pos_diffs = target_pts[None, ...] - origin_pts[:, None, :] norms = np.sqrt(np.sum(np.square(pos_diffs), -1, keepdims=True)) unit_diffs = pos_diffs / norms # Dot product between unit vector in middle of cone and the vector dot_cone_diff = np.sum(unit_diffs * cone_vec[:, None, :], -1) angle_between = np.arccos(dot_cone_diff) # Right now the only thing that should be nan will be targets that are on the origin point # This can only happen for the origin looking at itself, so just make this always true angle_between[np.isnan(angle_between)] = 0. return np.abs(normalize_angles(angle_between)) <= cone_angle
def observation_step(self, env, sim): qpos = sim.data.qpos.copy() qvel = sim.data.qvel.copy() if self.make_static: s_cylinder_geom_idxs = np.expand_dims(self.s_cylinder_geom_idxs, -1) s_cylinder_xpos = sim.data.geom_xpos[self.s_cylinder_geom_idxs] obs = { 'static_cylinder_geom_idxs': s_cylinder_geom_idxs, 'static_cylinder_xpos': s_cylinder_xpos } else: m_cylinder_geom_idxs = np.expand_dims(self.m_cylinder_geom_idxs, -1) m_cylinder_xpos = sim.data.geom_xpos[self.m_cylinder_geom_idxs] m_cylinder_qpos = qpos[self.m_cylinder_qpos_idxs] m_cylinder_qvel = qvel[self.m_cylinder_qvel_idxs] mc_angle = normalize_angles(m_cylinder_qpos[:, 3:]) polar_angle = np.concatenate( [np.cos(mc_angle), np.sin(mc_angle)], -1) m_cylinder_qpos = np.concatenate( [m_cylinder_qpos[:, :3], polar_angle], -1) m_cylinder_obs = np.concatenate([m_cylinder_qpos, m_cylinder_qvel], -1) obs = { 'moveable_cylinder_geom_idxs': m_cylinder_geom_idxs, 'moveable_cylinder_xpos': m_cylinder_xpos, 'moveable_cylinder_obs': m_cylinder_obs } return obs
def observation_step(self, env, sim): qpos = sim.data.qpos.copy() qvel = sim.data.qvel.copy() ramp_geom_idxs = np.expand_dims(self.ramp_geom_idxs, -1) ramp_qpos = qpos[self.ramp_qpos_idxs] ramp_qvel = qvel[self.ramp_qvel_idxs] ramp_angle = normalize_angles(ramp_qpos[:, 3:]) polar_angle = np.concatenate( [np.cos(ramp_angle), np.sin(ramp_angle)], -1) if self.polar_obs: ramp_qpos = np.concatenate([ramp_qpos[:, :3], polar_angle], -1) ramp_obs = np.concatenate([ramp_qpos, ramp_qvel], -1) if self.pad_ramp_size: ramp_obs = np.concatenate( [ramp_obs, np.zeros((ramp_obs.shape[0], 3))], -1) obs = { 'ramp_obs': ramp_obs, 'ramp_angle': ramp_angle, 'ramp_geom_idxs': ramp_geom_idxs, 'ramp_pos': ramp_qpos[:, :3] } return obs
def observation_step(self, env, sim): qpos = sim.data.qpos.copy() qvel = sim.data.qvel.copy() # Agent agent_qpos = qpos[self.agent_qpos_idxs] agent_qvel = qvel[self.agent_qvel_idxs] # Normalize angles in range 0 to 6.28 agent_angle = (clip_angle_range(agent_qpos[:, [-1]]) + 4.71239) % 6.28319 agent_qpos[:, [-1]] = agent_angle = normalize_angles(agent_angle) # Normalize positions in range 0 to 1, angle in range -1 to 1 agent_qpos[:, [0]] = agent_qpos[:, [0]] / self.placement_size[0] agent_qpos[:, [1]] = agent_qpos[:, [1]] / self.placement_size[1] agent_qpos[:, [-1]] = agent_qpos[:, [-1]] / np.pi # Normalize qvel in range -1 to 1 max_speed = np.sqrt(self.action_scale[0]**2 + self.action_scale[1]**2) agent_qvel[:, [0]] = agent_qvel[:, [0]] / max_speed agent_qvel[:, [1]] = agent_qvel[:, [1]] / max_speed agent_qvel[:, [3]] = agent_qvel[:, [3]] / 4.71239 # Form observation, remove z-pos from qpos and qvel agent_qpos = np.array([[aqpos[0], aqpos[1], aqpos[3]] for aqpos in agent_qpos]) agent_qvel = np.array([[aqvel[0], aqvel[1], aqvel[3]] for aqvel in agent_qvel]) agent_qpos_qvel = np.concatenate([agent_qpos, agent_qvel], -1) # Velocimeter info velocimeter_info = [] for i in range(self.n_agents): velocimeter_idx = sim.model.sensor_name2id(f"agent{i}:velocimeter") velocimeter_data = sim.data.sensordata[ velocimeter_idx:velocimeter_idx + 2] velocimeter_data[0] /= self.action_scale[0] velocimeter_data[1] /= self.action_scale[1] velocimeter_info.append(velocimeter_data) agent_local_qvel = np.array(velocimeter_info) obs = { 'agent_qpos_qvel': agent_qpos_qvel, 'agent_angle': agent_angle, 'agent_pos': agent_qpos[:, :3], 'agent_local_qvel': agent_local_qvel } return obs
def observation_step(self, env, sim): qpos = sim.data.qpos.copy() qvel = sim.data.qvel.copy() agent_qpos = qpos[self.agent_qpos_idxs] agent_qvel = qvel[self.agent_qvel_idxs] agent_angle = agent_qpos[:, [ -1 ]] - np.pi / 2 # Rotate the angle to match visual front agent_qpos_qvel = np.concatenate([agent_qpos, agent_qvel], -1) polar_angle = np.concatenate( [np.cos(agent_angle), np.sin(agent_angle)], -1) if self.polar_obs: agent_qpos = np.concatenate([agent_qpos[:, :-1], polar_angle], -1) agent_angle = normalize_angles(agent_angle) obs = { 'agent_qpos_qvel': agent_qpos_qvel, 'agent_angle': agent_angle, 'agent_pos': agent_qpos[:, :3] } return obs