def build_world_step(self, env, floor, floor_size):
env.metadata['n_agents'] = self.n_agents
successful_placement = True
for i in range(self.n_agents):
env.metadata.pop(f"agent{i}_initpos", None)
for i in range(self.n_agents):
obj = ObjFromXML("particle_hinge", name=f"agent{i}")
if self.friction is not None:
obj.add_transform(
set_geom_attr_transform('friction', self.friction))
if self.color is not None:
_color = (self.color[i] if isinstance(
self.color[0], (list, tuple, np.ndarray)) else self.color)
obj.add_transform(set_geom_attr_transform('rgba', _color))
if not self.damp_z:
obj.add_transform(set_joint_damping_transform(1, 'tz'))
if self.placement_fn is not None:
_placement_fn = (self.placement_fn[i] if isinstance(
self.placement_fn, list) else self.placement_fn)
obj_size = get_size_from_xml(obj)
pos, pos_grid = rejection_placement(env, _placement_fn,
floor_size, obj_size)
if pos is not None:
floor.append(obj, placement_xy=pos)
# store spawn position in metadata. This allows sampling subsequent agents
# close to previous agents
env.metadata[f"agent{i}_initpos"] = pos_grid
else:
successful_placement = False
else:
floor.append(obj)
return successful_placement
def build_world_step(self, env, floor, floor_size):
env.metadata['food_size'] = self.food_size
self.curr_n_food = env._random_state.randint(self.n_food[0], self.n_food[1] + 1)
env.metadata['max_n_food'] = self.n_food[1]
env.metadata['curr_n_food'] = self.curr_n_food
successful_placement = True
for i in range(self.curr_n_food):
env.metadata.pop(f"food{i}_initpos", None)
# Add food sites
for i in range(self.curr_n_food):
if self.placement_fn is not None:
_placement_fn = (self.placement_fn[i]
if isinstance(self.placement_fn, list)
else self.placement_fn)
pos, pos_grid = rejection_placement(env, _placement_fn, floor_size,
np.array([self.food_size, self.food_size]))
if pos is not None:
floor.mark(f"food{i}", relative_xyz=np.append(pos, [self.food_size / 2]),
size=(self.food_size, self.food_size, self.food_size),
rgba=(0., 1., 0., 1.))
# store spawn position in metadata. This allows sampling subsequent food items
# close to previous food items
env.metadata[f"food{i}_initpos"] = pos_grid
else:
successful_placement = False
else:
floor.mark(f"food{i}", rgba=(0., 1., 0., 1.),
size=(self.food_size, self.food_size, self.food_size))
return successful_placement
def build_world_step(self, env, floor, floor_size):
env.metadata['box_size'] = self.box_size
self.curr_n_boxes = env._random_state.randint(self.n_boxes[0],
self.n_boxes[1] + 1)
env.metadata['curr_n_boxes'] = np.zeros((self.n_boxes[1]))
env.metadata['curr_n_boxes'][:self.curr_n_boxes] = 1
env.metadata['curr_n_boxes'] = env.metadata['curr_n_boxes'].astype(
np.bool)
self.curr_n_elongated_boxes = env._random_state.randint(
self.n_elongated_boxes[0],
min(self.n_elongated_boxes[1], self.curr_n_boxes) + 1)
self.box_size_array = self.box_size * np.ones((self.curr_n_boxes, 3))
if self.curr_n_elongated_boxes > 0:
# sample number of x-aligned boxes
n_xaligned = env._random_state.randint(
self.curr_n_elongated_boxes + 1)
self.box_size_array[:n_xaligned, :] = self.box_size * np.array(
[3.3, 0.3, 1.0])
self.box_size_array[n_xaligned:self.curr_n_elongated_boxes, :] = (
self.box_size * np.array([0.3, 3.3, 1.0]))
env.metadata['box_size_array'] = self.box_size_array
successful_placement = True
for i in range(self.curr_n_boxes):
char = chr(ord('A') + i % 26)
geom = Geom("box",
self.box_size_array[i, :],
name=f'moveable_box{i}')
geom.set_material(Material(texture="chars/" + char + ".png"))
geom.add_transform(set_geom_attr_transform('mass', self.box_mass))
if self.mark_box_corners:
for j, (x, y) in enumerate([[0, 0], [0, 1], [1, 0], [1, 1]]):
geom.mark(f'moveable_box{i}_corner{j}',
relative_xyz=(x, y, 0.5),
rgba=[1., 1., 1., 0.])
if self.friction is not None:
geom.add_transform(
set_geom_attr_transform('friction', self.friction))
if self.box_only_z_rot:
geom.add_transform(
remove_hinge_axis_transform(np.array([1.0, 0.0, 0.0])))
geom.add_transform(
remove_hinge_axis_transform(np.array([0.0, 1.0, 0.0])))
if self.placement_fn is not None:
_placement_fn = (self.placement_fn[i] if isinstance(
self.placement_fn, list) else self.placement_fn)
pos, _ = rejection_placement(env, _placement_fn, floor_size,
self.box_size_array[i, :2])
if pos is not None:
floor.append(geom, placement_xy=pos)
else:
successful_placement = False
else:
floor.append(geom)
return successful_placement
def build_world_step(self, env, floor, floor_size):
default_name = 'static_cylinder' if self.make_static else 'moveable_cylinder'
diameter = env._random_state.uniform(self.diameter[0],
self.diameter[1])
height = env._random_state.uniform(self.height[0], self.height[1])
obj_size = (diameter, height, 0)
successful_placement = True
for i in range(self.n_objects):
geom = Geom('cylinder',
obj_size,
name=f'{default_name}{i}',
rgba=self.rgba)
if self.make_static:
geom.mark_static()
if self.placement_fn is not None:
_placement_fn = (self.placement_fn[i] if isinstance(
self.placement_fn, list) else self.placement_fn)
pos, _ = rejection_placement(env, _placement_fn, floor_size,
diameter * np.ones(2))
if pos is not None:
floor.append(geom, placement_xy=pos)
else:
successful_placement = False
else:
floor.append(geom)
return successful_placement
def build_world_step(self, env, floor, floor_size):
successful_placement = True
env.metadata['curr_n_ramps'] = np.ones((self.n_ramps)).astype(np.bool)
for i in range(self.n_ramps):
char = chr(ord('A') + i % 26)
geom = geom = ObjFromXML('ramp', name=f"ramp{i}")
geom.set_material(Material(texture="chars/" + char + ".png"))
if self.friction is not None:
geom.add_transform(
set_geom_attr_transform('friction', self.friction))
if self.placement_fn is not None:
_placement_fn = (self.placement_fn[i] if isinstance(
self.placement_fn, list) else self.placement_fn)
pos, _ = rejection_placement(env, _placement_fn, floor_size,
get_size_from_xml(geom))
if pos is not None:
floor.append(geom, placement_xy=pos)
else:
successful_placement = False
else:
floor.append(geom)
return successful_placement
def build_world_step(self, env, floor, floor_size):
self.curr_n_sites = env._random_state.randint(self.n_sites[0],
self.n_sites[1] + 1)
self.curr_n_elongated_sites = env._random_state.randint(
self.n_elongated_sites[0], self.n_elongated_sites[1] + 1)
env.metadata['curr_n_sites'] = self.curr_n_sites
env.metadata['curr_n_elongated_sites'] = self.curr_n_elongated_sites
self.site_size_array = self.site_size * np.ones((self.curr_n_sites, 2))
if self.curr_n_elongated_sites > 0:
n_xaligned = env._random_state.randint(
self.curr_n_elongated_sites + 1)
self.site_size_array[:n_xaligned, :] = self.site_size * np.array(
[3.3, 0.3])
self.site_size_array[n_xaligned:self.curr_n_elongated_sites, :] = (
self.site_size * np.array([0.3, 3.3]))
successful_placement = True
for i in range(self.curr_n_sites):
if self.placement_fn is not None:
_placement_fn = (self.placement_fn[i] if isinstance(
self.placement_fn, list) else self.placement_fn)
pos, _ = rejection_placement(env, _placement_fn, floor_size,
self.site_size_array[i])
if pos is not None:
self._mark_site_square(floor, floor_size,
f'{self.site_name}{i}',
(pos[0], pos[1], self.site_height),
self.site_size_array[i])
else:
successful_placement = False
else:
# place the site so that all the corners are still within the play area
pos_min = self.site_size_array[i].max() / (floor_size *
1.1) / 2
pos = env._random_state.uniform(pos_min, 1 - pos_min, 2)
self._mark_site_square(floor, floor_size,
f'{self.site_name}{i}',
(pos[0], pos[1], self.site_height),
self.site_size_array[i])
return successful_placement