code: 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"))
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  from absl import flags
  import tensorflow.compat.v2
  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
  from tf_agents.utils import common
  from tf_agents.utils import nest_utils
  from tf_agents.utils import test_utils
label_desc: Use image node_id 6 for calling the NestMapTest obj's underlying member method code with example usage: obj.testPolicySaverCompatibility() without return types