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
