def insight(sim,
geom1_id,
site1_id=None,
geom2_id=None,
pt1=None,
pt2=None,
dist_thresh=np.inf,
check_body=False):
'''
Check if geom2 or pt2 is in line of sight of geom1.
Args:
sim: Mujoco sim object
geom1 (int): geom id
geom2 (int): geom id
pt2 (tuple): xy point
dist_thresh (float): Adds a distance threshold for vision. Objects beyond the threshold
are considered out of sight.
check_body (bool): Check whether the raycast hit any geom in the body that geom2 is in
rather than if it just hit geom2
'''
if geom1_id is not None:
dist, collision_geom = raycast(sim,
geom1_id=geom1_id,
geom2_id=geom2_id,
pt2=pt2)
elif site1_id is not None:
dist, collision_geom = raycast(sim,
site1_id=site1_id,
geom2_id=geom2_id,
pt2=pt2)
else:
dist, collision_geom = raycast(sim,
pt1=pt1,
geom2_id=geom2_id,
pt2=pt2)
if geom2_id is not None:
if check_body:
body2_id, collision_body_id = sim.model.geom_bodyid[[
geom2_id, collision_geom
]]
return (collision_body_id == body2_id and dist < dist_thresh)
else:
return (collision_geom == geom2_id and dist < dist_thresh)
else:
pt1 = sim.data.geom_xpos[geom1_id]
dist_pt2 = np.linalg.norm(pt2 - pt1)
# if dist == -1 then we're raycasting from a geom to a point within itself,
# and all objects have line of sight of themselves.
return (dist == -1.0 or dist > dist_pt2) and dist_pt2 < dist_thresh
def observation(self, obs):
sim = self.unwrapped.sim
agent_pos = sim.data.body_xpos[self.agent_body_ids]
lidar_endpoints = agent_pos[:, None, :] + self.lidar_rays
# Would be nice to vectorize in the future with better mujoco-py interface
lidar = np.zeros((self.n_agents, self.n_lidar_per_agent))
for i in range(self.n_agents):
for j in range(self.n_lidar_per_agent):
lidar[i, j] = raycast(sim,
geom1_id=self.agent_geom_ids[i],
pt2=lidar_endpoints[i, j],
geom_group=None)[0]
lidar[lidar < 0.0] = self.lidar_range
if self.compress_lidar_scale is not None:
obs['lidar'] = (
self.compress_lidar_scale *
np.tanh(lidar[..., None] / self.compress_lidar_scale))
else:
obs['lidar'] = lidar[..., None]
if self.visualize_lidar:
# recalculate lidar endpoints
lidar_endpoints = agent_pos[:, None, :] + \
lidar[:, :, None] / self.lidar_range * self.lidar_rays
self.place_lidar_ray_markers(agent_pos, lidar_endpoints)
sim.model.site_rgba[self.lidar_ids, :] = np.array(
[0.0, 0.0, 1.0, 0.2])
sim.forward()
return obs
def step(self, action):
obs, rew, done, info = self.env.step(action)
target_geom = obs['static_cylinder_geom_idxs'][0, 0]
rew = rew + np.zeros((self.unwrapped.n_agents, ))
for pt in self.ray_start_points:
_, collision_geom = raycast(self.sim, pt1=pt, geom2_id=target_geom)
if collision_geom == target_geom:
rew -= 1
rew *= self.reward_scale
return obs, rew, done, info
