def testMixOfNonRecurrentAndRecurrent(self):
sequential = sequential_lib.Sequential(
[
tf.keras.layers.Dense(2),
tf.keras.layers.LSTM(
2, return_state=True, return_sequences=True),
tf.keras.layers.RNN(
tf.keras.layers.StackedRNNCells([
tf.keras.layers.LSTMCell(1),
tf.keras.layers.LSTMCell(32),
], ),
return_state=True,
return_sequences=True,
),
# Convert inner dimension to [4, 4, 2] for convolution.
inner_reshape.InnerReshape([32], [4, 4, 2]),
tf.keras.layers.Conv2D(2, 3),
# Convert 3 inner dimensions to [?] for RNN.
inner_reshape.InnerReshape([None] * 3, [-1]),
tf.keras.layers.GRU(
2, return_state=True, return_sequences=True),
dynamic_unroll_layer.DynamicUnroll(
tf.keras.layers.LSTMCell(2)),
],
input_spec=tf.TensorSpec((3, ), tf.float32))
self.assertEqual(sequential.input_tensor_spec,
tf.TensorSpec((3, ), tf.float32))
output_spec = sequential.create_variables()
self.assertEqual(output_spec, tf.TensorSpec((2, ), dtype=tf.float32))
tf.nest.map_structure(
self.assertEqual,
sequential.state_spec,
(
[ # LSTM
tf.TensorSpec((2, ), tf.float32),
tf.TensorSpec((2, ), tf.float32),
],
( # RNN(StackedRNNCells)
[
tf.TensorSpec((1, ), tf.float32),
tf.TensorSpec((1, ), tf.float32),
],
[
tf.TensorSpec((32, ), tf.float32),
tf.TensorSpec((32, ), tf.float32),
],
),
# GRU
tf.TensorSpec((2, ), tf.float32),
[ # DynamicUnroll
tf.TensorSpec((2, ), tf.float32),
tf.TensorSpec((2, ), tf.float32),
]))
inputs = tf.ones((8, 10, 3), dtype=tf.float32)
outputs, _ = sequential(inputs)
self.assertEqual(outputs.shape, tf.TensorShape([8, 10, 2]))
def testIncompatibleShapes(self):
with self.assertRaisesRegex(ValueError, 'must have known rank'):
inner_reshape.InnerReshape(tf.TensorShape(None), [1])
with self.assertRaisesRegex(ValueError,
'Mismatched number of elements'):
inner_reshape.InnerReshape([1, 2], [])
with self.assertRaisesRegex(ValueError, r'Shapes.*are incompatible'):
inner_reshape.InnerReshape([1], [1, 1])(np.ones((2, 3)))
def testLearnerRaiseExceptionOnMismatchingBatchSetup(self):
obs_spec = tensor_spec.TensorSpec([2], tf.float32)
time_step_spec = ts.time_step_spec(obs_spec)
action_spec = tensor_spec.BoundedTensorSpec([], tf.int32, 0, 1)
flat_action_spec = tf.nest.flatten(action_spec)[0]
num_actions = flat_action_spec.maximum - flat_action_spec.minimum + 1
network = sequential.Sequential([
tf.keras.layers.Dense(num_actions, dtype=tf.float32),
inner_reshape.InnerReshape([None], [num_actions])
])
agent = behavioral_cloning_agent.BehavioralCloningAgent(
time_step_spec, action_spec, cloning_network=network, optimizer=None)
with self.assertRaisesRegex(
RuntimeError,
(r'The slot variable initialization failed. The learner assumes all '
r'experience tensors required an `outer_rank = \(None, '
r'agent.train_sequence_length\)`\. If that\'s not the case for your '
r'agent try setting `run_optimizer_variable_init=False`\.')):
learner.Learner(
root_dir=os.path.join(self.create_tempdir().full_path, 'learner'),
train_step=train_utils.create_train_step(),
agent=agent)
def create_sequential_critic_network(obs_fc_layer_units, action_fc_layer_units,
joint_fc_layer_units):
"""Create a sequential critic network."""
# Split the inputs into observations and actions.
def split_inputs(inputs):
return {'observation': inputs[0], 'action': inputs[1]}
# Create an observation network.
obs_network = (create_fc_network(obs_fc_layer_units)
if obs_fc_layer_units else create_identity_layer())
# Create an action network.
action_network = (create_fc_network(action_fc_layer_units)
if action_fc_layer_units else create_identity_layer())
# Create a joint network.
joint_network = (create_fc_network(joint_fc_layer_units)
if joint_fc_layer_units else create_identity_layer())
# Final layer.
value_layer = tf.keras.layers.Dense(1, kernel_initializer='glorot_uniform')
return sequential.Sequential([
tf.keras.layers.Lambda(split_inputs),
nest_map.NestMap({
'observation': obs_network,
'action': action_network
}),
nest_map.NestFlatten(),
tf.keras.layers.Concatenate(), joint_network, value_layer,
inner_reshape.InnerReshape(current_shape=[1], new_shape=[])
],
name='sequential_critic')
def testLearnerRaiseExceptionOnMismatchingBatchSetup(self):
obs_spec = tensor_spec.TensorSpec([2], tf.float32)
time_step_spec = ts.time_step_spec(obs_spec)
action_spec = tensor_spec.BoundedTensorSpec([], tf.int32, 0, 1)
flat_action_spec = tf.nest.flatten(action_spec)[0]
num_actions = flat_action_spec.maximum - flat_action_spec.minimum + 1
network = sequential.Sequential([
tf.keras.layers.Dense(num_actions, dtype=tf.float32),
inner_reshape.InnerReshape([None], [num_actions])
])
agent = behavioral_cloning_agent.BehavioralCloningAgent(
time_step_spec,
action_spec,
cloning_network=network,
optimizer=None)
with self.assertRaisesRegex(
ValueError,
'All of the Tensors in `value` must have one outer dimension.'
):
learner.Learner(root_dir=os.path.join(
self.create_tempdir().full_path, 'learner'),
train_step=train_utils.create_train_step(),
agent=agent)
def create_critic_network(obs_fc_layer_units, action_fc_layer_units,
joint_fc_layer_units):
"""Create a critic network for DDPG."""
def split_inputs(inputs):
return {'observation': inputs[0], 'action': inputs[1]}
obs_network = create_fc_network(
obs_fc_layer_units) if obs_fc_layer_units else create_identity_layer()
action_network = create_fc_network(
action_fc_layer_units
) if action_fc_layer_units else create_identity_layer()
joint_network = create_fc_network(
joint_fc_layer_units
) if joint_fc_layer_units else create_identity_layer()
value_fc_layer = tf.keras.layers.Dense(
1,
activation=None,
kernel_initializer=tf.keras.initializers.RandomUniform(minval=-0.003,
maxval=0.003))
return sequential.Sequential([
tf.keras.layers.Lambda(split_inputs),
nest_map.NestMap({
'observation': obs_network,
'action': action_network
}),
nest_map.NestFlatten(),
tf.keras.layers.Concatenate(), joint_network, value_fc_layer,
inner_reshape.InnerReshape([1], [])
])
def testInnerReshapeUnknowns(self):
layer = inner_reshape.InnerReshape([None, None], [-1])
out = layer(np.arange(3 * 20).reshape(3, 4, 5))
self.assertAllEqual(self.evaluate(out),
np.arange(3 * 20).reshape(3, 20))
out = layer(np.arange(6 * 20).reshape(2, 3, 4, 5))
self.assertAllEqual(self.evaluate(out),
np.arange(6 * 20).reshape(2, 3, 20))
def testInnerReshapeSimple(self):
layer = inner_reshape.InnerReshape([3, 4], [12])
out = layer(np.arange(2 * 12).reshape(2, 3, 4))
self.assertAllEqual(self.evaluate(out),
np.arange(2 * 12).reshape(2, 12))
out = layer(np.arange(4 * 12).reshape(2, 2, 3, 4))
self.assertAllEqual(self.evaluate(out),
np.arange(4 * 12).reshape(2, 2, 12))
def testCreateAndCall(self):
net = sequential.Sequential([
nest_map.NestMap(
{'inp1': tf.keras.layers.Dense(8),
'inp2': sequential.Sequential([
tf.keras.layers.Conv2D(2, 3),
# Convert 3 inner dimensions to [8] for RNN.
inner_reshape.InnerReshape([None] * 3, [8]),
]),
'inp3': tf.keras.layers.LSTM(
8, return_state=True, return_sequences=True)}),
nest_map.NestFlatten(),
tf.keras.layers.Add()])
self.assertEqual(
net.state_spec,
({
'inp1': (),
'inp2': (),
'inp3': (2 * [tf.TensorSpec(shape=(8,), dtype=tf.float32)],),
},))
output_spec = net.create_variables(
{
'inp1': tf.TensorSpec(shape=(3,), dtype=tf.float32),
'inp2': tf.TensorSpec(shape=(4, 4, 2,), dtype=tf.float32),
'inp3': tf.TensorSpec(shape=(3,), dtype=tf.float32),
})
self.assertEqual(output_spec, tf.TensorSpec(shape=(8,), dtype=tf.float32))
inputs = {
'inp1': tf.ones((8, 10, 3), dtype=tf.float32),
'inp2': tf.ones((8, 10, 4, 4, 2), dtype=tf.float32),
'inp3': tf.ones((8, 10, 3), dtype=tf.float32)
}
output, next_state = net(inputs)
self.assertEqual(output.shape, tf.TensorShape([8, 10, 8]))
self.assertEqual(
tf.nest.map_structure(lambda t: t.shape, next_state),
({
'inp1': (),
'inp2': (),
'inp3': (2 * [tf.TensorShape([8, 8])],),
},))
# Test passing in a state.
output, next_state = net(inputs, next_state)
self.assertEqual(output.shape, tf.TensorShape([8, 10, 8]))
self.assertEqual(
tf.nest.map_structure(lambda t: t.shape, next_state),
({
'inp1': (),
'inp2': (),
'inp3': (2 * [tf.TensorShape([8, 8])],),
},))
def get_dummy_net(action_spec, observation_spec=None):
flat_action_spec = tf.nest.flatten(action_spec)[0]
if flat_action_spec.dtype.is_integer:
# Emitting discrete actions.
num_actions = flat_action_spec.maximum - flat_action_spec.minimum + 1
kernel_initializer = tf.constant_initializer([[2, 1], [1, 1]])
bias_initializer = tf.constant_initializer([[1], [1]])
final_shape = [num_actions]
else:
# Emitting continuous vectors.
num_actions = np.prod(flat_action_spec.shape)
kernel_initializer = None
bias_initializer = None
final_shape = flat_action_spec.shape
return sequential.Sequential([
tf.keras.layers.Dense(num_actions,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
dtype=tf.float32),
inner_reshape.InnerReshape([None], final_shape)
],
input_spec=observation_spec)
def testMixOfNonRecurrentAndRecurrent(self):
sequential = sequential_lib.Sequential(
[
tf.keras.layers.Dense(2),
tf.keras.layers.LSTM(
2, return_state=True, return_sequences=True),
tf.keras.layers.RNN(
tf.keras.layers.StackedRNNCells([
tf.keras.layers.LSTMCell(1),
tf.keras.layers.LSTMCell(32),
], ),
return_state=True,
return_sequences=True,
),
# Convert inner dimension to [4, 4, 2] for convolution.
inner_reshape.InnerReshape([32], [4, 4, 2]),
tf.keras.layers.Conv2D(2, 3),
# Convert 3 inner dimensions to [?] for RNN.
inner_reshape.InnerReshape([None] * 3, [-1]),
tf.keras.layers.GRU(
2, return_state=True, return_sequences=True),
dynamic_unroll_layer.DynamicUnroll(
tf.keras.layers.LSTMCell(2)),
tf.keras.layers.Lambda(
lambda x: tfd.MultivariateNormalDiag(loc=x, scale_diag=x)),
],
input_spec=tf.TensorSpec((3, ), tf.float32)) # pytype: disable=wrong-arg-types
self.assertEqual(sequential.input_tensor_spec,
tf.TensorSpec((3, ), tf.float32))
output_spec = sequential.create_variables()
self.assertIsInstance(output_spec,
distribution_utils.DistributionSpecV2)
output_event_spec = output_spec.event_spec
self.assertEqual(output_event_spec,
tf.TensorSpec((2, ), dtype=tf.float32))
tf.nest.map_structure(
self.assertEqual,
sequential.state_spec,
(
[ # LSTM
tf.TensorSpec((2, ), tf.float32),
tf.TensorSpec((2, ), tf.float32),
],
( # RNN(StackedRNNCells)
[
tf.TensorSpec((1, ), tf.float32),
tf.TensorSpec((1, ), tf.float32),
],
[
tf.TensorSpec((32, ), tf.float32),
tf.TensorSpec((32, ), tf.float32),
],
),
# GRU
tf.TensorSpec((2, ), tf.float32),
[ # DynamicUnroll
tf.TensorSpec((2, ), tf.float32),
tf.TensorSpec((2, ), tf.float32),
]))
inputs = tf.ones((8, 10, 3), dtype=tf.float32)
dist, _ = sequential(inputs)
outputs = dist.sample()
self.assertEqual(outputs.shape, tf.TensorShape([8, 10, 2]))
