def scenery(geometries, n_agents=1, device='cuda', random=np.random):
agentlines = np.tile(agent_model(), (n_agents, 1, 1))
agentcolors = np.tile(agent_colors(), (n_agents, 1))
data = []
for g in geometries:
lights = random_lights(g.lights)
lines = np.concatenate([agentlines, g.walls])
textures, texwidths = init_textures(agentlines, agentcolors, g.walls,
random)
data.append(
arrdict.arrdict(lights=arrdict.arrdict(vals=lights,
widths=len(lights)),
lines=arrdict.arrdict(vals=lines,
widths=len(lines)),
textures=arrdict.arrdict(vals=textures,
widths=texwidths)))
data = arrdict.torchify(arrdict.cat(data)).to(device)
lights = ragged.Ragged(**data['lights'])
scenery = core.cuda.Scenery(n_agents=n_agents,
lights=lights,
lines=ragged.Ragged(**data['lines']),
textures=ragged.Ragged(**data['textures']),
model=arrdict.torchify(
agent_model()).to(device))
core.cuda.bake(scenery)
return scenery
def __init__(self, geometries, core, *args, n_spawns=100, **kwargs):
self.core = core
positions = random_empty_positions(geometries, core.n_agents, n_spawns)
angles = core.random.uniform(
-180, +180, (len(geometries), core.n_agents, n_spawns))
self._spawns = arrdict.torchify(
arrdict.arrdict(positions=positions,
angles=angles)).to(core.device)
def __init__(self, geometries, core, n_spawns=100):
"""Respawns agents to random empty locations in the geometry.
:param geometries: The :ref:`geometry <geometry>` to use when calculating the spawn locations. Should be a list
with one for each environment.
:param core: The :class:`~megastep.core.Core` used by the environment.
:param n_spawns: The number of spawns to choose between for each agent. This is precomputed when the class
is created, so that the respawns themselves are fast.
"""
self.core = core
positions = random_empty_positions(geometries, core.n_agents, n_spawns)
angles = core.random.uniform(-180, +180, (len(geometries), core.n_agents, n_spawns))
self._spawns = arrdict.torchify(arrdict.arrdict(positions=positions, angles=angles)).to(core.device)
def __init__(self, n_envs, n_agents, *args, **kwargs):
geometries = cubicasa.sample(max(n_envs // 4, 1))
scenery = scene.scenery(geometries, n_agents)
self.core = core.Core(scenery, *args, res=4 * 128, fov=70, **kwargs)
self._rgb = modules.RGB(self.core, n_agents=1, subsample=4)
self._depth = modules.Depth(self.core, n_agents=1, subsample=4)
self._imu = modules.IMU(self.core, n_agents=1)
self._movement = modules.MomentumMovement(self.core, n_agents=1)
self._spawner = modules.RandomSpawns(geometries, self.core)
self.action_space = self._movement.space
self.obs_space = dotdict.dotdict(rgb=self._rgb.space,
d=self._depth.space,
imu=self._imu.space,
health=spaces.MultiVector(1, 1))
self._bounds = arrdict.torchify(
np.stack([g.masks.shape * g.res
for g in geometries])).to(self.core.device)
self._health = self.core.agent_full(np.nan)
self._damage = self.core.agent_full(np.nan)
self.n_envs = self.core.n_envs * self.core.n_agents
self.device = self.core.device
def torchify(self):
"""Applies :func:`arrdict.torchify` to the backing arrays and returns a new :class:`cuda.Ragged$ND` for them"""
return Ragged(arrdict.torchify(self.vals),
arrdict.torchify(self.widths))