code: # coding=utf-8
  # Copyright 2020 The TF-Agents Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     https://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """Tests for tf_agents.networks.nest_map."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import os
  
  from absl import flags
  import tensorflow.compat.v2 as tf
  
  from tf_agents.keras_layers import inner_reshape
  from tf_agents.networks import nest_map
  from tf_agents.networks import sequential
  from tf_agents.policies import policy_saver
  from tf_agents.policies import tf_policy
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step as ts
  from tf_agents.utils import common
  from tf_agents.utils import nest_utils
  from tf_agents.utils import test_utils
  
  FLAGS = flags.FLAGS
  
  
  class MyPolicy(tf_policy.TFPolicy):
  
    def __init__(self, time_step_spec, net):
      super(MyPolicy, self).__init__(
          time_step_spec,
          action_spec=tf.TensorSpec((None,), tf.float32))
      self._net = net
  
    def _action(self, time_step, policy_state=(), seed=None):
      out, _ = self._net(time_step.observation)
      out = tf.math.add(*tf.nest.flatten(out))
      return policy_step.PolicyStep(out, (), ())
  
  
  class NestFlattenTest(test_utils.TestCase):
  
    def testNestFlatten(self):
      layer = nest_map.NestFlatten()
      outputs = layer({'a': 1, 'b': 2})
      self.assertEqual(self.evaluate(outputs), [1, 2])
  
  
  class NestMapTest(test_utils.TestCase):
  
    def setUp(self):
      if not common.has_eager_been_enabled():
        self.skipTest('Only supported in TF2.x.')
      super(NestMapTest, self).setUp()
  
    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 testNestedNest(self):
      # layer structure: {'a': {'b': .}}
      net = nest_map.NestMap(
          {'a': nest_map.NestMap(
              {'b': tf.keras.layers.Dense(8)})})
      net.create_variables({'a': {'b': tf.TensorSpec((1,), dtype=tf.float32)}})
  
    def testNestedNestWithNestedState(self):
      # layer structure: (., {'a': {'b': .}})
      net = nest_map.NestMap(
          (tf.keras.layers.Dense(7),
           {'a': nest_map.NestMap(
               {'b': tf.keras.layers.LSTM(
                   8, return_state=True, return_sequences=True)})}))
      # TODO(b/177337002): remove the forced tuple wrapping the LSTM
      # state once we make a generic KerasWrapper network and clean up
      # Sequential and NestMap to use that instead of singleton Sequential.
      out, state = net(
          (tf.ones((1, 2)), {'a': {'b': tf.ones((1, 2))}}),
          network_state=((), {'a': {'b': ((tf.ones((1, 8)), tf.ones((1, 8))),)}}))
      nest_utils.assert_matching_dtypes_and_inner_shapes(
          out,
          (
              tf.TensorSpec(dtype=tf.float32, shape=(7,)),
              {'a': {'b': tf.TensorSpec(dtype=tf.float32, shape=(8,))}}
          ),
          caller=self, tensors_name='out', specs_name='out_expected')
      nest_utils.assert_matching_dtypes_and_inner_shapes(
          state,
          (
              (),
              {'a': {'b': ((tf.TensorSpec(dtype=tf.float32, shape=(8,)),
                            tf.TensorSpec(dtype=tf.float32, shape=(8,))),)}}
          ),
          caller=self, tensors_name='state', specs_name='state_expected')
  
    def testIncompatibleStructureInputs(self):
      with self.assertRaisesRegex(
          TypeError,
          r'`nested_layers` and `input_spec` do not have matching structures'):
        nest_map.NestMap(
            [tf.keras.layers.Dense(8)],
            input_spec={'ick': tf.TensorSpec(8, tf.float32)})
  
      with self.assertRaisesRegex(
          TypeError,
          r'`self.nested_layers` and `inputs` do not have matching structures'):
        net = nest_map.NestMap([tf.keras.layers.Dense(8)])
        net.create_variables({'ick': tf.TensorSpec((1,), dtype=tf.float32)})
  
      with self.assertRaisesRegex(
          TypeError,
          r'`self.nested_layers` and `inputs` do not have matching structures'):
        net = nest_map.NestMap([tf.keras.layers.Dense(8)])
        net({'ick': tf.constant([[1.0]])})
  
      with self.assertRaisesRegex(
          ValueError,
          r'`network_state` and `state_spec` do not have matching structures'):
        net = nest_map.NestMap(
            tf.keras.layers.LSTM(8, return_state=True, return_sequences=True))
        net(tf.ones((1, 2)), network_state=(tf.ones((1, 1)), ()))
  
    def testPolicySaverCompatibility(self):
      observation_spec = {
          'a': tf.TensorSpec(4, tf.float32),
          'b': tf.TensorSpec(3, tf.float32)
      }
      time_step_tensor_spec = ts.time_step_spec(observation_spec)
      net = nest_map.NestMap(
          {'a': tf.keras.layers.LSTM(8, return_state=True, return_sequences=True),
           'b': tf.keras.layers.Dense(8)})
      net.create_variables(observation_spec)
      policy = MyPolicy(time_step_tensor_spec, net)
  
      sample = tensor_spec.sample_spec_nest(
          time_step_tensor_spec, outer_dims=(5,))
  
      step = policy.action(sample)
      self.assertEqual(step.action.shape.as_list(), [5, 8])
  
      train_step = common.create_variable('train_step')
      saver = policy_saver.PolicySaver(policy, train_step=train_step)
      self.initialize_v1_variables()
  
      with self.cached_session():
        saver.save(os.path.join(FLAGS.test_tmpdir, 'nest_map_model'))
  
  
  if __name__ == '__main__':
    test_utils.main()
  
apis: tensorflow.compat.v2.keras.layers.LSTM tensorflow.compat.v2.keras.layers.Conv2D tensorflow.compat.v2.nest.flatten tensorflow.compat.v2.keras.layers.Dense tensorflow.compat.v2.constant tensorflow.compat.v2.keras.layers.Add tensorflow.compat.v2.TensorShape tensorflow.compat.v2.ones tensorflow.compat.v2.TensorSpec tensorflow.compat.v2.nest.map_structure