def on_reset(self):
goal_xyhw = DEFAULT_GOAL_XYHW
if self.rand_poses_minor or self.rand_poses_full:
# randomise width and height of goal region
# (unfortunately this has to be done before pose randomisation b/c
# I don't have an easy way of changing size later)
if self.rand_poses_minor:
hw_bound = self.JITTER_TARGET_BOUND
else:
hw_bound = None
sampled_hw = geom.randomise_hw(self.RAND_GOAL_MIN_SIZE,
self.RAND_GOAL_MAX_SIZE,
self.rng,
current_hw=goal_xyhw[2:],
linf_bound=hw_bound)
goal_xyhw = (*goal_xyhw[:2], *sampled_hw)
# colour the goal region
if self.rand_goal_colour:
goal_colour = self.rng.choice(
np.asarray(en.SHAPE_COLOURS, dtype='object'))
else:
goal_colour = DEFAULT_GOAL_COLOUR
# place the goal region
assert len(goal_xyhw) == 4, goal_xyhw
goal = en.GoalRegion(*goal_xyhw, goal_colour)
self.add_entities([goal])
self.__goal_ref = goal
# now place the robot
default_robot_pos, default_robot_angle = DEFAULT_ROBOT_POSE
robot = self._make_robot(default_robot_pos, default_robot_angle)
self.add_entities([robot])
self.__robot_ent_index = en.EntityIndex([robot])
if self.rand_poses_minor or self.rand_poses_full:
if self.rand_poses_minor:
# limit amount by which position and rotation can be randomised
pos_limits = self.JITTER_POS_BOUND
rot_limits = [None, self.JITTER_ROT_BOUND]
else:
# no limits, can randomise as much as needed
assert self.rand_poses_full
pos_limits = rot_limits = None
geom.pm_randomise_all_poses(self._space,
(self.__goal_ref, self._robot),
self.ARENA_BOUNDS_LRBT,
rng=self.rng,
rand_pos=True,
rand_rot=(False, True),
rel_pos_linf_limits=pos_limits,
rel_rot_limits=rot_limits)
def on_reset(self):
# make the robot
robot_pos = np.asarray((-0.5, 0.1))
robot_angle = -math.pi * 1.2
# if necessary, robot pose is randomised below
robot = self._make_robot(robot_pos, robot_angle)
# set up target colour/region/pose
if self.rand_target_colour:
target_colour = self.rng.choice(en.SHAPE_COLOURS)
else:
target_colour = en.ShapeColour.GREEN
distractor_colours = [
c for c in en.SHAPE_COLOURS if c != target_colour
]
target_h = 0.7
target_w = 0.6
target_x = 0.1
target_y = 0.7
if self.rand_layout_minor or self.rand_layout_full:
if self.rand_layout_minor:
hw_bound = self.JITTER_TARGET_BOUND
else:
hw_bound = None
target_h, target_w = geom.randomise_hw(self.RAND_GOAL_MIN_SIZE,
self.RAND_GOAL_MAX_SIZE,
self.rng,
current_hw=(target_h,
target_w),
linf_bound=hw_bound)
sensor = en.GoalRegion(target_x, target_y, target_h, target_w,
target_colour)
self.add_entities([sensor])
self.__sensor_ref = sensor
# set up spec for remaining blocks
default_target_types = [
en.ShapeType.STAR,
en.ShapeType.SQUARE,
]
default_distractor_types = [
[],
[en.ShapeType.PENTAGON],
[en.ShapeType.CIRCLE, en.ShapeType.PENTAGON],
]
default_target_poses = [
# (x, y, theta)
(0.8, -0.7, 2.37),
(-0.68, 0.72, 1.28),
]
default_distractor_poses = [
# (x, y, theta)
[],
[(-0.05, -0.2, -1.09)],
[(-0.75, -0.55, 2.78), (0.3, -0.82, -1.15)],
]
if self.rand_shape_count:
target_count = self.rng.randint(1, 2 + 1)
distractor_counts = [
self.rng.randint(0, 2 + 1) for c in distractor_colours
]
else:
target_count = len(default_target_types)
distractor_counts = [len(lst) for lst in default_distractor_types]
if self.rand_shape_type:
shape_types_np = np.asarray(en.SHAPE_TYPES, dtype='object')
target_types = [
self.rng.choice(shape_types_np) for _ in range(target_count)
]
distractor_types = [[
self.rng.choice(shape_types_np) for _ in range(dist_count)
] for dist_count in distractor_counts]
else:
target_types = default_target_types
distractor_types = default_distractor_types
if self.rand_layout_full:
# will do post-hoc randomisation at the end
target_poses = [(0, 0, 0)] * target_count
distractor_poses = [[(0, 0, 0)] * dcount
for dcount in distractor_counts]
else:
target_poses = default_target_poses
distractor_poses = default_distractor_poses
assert len(target_types) == target_count
assert len(target_poses) == target_count
assert len(distractor_types) == len(distractor_counts)
assert len(distractor_types) == len(distractor_colours)
assert len(distractor_types) == len(distractor_poses)
assert all(
len(types) == dcount
for types, dcount in zip(distractor_types, distractor_counts))
assert all(
len(poses) == dcount
for poses, dcount in zip(distractor_poses, distractor_counts))
self.__target_shapes = [
self._make_shape(shape_type=shape_type,
colour_name=target_colour,
init_pos=(shape_x, shape_y),
init_angle=shape_angle)
for shape_type, (shape_x, shape_y,
shape_angle) in zip(target_types, target_poses)
]
self.__distractor_shapes = []
for dist_colour, dist_types, dist_poses \
in zip(distractor_colours, distractor_types, distractor_poses):
for shape_type, (shape_x, shape_y, shape_angle) \
in zip(dist_types, dist_poses):
dist_shape = self._make_shape(shape_type=shape_type,
colour_name=dist_colour,
init_pos=(shape_x, shape_y),
init_angle=shape_angle)
self.__distractor_shapes.append(dist_shape)
shape_ents = self.__target_shapes + self.__distractor_shapes
self.add_entities(shape_ents)
# add this last so it shows up on top, but before layout randomisation,
# since it needs to be added to the space before randomising
self.add_entities([robot])
if self.rand_layout_minor or self.rand_layout_full:
all_ents = (sensor, robot, *shape_ents)
if self.rand_layout_minor:
# limit amount by which position and rotation can be randomised
pos_limits = self.JITTER_POS_BOUND
rot_limits = self.JITTER_ROT_BOUND
else:
# no limits, can randomise as much as needed
assert self.rand_layout_full
pos_limits = rot_limits = None
# randomise rotations of all entities but goal region
rand_rot = [False] + [True] * (len(all_ents) - 1)
geom.pm_randomise_all_poses(self._space,
all_ents,
self.ARENA_BOUNDS_LRBT,
rng=self.rng,
rand_pos=True,
rand_rot=rand_rot,
rel_pos_linf_limits=pos_limits,
rel_rot_limits=rot_limits)
# set up index for lookups
self.__ent_index = en.EntityIndex(shape_ents)
def on_reset(self):
robot_pos, robot_angle = DEFAULT_ROBOT_POSE
robot = self._make_robot(robot_pos, robot_angle)
block_colours = DEFAULT_BLOCK_COLOURS
region_colours = DEFAULT_REGION_COLOURS
block_shapes = DEFAULT_BLOCK_SHAPES
block_poses = DEFAULT_BLOCK_POSES
region_xyhws = DEFAULT_REGION_XYHWS
# randomise count
n_regions = len(block_colours)
if self.rand_count:
n_regions = self.rng.randint(MIN_REGIONS, MAX_REGIONS + 1)
block_poses = block_poses[:1] * n_regions
region_xyhws = region_xyhws[:1] * n_regions
# randomise colours
if self.rand_colours:
region_colours = self.rng.choice(en.SHAPE_COLOURS,
size=n_regions).tolist()
block_colours = list(region_colours)
# randomly choose one block to be the odd one out
odd_idx = self.rng.randint(len(block_colours))
new_col_idx = self.rng.randint(len(en.SHAPE_COLOURS) - 1)
if en.SHAPE_COLOURS[new_col_idx] == block_colours[odd_idx]:
new_col_idx += 1
block_colours[odd_idx] = en.SHAPE_COLOURS[new_col_idx]
# randomise shapes
if self.rand_shapes:
block_shapes = self.rng.choice(en.SHAPE_TYPES,
size=n_regions).tolist()
# create all regions, randomising their height/width first if necessary
if self.rand_layout_minor or self.rand_layout_full:
if self.rand_layout_minor:
hw_bound = self.JITTER_TARGET_BOUND
else:
hw_bound = None
rand_hw = functools.partial(geom.randomise_hw,
MIN_GOAL_SIZE,
MAX_GOAL_SIZE,
self.rng,
linf_bound=hw_bound)
region_xyhws = [(x, y, *rand_hw(current_hw=hw))
for x, y, *hw in region_xyhws]
sensors = [
en.GoalRegion(*xyhw, colour)
for colour, xyhw in zip(region_colours, region_xyhws)
]
self.add_entities(sensors)
self._sensors = sensors
# set up blocks
blocks = []
self._target_blocks = []
for bshape, bcol, tcol, (bpos,
bangle) in zip(block_shapes, block_colours,
region_colours, block_poses):
new_block = self._make_shape(shape_type=bshape,
colour_name=bcol,
init_pos=bpos,
init_angle=bangle)
blocks.append(new_block)
if bcol != tcol:
# we don't want this region to contain any blocks
self._target_blocks.append([])
else:
# we want this region to keep the original block
self._target_blocks.append([new_block])
self.add_entities(blocks)
# add robot last (for draw order reasons)
self.add_entities([robot])
if self.rand_layout_minor or self.rand_layout_full:
sensors_robot = (*sensors, robot)
if self.rand_layout_minor:
pos_limits = self.JITTER_POS_BOUND
rot_limit = self.JITTER_ROT_BOUND
else:
assert self.rand_layout_full
pos_limits = rot_limit = None
# don't randomise rotations of goal regions
rand_rot = [False] * n_regions + [True]
# don't collision check with blocks
block_pm_shapes = sum((b.shapes for b in blocks), [])
# randomise positions of goal regions and robot
geom.pm_randomise_all_poses(self._space,
sensors_robot,
self.ARENA_BOUNDS_LRBT,
rng=self.rng,
rand_pos=True,
rand_rot=rand_rot,
rel_pos_linf_limits=pos_limits,
rel_rot_limits=rot_limit,
ignore_shapes=block_pm_shapes)
# shift blocks out of the way of each other
for block, sensor in zip(blocks, sensors):
geom.pm_shift_bodies(self._space,
block.bodies,
position=sensor.bodies[0].position)
# randomise each inner block position separately
for block, sensor, (_, _,
*sensor_hw) in zip(blocks, sensors,
region_xyhws):
block_pos_limit = max(0, min(sensor_hw) / 2 - self.SHAPE_RAD)
if self.rand_layout_minor:
block_pos_limit = min(self.JITTER_POS_BOUND,
block_pos_limit)
geom.pm_randomise_pose(self._space,
block.bodies,
self.ARENA_BOUNDS_LRBT,
ignore_shapes=sensor.shapes,
rng=self.rng,
rand_pos=True,
rand_rot=True,
rel_pos_linf_limit=block_pos_limit,
rel_rot_limit=rot_limit)
# block lookup index
self._block_index = en.EntityIndex(blocks)
def on_reset(self):
# make the robot (pose is randomised at end, if necessary)
robot_pos, robot_angle = DEFAULT_ROBOT_POSE
robot = self._make_robot(robot_pos, robot_angle)
# sort out shapes and colours of blocks, including the target
# shape/colour
query_colour = DEFAULT_QUERY_COLOUR
query_shape = DEFAULT_QUERY_SHAPE
out_block_colours = DEFAULT_OUT_BLOCK_COLOURS
out_block_shapes = DEFAULT_OUT_BLOCK_SHAPES
n_out_blocks = len(DEFAULT_OUT_BLOCK_COLOURS)
if self.rand_count:
# have between 1 and 5 randomly generated blocks, plus one outside
# block that always matches the query block
n_out_blocks = self.rng.randint(1, 5 + 1) + 1
n_distractors = n_out_blocks - 1
if self.rand_colours:
query_colour = self.rng.choice(en.SHAPE_COLOURS)
out_block_colours = self.rng.choice(en.SHAPE_COLOURS,
size=n_distractors).tolist()
# last block always matches the query
out_block_colours.append(query_colour)
if self.rand_shapes:
query_shape = self.rng.choice(en.SHAPE_TYPES)
out_block_shapes = self.rng.choice(en.SHAPE_TYPES,
size=n_distractors).tolist()
out_block_shapes.append(query_shape)
# create the target region
target_region_xyhw = DEFAULT_TARGET_REGION_XYHW
if self.rand_layout_minor or self.rand_layout_full:
if self.rand_layout_minor:
hw_bound = self.JITTER_TARGET_BOUND
else:
hw_bound = None
target_hw = geom.randomise_hw(self.RAND_GOAL_MIN_SIZE,
self.RAND_GOAL_MAX_SIZE,
self.rng,
current_hw=target_region_xyhw[2:],
linf_bound=hw_bound)
target_region_xyhw = (*target_region_xyhw[:2], *target_hw)
sensor = en.GoalRegion(*target_region_xyhw, query_colour)
self.add_entities([sensor])
self.__sensor_ref = sensor
# set up poses, as necessary
query_block_pose = DEFAULT_QUERY_BLOCK_POSE
out_block_poses = DEFAULT_OUT_BLOCK_POSES
if self.rand_count:
# this will be filled out when randomising the layout
out_block_poses = [((0, 0), 0)] * n_out_blocks
# The "outside blocks" are those which need to be outside the goal
# region in the initial state. That includes the clone of the query
# block that we intentionally inserted above.
outside_blocks = []
# we also keep a list of blocks that count as target blocks
self.__target_set = set()
for bshape, bcol, (bpos, bangle) in zip(out_block_shapes,
out_block_colours,
out_block_poses):
new_block = self._make_shape(shape_type=bshape,
colour_name=bcol,
init_pos=bpos,
init_angle=bangle)
outside_blocks.append(new_block)
if bcol == query_colour and bshape == query_shape:
self.__target_set.add(new_block)
self.add_entities(outside_blocks)
# now add the query block
assert len(query_block_pose) == 2
query_block = self._make_shape(shape_type=query_shape,
colour_name=query_colour,
init_pos=query_block_pose[0],
init_angle=query_block_pose[1])
self.__target_set.add(query_block)
self.add_entities([query_block])
# these are shapes that shouldn't end up in the goal region
self.__distractor_set = set(outside_blocks) - self.__target_set
# add robot last, but before layout randomisation
self.add_entities([robot])
if self.rand_layout_minor or self.rand_layout_full:
all_ents = (sensor, robot, *outside_blocks)
if self.rand_layout_minor:
# limit amount by which position and rotation can be randomised
pos_limits = self.JITTER_POS_BOUND
rot_limit = self.JITTER_ROT_BOUND
else:
# no limits, can randomise as much as needed
assert self.rand_layout_full
pos_limits = rot_limit = None
# randomise rotations of all entities but goal region
rand_rot = [False] + [True] * (len(all_ents) - 1)
geom.pm_randomise_all_poses(self._space,
all_ents,
self.ARENA_BOUNDS_LRBT,
rng=self.rng,
rand_pos=True,
rand_rot=rand_rot,
rel_pos_linf_limits=pos_limits,
rel_rot_limits=rot_limit,
ignore_shapes=query_block.shapes)
# randomise the query block last, since it must be mostly inside
# the (randomly-placed) sensor region
query_pos_limit = max(
0,
min(target_region_xyhw[2:]) / 2 - self.SHAPE_RAD / 2)
if self.rand_layout_minor:
query_pos_limit = min(self.JITTER_POS_BOUND, query_pos_limit)
geom.pm_shift_bodies(self._space,
query_block.bodies,
position=sensor.bodies[0].position)
geom.pm_randomise_pose(self._space,
query_block.bodies,
self.ARENA_BOUNDS_LRBT,
ignore_shapes=sensor.shapes,
rng=self.rng,
rand_pos=True,
rand_rot=True,
rel_pos_linf_limit=query_pos_limit,
rel_rot_limit=rot_limit)
# block lookup index
self.__block_index = en.EntityIndex([query_block, *outside_blocks])
def on_reset(self):
# make the robot at default position (will be randomised at end if
# rand_layout is true)
robot = self._make_robot(*self.DEFAULT_ROBOT_POSE)
# 3x blue & 2x of each other colour
default_colours = self.DEFAULT_BLOCK_COLOURS
# 3x pentagon & 2x of each other shape type
default_shape_types = self.DEFAULT_BLOCK_SHAPES
# these were generated by randomly scattering shapes about the chosen
# default robot position and then rounding down values a bit
default_poses = self.DEFAULT_BLOCK_POSES
default_n_shapes = len(default_colours)
if self.rand_shape_count:
n_shapes = self.rng.randint(7, 10 + 1)
# rand_shape_count=True implies rand_layout=True, so these MUST be
# randomised at the end
poses = [((0, 0), 0)] * n_shapes
else:
n_shapes = default_n_shapes
# if rand_layout=True, these will be randomised at the end
poses = default_poses
if self.rand_shape_colour:
# make sure we have at least one of each colour
colours = list(en.SHAPE_COLOURS)
colours.extend([
self.rng.choice(en.SHAPE_COLOURS)
for _ in range(n_shapes - len(colours))
])
self.rng.shuffle(colours)
else:
colours = default_colours
if self.rand_shape_type:
# make sure we have at least one of each type, too
shape_types = list(en.SHAPE_TYPES)
shape_types.extend([
self.rng.choice(en.SHAPE_TYPES)
for _ in range(n_shapes - len(shape_types))
])
self.rng.shuffle(shape_types)
else:
shape_types = default_shape_types
assert len(poses) == n_shapes
assert len(colours) == n_shapes
assert len(shape_types) == n_shapes
shape_ents = []
for ((x, y), angle), colour, shape_type \
in zip(poses, colours, shape_types):
shape = self._make_shape(shape_type=shape_type,
colour_name=colour,
init_pos=(x, y),
init_angle=angle)
shape_ents.append(shape)
self.add_entities(shape_ents)
# make index mapping characteristic values to blocks
if self.cluster_by == self.ClusterBy.COLOUR:
c_values_list = np.asarray(colours, dtype='object')
self.__characteristic_values = np.unique(c_values_list)
elif self.cluster_by == self.ClusterBy.TYPE:
c_values_list = np.asarray(shape_types, dtype='object')
self.__characteristic_values = np.unique(c_values_list)
else:
raise NotImplementedError(
f"don't know how to cluster by '{self.cluster_by}'")
self.__blocks_by_characteristic = {}
assert len(c_values_list) == len(shape_ents)
for shape, c_value in zip(shape_ents, c_values_list):
c_list = self.__blocks_by_characteristic.setdefault(c_value, [])
c_list.append(shape)
# as in match_regions.py, this should be added after all shapes so it
# appears on top, but before layout randomisation so that it gets added
# to the space correctly
self.add_entities([robot])
if self.rand_layout_full or self.rand_layout_minor:
if self.rand_layout_full:
pos_limit = rot_limit = None
else:
pos_limit = self.JITTER_POS_BOUND
rot_limit = self.JITTER_ROT_BOUND
geom.pm_randomise_all_poses(space=self._space,
entities=[robot, *shape_ents],
arena_lrbt=self.ARENA_BOUNDS_LRBT,
rng=self.rng,
rand_pos=True,
rand_rot=True,
rel_pos_linf_limits=pos_limit,
rel_rot_limits=rot_limit)
# set up index for lookups
self.__ent_index = en.EntityIndex(shape_ents)