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]))