def _get_sprites(self):
"""Get list of sprites."""
sprites = [
sprite.Sprite(x=0.75,
y=0.95,
shape='spoke_6',
scale=0.2,
c0=20,
c1=50,
c2=80),
sprite.Sprite(x=0.2,
y=0.3,
shape='triangle',
scale=0.1,
c0=150,
c1=255,
c2=100),
sprite.Sprite(x=0.7,
y=0.5,
shape='square',
scale=0.3,
c0=0,
c1=255,
c2=0),
sprite.Sprite(x=0.5,
y=0.5,
shape='square',
scale=0.3,
c0=255,
c1=0,
c2=0),
]
return sprites
def testFilterDistrib(self):
sprites = [
sprite.Sprite(x=0.45, y=0.45, c0=64),
sprite.Sprite(x=0.45, y=0.55, c0=128),
sprite.Sprite(x=0.55, y=0.45, c0=192),
sprite.Sprite(x=0.4, y=0.4, c0=255),
]
filter_distribs = [
distribs.Continuous('c0', 0, 65), # selects sprites[:1]
distribs.Continuous('c0', 0, 129), # selects sprites[:2]
distribs.Continuous('c0', 0, 193), # selects sprites[:3]
distribs.Continuous('c0', 0, 256), # selects sprites[:4]
distribs.Continuous('c0', 65, 256), # selects sprites[1:4]
]
task_list = [
tasks.FindGoalPosition(filter_distrib=x,
goal_position=(0.5, 0.5),
terminate_distance=0.1)
for x in filter_distribs
]
rewards = [1.5, 2.9, 4.4, 2.3, 0.9]
successes = [True, True, True, False, False]
for t, r, s in zip(task_list, rewards, successes):
self.assertAlmostEqual(t.reward(sprites), r, delta=0.1)
self.assertEqual(t.success(sprites), s)
def testMoreSprites(self, positions_0, positions_1, reward):
sprites = [sprite.Sprite(x=p[0], y=p[1], c0=75) for p in positions_0]
sprites.extend(
[sprite.Sprite(x=p[0], y=p[1], c0=225) for p in positions_1])
cluster_distribs = [
distribs.Continuous('c0', 50, 100),
distribs.Continuous('c0', 200, 250),
]
task = tasks.Clustering(cluster_distribs=cluster_distribs)
self.assertAlmostEqual(task.reward(sprites), reward, delta=0.1)
def setUp(self):
super(ClusteringTest, self).setUp()
self.sprites = [
sprite.Sprite(x=0.2, y=0.2, c0=64),
sprite.Sprite(x=0.3, y=0.3, c0=128),
sprite.Sprite(x=0.8, y=0.9, c0=192),
sprite.Sprite(x=0.9, y=0.8, c0=255),
]
self.cluster_distribs = [
distribs.Continuous('c0', 0, 129),
distribs.Continuous('c0', 190, 256),
]
def test4Sprites(self, positions, reward, success):
sprites = [
sprite.Sprite(x=positions[0][0], y=positions[0][1], c0=64),
sprite.Sprite(x=positions[1][0], y=positions[1][1], c0=128),
sprite.Sprite(x=positions[2][0], y=positions[2][1], c0=192),
sprite.Sprite(x=positions[3][0], y=positions[3][1], c0=255),
]
cluster_distribs = [
distribs.Continuous('c0', 0, 129),
distribs.Continuous('c0', 190, 256),
]
task = tasks.Clustering(cluster_distribs=cluster_distribs)
self.assertAlmostEqual(task.reward(sprites), reward, delta=0.1)
self.assertEqual(task.success(sprites), success)
def testContainsPoint(self, x, y, shape, angle, scale, containment):
linspace = np.linspace(0.2, 0.8, 4)
grid = np.stack(np.meshgrid(linspace, linspace), axis=-1)
s = sprite.Sprite(x=x, y=y, shape=shape, angle=angle, scale=scale)
eval_containment = [[s.contains_point(p) for p in row] for row in grid] # pylint: disable=g-complex-comprehension
self.assertSequenceEqual(list(np.ravel(eval_containment)),
list(np.ravel(containment)))
def __getitem__(self, index):
sampled_latents = self.factor_dist.sample()
idx = np.random.choice(len(self.csv_dict['id']), p=None)
sprites = []
latents = []
for i in range(self.sequence_len):
curr_latents = sampled_latents.copy()
csv_vals = [
self.csv_dict['x'][idx][i], self.csv_dict['y'][idx][i],
self.csv_dict['area'][idx][i]
]
curr_latents['x'] = csv_vals[0] / 64
curr_latents['y'] = csv_vals[1] / 64
curr_latents['scale'] = np.sqrt(csv_vals[2] / (64**2))
sprites.append(sprite.Sprite(**curr_latents))
latents.append(curr_latents)
first_sample = np.transpose(
self.renderer.render(sprites=[sprites[0]]).astype(np.float32) /
255., (2, 0, 1))
second_sample = np.transpose(
self.renderer.render(sprites=[sprites[1]]).astype(np.float32) /
255., (2, 0, 1))
latents1 = np.array(
[self.convert_cat(item) for item in latents[0].values()])
latents2 = np.array(
[self.convert_cat(item) for item in latents[1].values()])
return first_sample, second_sample, latents1, latents2
def make_object_under_test(self, renderer):
self.environment = environment.Environment(
task=tasks.NoReward(),
action_space=action_spaces.SelectMove(),
renderers={'obs': renderer},
init_sprites=lambda: [sprite.Sprite(c0=255)],
max_episode_length=7)
def testNoFilteredSprites(self):
sprites = [sprite.Sprite(x=0.45, y=0.45, c0=255)]
filter_distrib = distribs.Continuous('c0', 0, 254)
r = tasks.FindGoalPosition(filter_distrib=filter_distrib,
goal_position=(0.5, 0.5),
terminate_distance=0.1).reward(sprites)
self.assertTrue(np.isnan(r))
def testRenderOne(self):
sprite = sprite_lib.Sprite(
x=0.1, y=0.3, shape='square', scale=0.5, c0=0, c1=0, c2=255)
renderer = handcrafted.SpritePassthrough()
observation = renderer.render(sprites=[sprite])
self.assertEqual(observation, [sprite])
def testSequence(self):
sprites = [
sprite_lib.Sprite(x=np.random.rand(), y=np.random.rand())
for _ in range(5)
]
renderer = handcrafted.SpriteFactors()
renderer.render(sprites=sprites)
def testMoveSprites(self):
"""Take a series of actions and repeatedly check sprite motions."""
action_space = action_spaces.SelectMove(scale=0.5)
sprites = [sprite.Sprite(x=0.55, y=0.5), sprite.Sprite(x=0.5, y=0.5)]
# Move second (top) sprite
action_space.step(
np.array([0.52, 0.52, 0.5, 0.48]), sprites, keep_in_frame=False)
self.assertTrue(np.allclose(sprites[0].position, [0.55, 0.5], atol=1e-5))
self.assertTrue(np.allclose(sprites[1].position, [0.5, 0.49], atol=1e-5))
# Move first (bottom) sprite
action_space.step(
np.array([0.58, 0.5, 0.9, 0.9]), sprites, keep_in_frame=False)
self.assertTrue(np.allclose(sprites[0].position, [0.75, 0.7], atol=1e-5))
self.assertTrue(np.allclose(sprites[1].position, [0.5, 0.49], atol=1e-5))
# Move neither sprite
action_space.step(
np.array([0.58, 0.5, 0.9, 0.9]), sprites, keep_in_frame=False)
self.assertTrue(np.allclose(sprites[0].position, [0.75, 0.7], atol=1e-5))
self.assertTrue(np.allclose(sprites[1].position, [0.5, 0.49], atol=1e-5))
# Move second (top) sprite
action_space.step(
np.array([0.5, 0.5, 0.2, 0.5]), sprites, keep_in_frame=False)
self.assertTrue(np.allclose(sprites[0].position, [0.75, 0.7], atol=1e-5))
self.assertTrue(np.allclose(sprites[1].position, [0.35, 0.49], atol=1e-5))
# Move first (bottom) sprite
action_space.step(
np.array([0.78, 0.74, 0.9, 0.9]), sprites, keep_in_frame=False)
self.assertTrue(np.allclose(sprites[0].position, [0.95, 0.9], atol=1e-5))
self.assertTrue(np.allclose(sprites[1].position, [0.35, 0.49], atol=1e-5))
# Move first (bottom) sprite
action_space.step(
np.array([0.92, 0.9, 0.9, 0.5]), sprites, keep_in_frame=True)
self.assertTrue(np.allclose(sprites[0].position, [1., 0.9], atol=1e-5))
self.assertTrue(np.allclose(sprites[1].position, [0.35, 0.49], atol=1e-5))
# Move first (bottom) sprite
action_space.step(
np.array([0.98, 0.9, 0.7, 0.9]), sprites, keep_in_frame=False)
self.assertTrue(np.allclose(sprites[0].position, [1.1, 1.1], atol=1e-5))
self.assertTrue(np.allclose(sprites[1].position, [0.35, 0.49], atol=1e-5))
def testColorToRGB(self):
s = sprite.Sprite(x=0.5, y=0.5, shape='square', c0=0.2, c1=0.5, c2=0.5)
def _color_to_rgb(c):
return tuple((255 * np.array(colorsys.hsv_to_rgb(*c))).astype(np.uint8))
renderer = pil_renderer.PILRenderer(
image_size=(64, 64), color_to_rgb=_color_to_rgb)
image = renderer.render([s])
self.assertSequenceEqual(list(image[32, 32]), [114, 127, 63])
def _generate():
n = num_sprites() if callable(num_sprites) else num_sprites
sprites = [sprite.Sprite(**factor_dist.sample()) for _ in range(n)]
if fix_colors:
colors = [[0.9,0.6,0.95], [0.55,0.6,0.95], [0.27,0.6,0.95]]
np.random.shuffle(colors)
for idx,s in enumerate(sprites):
s.color = colors[idx]
return sprites
def make_object_under_test(self):
"""Environment creator used by test_utils.EnvironmentTestMixin."""
env = environment.Environment(
task=tasks.NoReward(),
action_space=action_spaces.SelectMove(),
renderers={},
init_sprites=lambda: [sprite.Sprite(c0=255)],
max_episode_length=7)
return env
def testAttributesSingleton(self, x, y, shape, c0, c1, c2, scale, angle):
sprite = sprite_lib.Sprite(
x=x, y=y, shape=shape, c0=c0, c1=c1, c2=c2, scale=scale, angle=angle)
renderer = handcrafted.SpriteFactors()
outputs = renderer.render(sprites=[sprite])[0]
self.assertEqual(outputs['shape'], const.ShapeType[shape].value)
for (name, value) in (('x', x), ('y', y), ('c0', c0), ('c1', c1),
('c2', c2), ('scale', scale), ('angle', angle)):
self.assertAlmostEqual(outputs[name], value, delta=1e-4)
def testBasicInitialization(self):
sprite.Sprite(x=0.2,
y=0.8,
shape='triangle',
angle=45,
scale=0.3,
c0=200,
c1=150,
c2=100,
x_vel=-0.2,
y_vel=0.1)
def sample_observations_from_factors(self, factors, random_state):
#Sample a batch of observations X given a batch of factors Y.
images = []
for f in factors:
if self.natural_discrete:
f_convert = []
for i in self.latent_factor_indices:
f_convert.append(
list(self.lab_encs[list(
self.dsprites.keys())[i]].inverse_transform(
[f[i]]))[0])
rendering = self.renderer.render(sprites=[
sprite.Sprite(**self.back_to_dict(f_convert, False))
])
else:
rendering = self.renderer.render(
sprites=[sprite.Sprite(**self.back_to_dict(f, True))])
images.append(
np.transpose(rendering.astype(np.float32) / 255., (2, 0, 1)))
return np.array(images)
def testObservationSpec(self, num_sprites, factors):
sprites = [sprite_lib.Sprite() for _ in range(num_sprites)]
renderer = handcrafted.SpriteFactors(factors=factors)
renderer.render(sprites=sprites)
obs_spec = renderer.observation_spec()
for v in obs_spec[0].values():
self.assertEqual(v.shape, ())
obs_spec_keys = [set(x) for x in obs_spec]
self.assertSequenceEqual(obs_spec_keys, num_sprites * [set(factors)])
def testResetAngle(self):
init_vertices = [[0.625, 0.625], [0.375, 0.625], [0.375, 0.375],
[0.625, 0.375]]
s = sprite.Sprite(angle=0, scale=0.25, shape='square')
self.assertSequenceAlmostEqual(
np.ravel(s.vertices), np.ravel(init_vertices), delta=1e-3)
s.angle = -45
rotated_vertices = [[0.677, 0.5], [0.5, 0.677], [0.323, 0.5], [0.5, 0.323]]
self.assertSequenceAlmostEqual(
np.ravel(s.vertices), np.ravel(rotated_vertices), delta=1e-3)
def testContainsPoint(self, x, y, shape, angle, scale, containment):
# As we use plots to prepare these tests, it's easier to write the matrix
# "in the wrong orientation" (i.e. with origin='lower') and flip it.
containment = np.flipud(containment)
linspace = np.linspace(0.1, 0.9, 4)
grid = np.stack(np.meshgrid(linspace, linspace), axis=-1)
s = sprite.Sprite(x=x, y=y, shape=shape, angle=angle, scale=scale)
eval_containment = np.array(
[[s.contains_point(p) for p in row] for row in grid])
self.assertTrue(np.allclose(eval_containment, containment))
def testResetScale(self):
s = sprite.Sprite(scale=0.25, shape='square')
init_vertices = [[0.625, 0.625], [0.375, 0.625], [0.375, 0.375],
[0.625, 0.375]]
self.assertSequenceAlmostEqual(
np.ravel(s.vertices), np.ravel(init_vertices), delta=1e-3)
s.scale = 0.5
scaled_vertices = [[0.531, 0.531], [0.469, 0.531], [0.469, 0.469],
[0.531, 0.469]]
self.assertSequenceAlmostEqual(
np.ravel(s.vertices), np.ravel(scaled_vertices), delta=1e-3)
def testResetShape(self):
s = sprite.Sprite(scale=0.25, shape='square')
square_vertices = [[0.625, 0.625], [0.375, 0.625], [0.375, 0.375],
[0.625, 0.375]]
self.assertSequenceAlmostEqual(
np.ravel(s.vertices), np.ravel(square_vertices), delta=1e-3)
s.shape = 'triangle'
triangle_vertices = [[0.5, 0.72], [0.31, 0.39], [0.69, 0.39]]
self.assertSequenceAlmostEqual(
np.ravel(s.vertices), np.ravel(triangle_vertices), delta=1e-2)
def _generate():
mean = [0, 0, 0, 0]
cov = [[1., 0.7, 0.7, 0.7], [0.7, 1., 0.7, 0.7], [0.7, 0.7, 1., 0.7],
[0.7, 0.7, 0.7, 1.]]
gaussian = np.random.multivariate_normal(mean, cov, 1)
uniform = norm.cdf(gaussian)
n = num_sprites() if callable(num_sprites) else num_sprites
sprites = [
sprite.Sprite(**factor_dist.sample(uniform)) for _ in range(n)
]
return sprites
def testKSprites(self, num_sprites):
sprites = [
sprite_lib.Sprite(x=np.random.rand(), y=np.random.rand())
for _ in range(num_sprites)
]
renderer = handcrafted.SpritePassthrough()
observation = renderer.render(sprites=sprites)
self.assertSequenceEqual(observation, sprites)
obs_spec = renderer.observation_spec()
self.assertTrue(obs_spec.shape, (num_sprites,))
def setUp(self):
super(MetaAggregatedTest, self).setUp()
self.subtasks = [
tasks.FindGoalPosition(filter_distrib=distribs.Continuous(
'c0', 0, 100),
goal_position=np.array([0.2, 0.2]),
terminate_distance=0.1),
tasks.FindGoalPosition(filter_distrib=distribs.Continuous(
'c0', 100, 200),
goal_position=np.array([0.5, 0.5]),
terminate_distance=0.1,
terminate_bonus=5.0),
tasks.FindGoalPosition(filter_distrib=distribs.Continuous(
'c0', 200, 256),
goal_position=np.array([0.8, 0.8]),
terminate_distance=0.1,
terminate_bonus=10.0),
]
self.success_sprites = [
sprite.Sprite(x=0.2, y=0.2, c0=50),
sprite.Sprite(x=0.5, y=0.45, c0=150),
sprite.Sprite(x=0.85, y=0.75, c0=250),
]
self.success_rewards = [5., 7.5, 11.5]
self.failure_sprites = [
sprite.Sprite(x=0.2, y=0.8, c0=50),
sprite.Sprite(x=0.3, y=0.45, c0=150),
sprite.Sprite(x=0.9, y=0.75, c0=250),
]
self.failure_rewards = [-25., -5.3, -0.6]
def testResetShape(self):
s = sprite.Sprite(angle=30, scale=0.25, shape='square')
square_vertices = [[0.653, 0.588], [0.412, 0.653], [0.347, 0.412],
[0.588, 0.347]]
self.assertSequenceAlmostEqual(np.ravel(s.vertices),
np.ravel(square_vertices),
delta=1e-3)
s.shape = 'triangle'
triangle_vertices = [[0.69, 0.61], [0.31, 0.61], [0.5, 0.281]]
self.assertSequenceAlmostEqual(np.ravel(s.vertices),
np.ravel(triangle_vertices),
delta=1e-3)
def testMoveSprites(self,
init_positions,
action,
final_positions,
keep_in_frame=True):
"""Take a series of actions and repeatedly check sprite motions."""
action_space = action_spaces.Embodied(step_size=0.1)
sprites = [
sprite.Sprite(x=pos[0], y=pos[1], shape='square', scale=0.15)
for pos in init_positions
]
action_space.step(action, sprites, keep_in_frame=keep_in_frame)
for s, p in zip(sprites, final_positions):
self.assertTrue(np.allclose(s.position, p, atol=1e-5))
def test3Clusters(self):
sprites = [
sprite.Sprite(x=0.2, y=0.2, c0=64),
sprite.Sprite(x=0.3, y=0.3, c0=64),
sprite.Sprite(x=0.8, y=0.9, c0=128),
sprite.Sprite(x=0.9, y=0.8, c0=128),
sprite.Sprite(x=0.8, y=0.9, c0=255),
sprite.Sprite(x=0.9, y=0.8, c0=255),
]
cluster_distribs = [
distribs.Continuous('c0', 0, 100),
distribs.Continuous('c0', 100, 150),
distribs.Continuous('c0', 200, 256),
]
task = tasks.Clustering(cluster_distribs=cluster_distribs)
self.assertAlmostEqual(task.reward(sprites), 17.5, delta=0.1)
def testTaskTermination(self):
task = tasks.FindGoalPosition(goal_position=(0.5, 0.5))
action_space = action_spaces.SelectMove()
env_renderers = {}
init_sprites = lambda: [sprite.Sprite(x=0.25, y=0.25, c0=255)]
env = environment.Environment(task, action_space, env_renderers,
init_sprites)
donothing_action = np.array([0.25, 0.25, 0.5, 0.5])
success_action = np.array([0.25, 0.25, 0.75, 0.75])
timestep = env.step(donothing_action)
self.assertTrue(timestep.first())
timestep = env.step(donothing_action)
self.assertTrue(timestep.mid())
timestep = env.step(success_action)
self.assertTrue(timestep.last())
timestep = env.step(success_action)
self.assertTrue(timestep.first())
