def testTrainStepNotSaved(self): network = q_network.QNetwork( input_tensor_spec=self._time_step_spec.observation, action_spec=self._action_spec) policy = q_policy.QPolicy(time_step_spec=self._time_step_spec, action_spec=self._action_spec, q_network=network) saver = policy_saver.PolicySaver(policy, batch_size=None) path = os.path.join(self.get_temp_dir(), 'save_model') saver.save(path) reloaded = tf.compat.v2.saved_model.load(path) self.assertIn('get_train_step', reloaded.signatures) train_step_value = self.evaluate(reloaded.train_step()) self.assertEqual(-1, train_step_value)
def testCheckpointSave(self): network = q_network.QNetwork( input_tensor_spec=self._time_step_spec.observation, action_spec=self._action_spec) policy = q_policy.QPolicy(time_step_spec=self._time_step_spec, action_spec=self._action_spec, q_network=network) saver = policy_saver.PolicySaver(policy, batch_size=None) path = os.path.join(self.get_temp_dir(), 'save_model') self.evaluate(tf.compat.v1.global_variables_initializer()) saver.save(path) checkpoint_path = os.path.join(self.get_temp_dir(), 'checkpoint') saver.save_checkpoint(checkpoint_path) self.assertTrue(tf.compat.v2.io.gfile.exists(checkpoint_path))
def testActionWithinBounds(self): bounded_action_spec = tensor_spec.BoundedTensorSpec([1], tf.int32, minimum=-6, maximum=-5) policy = q_policy.QPolicy(self._time_step_spec, bounded_action_spec, q_network=DummyNet()) observations = tf.constant([[1, 2], [3, 4]], dtype=tf.float32) time_step = ts.restart(observations, batch_size=2) action_step = policy.action(time_step) self.assertEqual(action_step.action.shape.as_list(), [2, 1]) self.assertEqual(action_step.action.dtype, tf.int32) # Initialize all variables self.evaluate(tf.compat.v1.global_variables_initializer()) action = self.evaluate(action_step.action) self.assertTrue(np.all(action <= -5) and np.all(action >= -6))
def testLogits(self): tf.compat.v1.set_random_seed(1) wrapped = q_policy.QPolicy(self._time_step_spec, self._action_spec, q_network=DummyNet()) policy = boltzmann_policy.BoltzmannPolicy(wrapped, temperature=0.5) observations = tf.constant([[1, 2]], dtype=tf.float32) time_step = ts.restart(observations, batch_size=1) distribution_step = policy.distribution(time_step) logits = distribution_step.action.logits original_logits = wrapped.distribution(time_step).action.logits self.evaluate(tf.compat.v1.global_variables_initializer()) # The un-temperature'd logits would be 4 and 5.5, because it is (1 2) . (1 # 1) + 1 and (1 2) . (1.5 1.5) + 1. The temperature'd logits will be double # that. self.assertAllEqual([[[4., 5.5]]], self.evaluate(original_logits)) self.assertAllEqual([[[8., 11.]]], self.evaluate(logits))
def testSaveGetInitialState(self): if not tf.executing_eagerly(): self.skipTest( 'b/129079730: PolicySaver does not work in TF1.x yet') q_network = q_rnn_network.QRnnNetwork( input_tensor_spec=self._time_step_spec.observation, action_spec=self._action_spec) policy = q_policy.QPolicy(time_step_spec=self._time_step_spec, action_spec=self._action_spec, q_network=q_network) saver_nobatch = policy_saver.PolicySaver(policy, batch_size=None) path = os.path.join(tf.compat.v1.test.get_temp_dir(), 'save_model_initial_state_nobatch') saver_nobatch.save(path) reloaded_nobatch = tf.compat.v2.saved_model.load(path) self.assertIn('get_initial_state', reloaded_nobatch.signatures) reloaded_get_initial_state = ( reloaded_nobatch.signatures['get_initial_state']) self._compare_input_output_specs( reloaded_get_initial_state, expected_input_specs=(tf.TensorSpec(dtype=tf.int32, shape=(), name='batch_size'), ), expected_output_spec=policy.policy_state_spec, batch_input=False, batch_size=None) saver_batch = policy_saver.PolicySaver(policy, batch_size=3) path = os.path.join(tf.compat.v1.test.get_temp_dir(), 'save_model_initial_state_batch') saver_batch.save(path) reloaded_batch = tf.compat.v2.saved_model.load(path) self.assertIn('get_initial_state', reloaded_batch.signatures) reloaded_get_initial_state = reloaded_batch.signatures[ 'get_initial_state'] self._compare_input_output_specs( reloaded_get_initial_state, expected_input_specs=(), expected_output_spec=policy.policy_state_spec, batch_input=False, batch_size=3)
def _setup_as_discrete(self, time_step_spec, action_spec, loss_fn, epsilon_greedy): self._bc_loss_fn = loss_fn or self._discrete_loss if any(isinstance(d, distribution_utils.DistributionSpecV2) for d in tf.nest.flatten([self._network_output_spec])): # If the output of the cloning network contains a distribution. base_policy = actor_policy.ActorPolicy(time_step_spec, action_spec, self._cloning_network) else: # If the output of the cloning network is logits. base_policy = q_policy.QPolicy( time_step_spec, action_spec, q_network=self._cloning_network, validate_action_spec_and_network=False) policy = greedy_policy.GreedyPolicy(base_policy) collect_policy = epsilon_greedy_policy.EpsilonGreedyPolicy( base_policy, epsilon=epsilon_greedy) return policy, collect_policy
def testDeferredBatchingAction(self): # Construct policy without providing batch_size. tf_policy = q_policy.QPolicy(self._time_step_spec, self._action_spec, q_network=DummyNet(stateful=False)) policy = py_tf_policy.PyTFPolicy(tf_policy) # But time_steps have batch_size of 5 batch_size = 5 single_observation = np.array([1, 2], dtype=np.float32) time_steps = [ts.restart(single_observation)] * batch_size time_steps = fast_map_structure(lambda *arrays: np.stack(arrays), *time_steps) with self.test_session(): tf.global_variables_initializer().run() action_steps = policy.action(time_steps) self.assertEqual(action_steps.action.shape, (batch_size, )) self.assertAllEqual(action_steps.action, [1] * batch_size) self.assertAllEqual(action_steps.state, ())
def testTrainStepSaved(self): # We need to use one default session so that self.evaluate and the # SavedModel loader share the same session. with tf.compat.v1.Session().as_default(): network = q_network.QNetwork( input_tensor_spec=self._time_step_spec.observation, action_spec=self._action_spec) policy = q_policy.QPolicy(time_step_spec=self._time_step_spec, action_spec=self._action_spec, q_network=network) self.evaluate( tf.compat.v1.initializers.variables(policy.variables())) train_step = common.create_variable('train_step', initial_value=7) self.evaluate(tf.compat.v1.initializers.variables([train_step])) saver = policy_saver.PolicySaver(policy, batch_size=None, train_step=train_step) if tf.executing_eagerly(): step = saver.get_train_step() else: step = self.evaluate(saver.get_train_step()) self.assertEqual(7, step) path = os.path.join(self.get_temp_dir(), 'save_model') saver.save(path) reloaded = tf.compat.v2.saved_model.load(path) self.assertIn('get_train_step', reloaded.signatures) self.evaluate(tf.compat.v1.global_variables_initializer()) train_step_value = self.evaluate(reloaded.get_train_step()) self.assertEqual(7, train_step_value) train_step = train_step.assign_add(3) self.evaluate(train_step) saver.save(path) reloaded = tf.compat.v2.saved_model.load(path) self.evaluate(tf.compat.v1.global_variables_initializer()) train_step_value = self.evaluate(reloaded.get_train_step()) self.assertEqual(10, train_step_value)
def testUniqueSignatures(self): network = q_network.QNetwork( input_tensor_spec=self._time_step_spec.observation, action_spec=self._action_spec) policy = q_policy.QPolicy(time_step_spec=self._time_step_spec, action_spec=self._action_spec, q_network=network) saver = policy_saver.PolicySaver(policy, batch_size=None) action_signature_names = [ s.name for s in saver._signatures['action'].input_signature ] self.assertAllEqual( ['0/step_type', '0/reward', '0/discount', '0/observation'], action_signature_names) initial_state_signature_names = [ s.name for s in saver._signatures['get_initial_state'].input_signature ] self.assertAllEqual(['batch_size'], initial_state_signature_names)
def testTrainStepNotSaved(self): if not common.has_eager_been_enabled(): self.skipTest('Only supported in TF2.x. Step is required in TF1.x') network = q_network.QNetwork( input_tensor_spec=self._time_step_spec.observation, action_spec=self._action_spec) policy = q_policy.QPolicy(time_step_spec=self._time_step_spec, action_spec=self._action_spec, q_network=network) saver = policy_saver.PolicySaver(policy, batch_size=None) path = os.path.join(self.get_temp_dir(), 'save_model') saver.save(path) reloaded = tf.compat.v2.saved_model.load(path) self.assertIn('get_train_step', reloaded.signatures) train_step_value = self.evaluate(reloaded.get_train_step()) self.assertEqual(-1, train_step_value)
def testDeferredBatchingAction(self): if tf.executing_eagerly(): self.skipTest('b/123770140') # Construct policy without providing batch_size. tf_policy = q_policy.QPolicy(self._time_step_spec, self._action_spec, q_network=DummyNet(stateful=False)) policy = py_tf_policy.PyTFPolicy(tf_policy) # But time_steps have batch_size of 5 batch_size = 5 single_observation = np.array([1, 2], dtype=np.float32) time_steps = [ts.restart(single_observation)] * batch_size time_steps = fast_map_structure(lambda *arrays: np.stack(arrays), *time_steps) with self.cached_session(): self.evaluate(tf.compat.v1.global_variables_initializer()) action_steps = policy.action(time_steps) self.assertEqual(action_steps.action.shape, (batch_size, )) for a in action_steps.action: self.assertIn(a, (0, 1)) self.assertAllEqual(action_steps.state, ())
def testSaveAction(self, seeded, has_state, distribution_net, has_input_fn_and_spec): with tf.compat.v1.Graph().as_default(): tf.compat.v1.set_random_seed(self._global_seed) with tf.compat.v1.Session().as_default(): global_step = common.create_variable('train_step', initial_value=0) if distribution_net: network = actor_distribution_network.ActorDistributionNetwork( self._time_step_spec.observation, self._action_spec) policy = actor_policy.ActorPolicy( time_step_spec=self._time_step_spec, action_spec=self._action_spec, actor_network=network) else: if has_state: network = q_rnn_network.QRnnNetwork( input_tensor_spec=self._time_step_spec.observation, action_spec=self._action_spec, lstm_size=(40,)) else: network = q_network.QNetwork( input_tensor_spec=self._time_step_spec.observation, action_spec=self._action_spec) policy = q_policy.QPolicy( time_step_spec=self._time_step_spec, action_spec=self._action_spec, q_network=network) action_seed = 98723 batch_size = 3 action_inputs = tensor_spec.sample_spec_nest( (self._time_step_spec, policy.policy_state_spec), outer_dims=(batch_size,), seed=4) action_input_values = self.evaluate(action_inputs) action_input_tensors = tf.nest.map_structure(tf.convert_to_tensor, action_input_values) action_output = policy.action(*action_input_tensors, seed=action_seed) distribution_output = policy.distribution(*action_input_tensors) self.assertIsInstance( distribution_output.action, tfp.distributions.Distribution) self.evaluate(tf.compat.v1.global_variables_initializer()) action_output_dict = collections.OrderedDict( ((spec.name, value) for (spec, value) in zip( tf.nest.flatten(policy.policy_step_spec), tf.nest.flatten(action_output)))) # Check output of the flattened signature call. (action_output_value, action_output_dict) = self.evaluate( (action_output, action_output_dict)) distribution_output_value = self.evaluate(_sample_from_distributions( distribution_output)) input_fn_and_spec = None if has_input_fn_and_spec: input_fn_and_spec = (_convert_string_vector_to_action_input, tf.TensorSpec((7,), tf.string, name='example')) saver = policy_saver.PolicySaver( policy, batch_size=None, use_nest_path_signatures=False, seed=action_seed, input_fn_and_spec=input_fn_and_spec, train_step=global_step) path = os.path.join(self.get_temp_dir(), 'save_model_action') saver.save(path) with tf.compat.v1.Graph().as_default(): tf.compat.v1.set_random_seed(self._global_seed) with tf.compat.v1.Session().as_default(): reloaded = tf.compat.v2.saved_model.load(path) self.assertIn('action', reloaded.signatures) reloaded_action = reloaded.signatures['action'] if has_input_fn_and_spec: self._compare_input_output_specs( reloaded_action, expected_input_specs=input_fn_and_spec[1], expected_output_spec=policy.policy_step_spec, batch_input=True) else: self._compare_input_output_specs( reloaded_action, expected_input_specs=(self._time_step_spec, policy.policy_state_spec), expected_output_spec=policy.policy_step_spec, batch_input=True) # Reload action_input_values as tensors in the new graph. action_input_tensors = tf.nest.map_structure(tf.convert_to_tensor, action_input_values) action_input_spec = (self._time_step_spec, policy.policy_state_spec) function_action_input_dict = collections.OrderedDict( (spec.name, value) for (spec, value) in zip( tf.nest.flatten(action_input_spec), tf.nest.flatten(action_input_tensors))) # NOTE(ebrevdo): The graph-level seeds for the policy and the reloaded # model are equal, which in addition to seeding the call to action() and # PolicySaver helps ensure equality of the output of action() in both # cases. self.assertEqual(reloaded_action.graph.seed, self._global_seed) # The seed= argument for the SavedModel action call was given at # creation of the PolicySaver. if has_input_fn_and_spec: action_string_vector = _convert_action_input_to_string_vector( action_input_tensors) action_string_vector_values = self.evaluate(action_string_vector) reloaded_action_output_dict = reloaded_action(action_string_vector) reloaded_action_output = reloaded.action(action_string_vector) reloaded_distribution_output = reloaded.distribution( action_string_vector) self.assertIsInstance(reloaded_distribution_output.action, tfp.distributions.Distribution) else: # This is the flat-signature function. reloaded_action_output_dict = reloaded_action( **function_action_input_dict) # This is the non-flat function. reloaded_action_output = reloaded.action(*action_input_tensors) reloaded_distribution_output = reloaded.distribution( *action_input_tensors) self.assertIsInstance(reloaded_distribution_output.action, tfp.distributions.Distribution) if not has_state: # Try both cases: one with an empty policy_state and one with no # policy_state. Compare them. # NOTE(ebrevdo): The first call to .action() must be stored in # reloaded_action_output because this is the version being compared # later against the true action_output and the values will change # after the first call due to randomness. reloaded_action_output_no_input_state = reloaded.action( action_input_tensors[0]) reloaded_distribution_output_no_input_state = reloaded.distribution( action_input_tensors[0]) # Even with a seed, multiple calls to action will get different # values, so here we just check the signature matches. self.assertIsInstance( reloaded_distribution_output_no_input_state.action, tfp.distributions.Distribution) tf.nest.map_structure(self.match_dtype_shape, reloaded_action_output_no_input_state, reloaded_action_output) tf.nest.map_structure( self.match_dtype_shape, _sample_from_distributions( reloaded_distribution_output_no_input_state), _sample_from_distributions(reloaded_distribution_output)) self.evaluate(tf.compat.v1.global_variables_initializer()) (reloaded_action_output_dict, reloaded_action_output_value) = self.evaluate( (reloaded_action_output_dict, reloaded_action_output)) reloaded_distribution_output_value = self.evaluate( _sample_from_distributions(reloaded_distribution_output)) self.assertAllEqual(action_output_dict.keys(), reloaded_action_output_dict.keys()) for k in action_output_dict: if seeded: self.assertAllClose( action_output_dict[k], reloaded_action_output_dict[k], msg='\nMismatched dict key: %s.' % k) else: self.match_dtype_shape( action_output_dict[k], reloaded_action_output_dict[k], msg='\nMismatch dict key: %s.' % k) # With non-signature functions, we can check that passing a seed does # the right thing the second time. if seeded: tf.nest.map_structure(self.assertAllClose, action_output_value, reloaded_action_output_value) else: tf.nest.map_structure(self.match_dtype_shape, action_output_value, reloaded_action_output_value) tf.nest.map_structure(self.assertAllClose, distribution_output_value, reloaded_distribution_output_value) ## TFLite tests. # The converter must run outside of a TF1 graph context, even in # eager mode, to ensure the TF2 path is being executed. Only # works in TF2. if tf.compat.v1.executing_eagerly_outside_functions(): tflite_converter = tf.lite.TFLiteConverter.from_saved_model( path, signature_keys=['action']) tflite_converter.target_spec.supported_ops = [ tf.lite.OpsSet.TFLITE_BUILTINS, # TODO(b/111309333): Remove this when `has_input_fn_and_spec` # is `False` once TFLite has native support for RNG ops, atan, etc. tf.lite.OpsSet.SELECT_TF_OPS, ] tflite_serialized_model = tflite_converter.convert() tflite_interpreter = tf.lite.Interpreter( model_content=tflite_serialized_model) tflite_runner = tflite_interpreter.get_signature_runner('action') tflite_signature = tflite_interpreter.get_signature_list()['action'] if has_input_fn_and_spec: tflite_action_input_dict = { 'example': action_string_vector_values, } else: tflite_action_input_dict = collections.OrderedDict( (spec.name, value) for (spec, value) in zip( tf.nest.flatten(action_input_spec), tf.nest.flatten(action_input_values))) self.assertEqual( set(tflite_signature['inputs']), set(tflite_action_input_dict)) self.assertEqual( set(tflite_signature['outputs']), set(action_output_dict)) tflite_output = tflite_runner(**tflite_action_input_dict) self.assertAllClose(tflite_output, action_output_dict)
def testInferenceWithCheckpoint(self): if not common.has_eager_been_enabled(): self.skipTest('Only supported in TF2.x.') # Create and saved_model for a q_policy. network = q_network.QNetwork( input_tensor_spec=self._time_step_spec.observation, action_spec=self._action_spec) policy = q_policy.QPolicy(time_step_spec=self._time_step_spec, action_spec=self._action_spec, q_network=network) sample_input = self.evaluate( tensor_spec.sample_spec_nest(self._time_step_spec, outer_dims=(3, ))) saver = policy_saver.PolicySaver(policy, batch_size=None) path = os.path.join(self.get_temp_dir(), 'save_model') self.evaluate(tf.compat.v1.global_variables_initializer()) original_eval = self.evaluate(policy.action(sample_input)) saver.save(path) # Asign -1 to all variables in the policy. Making checkpoint different than # the initial saved_model. self.evaluate( tf.nest.map_structure(lambda v: v.assign(v * 0 + -1), policy.variables())) checkpoint_path = os.path.join(self.get_temp_dir(), 'checkpoint') saver.save_checkpoint(checkpoint_path) # Get an instance of the saved_model. reloaded_policy = tf.compat.v2.saved_model.load(path) self.evaluate( tf.compat.v1.initializers.variables( reloaded_policy.model_variables)) # Verify loaded saved_model variables are different than the current policy. model_variables = self.evaluate(policy.variables()) reloaded_model_variables = self.evaluate( reloaded_policy.model_variables) assert_np_not_equal = lambda a, b: self.assertFalse( np.equal(a, b).any()) tf.nest.map_structure(assert_np_not_equal, model_variables, reloaded_model_variables) # Update from checkpoint. checkpoint = tf.train.Checkpoint(policy=reloaded_policy) checkpoint_file_prefix = os.path.join(checkpoint_path, 'variables', 'variables') checkpoint.read( checkpoint_file_prefix).assert_existing_objects_matched() self.evaluate( tf.compat.v1.initializers.variables( reloaded_policy.model_variables)) # Verify variables are now equal. model_variables = self.evaluate(policy.variables()) reloaded_model_variables = self.evaluate( reloaded_policy.model_variables) assert_np_all_equal = lambda a, b: self.assertTrue( np.equal(a, b).all()) tf.nest.map_structure(assert_np_all_equal, model_variables, reloaded_model_variables) # Verify variable update affects inference. reloaded_eval = self.evaluate(reloaded_policy.action(sample_input)) tf.nest.map_structure(assert_np_not_equal, original_eval, reloaded_eval) current_eval = self.evaluate(policy.action(sample_input)) tf.nest.map_structure(assert_np_not_equal, current_eval, reloaded_eval)
def testSaveAction(self, seeded, has_state, distribution_net, has_input_fn_and_spec): with tf.compat.v1.Graph().as_default(): tf.compat.v1.set_random_seed(self._global_seed) with tf.compat.v1.Session().as_default(): global_step = common.create_variable('train_step', initial_value=0) if distribution_net: network = actor_distribution_network.ActorDistributionNetwork( self._time_step_spec.observation, self._action_spec) policy = actor_policy.ActorPolicy( time_step_spec=self._time_step_spec, action_spec=self._action_spec, actor_network=network) else: if has_state: network = q_rnn_network.QRnnNetwork( input_tensor_spec=self._time_step_spec.observation, action_spec=self._action_spec) else: network = q_network.QNetwork( input_tensor_spec=self._time_step_spec.observation, action_spec=self._action_spec) policy = q_policy.QPolicy( time_step_spec=self._time_step_spec, action_spec=self._action_spec, q_network=network) action_seed = 98723 batch_size = 3 action_inputs = tensor_spec.sample_spec_nest( (self._time_step_spec, policy.policy_state_spec), outer_dims=(batch_size, ), seed=4) action_input_values = self.evaluate(action_inputs) action_input_tensors = tf.nest.map_structure( tf.convert_to_tensor, action_input_values) action_output = policy.action(*action_input_tensors, seed=action_seed) self.evaluate(tf.compat.v1.global_variables_initializer()) action_output_dict = dict(((spec.name, value) for ( spec, value) in zip(tf.nest.flatten(policy.policy_step_spec), tf.nest.flatten(action_output)))) # Check output of the flattened signature call. (action_output_value, action_output_dict) = self.evaluate( (action_output, action_output_dict)) input_fn_and_spec = None if has_input_fn_and_spec: input_fn_and_spec = ( self._convert_string_vector_to_action_input, tf.TensorSpec((7, ), tf.string, name='example')) saver = policy_saver.PolicySaver( policy, batch_size=None, use_nest_path_signatures=False, seed=action_seed, input_fn_and_spec=input_fn_and_spec, train_step=global_step) path = os.path.join(self.get_temp_dir(), 'save_model_action') saver.save(path) with tf.compat.v1.Graph().as_default(): tf.compat.v1.set_random_seed(self._global_seed) with tf.compat.v1.Session().as_default(): reloaded = tf.compat.v2.saved_model.load(path) self.assertIn('action', reloaded.signatures) reloaded_action = reloaded.signatures['action'] if has_input_fn_and_spec: self._compare_input_output_specs( reloaded_action, expected_input_specs=input_fn_and_spec[1], expected_output_spec=policy.policy_step_spec, batch_input=True) else: self._compare_input_output_specs( reloaded_action, expected_input_specs=(self._time_step_spec, policy.policy_state_spec), expected_output_spec=policy.policy_step_spec, batch_input=True) # Reload action_input_values as tensors in the new graph. action_input_tensors = tf.nest.map_structure( tf.convert_to_tensor, action_input_values) action_input_spec = (self._time_step_spec, policy.policy_state_spec) function_action_input_dict = dict( (spec.name, value) for (spec, value) in zip(tf.nest.flatten(action_input_spec), tf.nest.flatten(action_input_tensors))) # NOTE(ebrevdo): The graph-level seeds for the policy and the reloaded # model are equal, which in addition to seeding the call to action() and # PolicySaver helps ensure equality of the output of action() in both # cases. self.assertEqual(reloaded_action.graph.seed, self._global_seed) def match_dtype_shape(x, y, msg=None): self.assertEqual(x.shape, y.shape, msg=msg) self.assertEqual(x.dtype, y.dtype, msg=msg) # The seed= argument for the SavedModel action call was given at # creation of the PolicySaver. if has_input_fn_and_spec: action_string_vector = self._convert_action_input_to_string_vector( action_input_tensors) reloaded_action_output_dict = reloaded_action( action_string_vector) reloaded_action_output = reloaded.action( action_string_vector) else: # This is the flat-signature function. reloaded_action_output_dict = reloaded_action( **function_action_input_dict) # This is the non-flat function. reloaded_action_output = reloaded.action( *action_input_tensors) if not has_state: # Try both cases: one with an empty policy_state and one with no # policy_state. Compare them. # NOTE(ebrevdo): The first call to .action() must be stored in # reloaded_action_output because this is the version being compared # later against the true action_output and the values will change # after the first call due to randomness. reloaded_action_output_no_input_state = reloaded.action( action_input_tensors[0]) # Even with a seed, multiple calls to action will get different # values, so here we just check the signature matches. tf.nest.map_structure( match_dtype_shape, reloaded_action_output_no_input_state, reloaded_action_output) self.evaluate(tf.compat.v1.global_variables_initializer()) (reloaded_action_output_dict, reloaded_action_output_value) = self.evaluate( (reloaded_action_output_dict, reloaded_action_output)) self.assertAllEqual(action_output_dict.keys(), reloaded_action_output_dict.keys()) for k in action_output_dict: if seeded: self.assertAllClose(action_output_dict[k], reloaded_action_output_dict[k], msg='\nMismatched dict key: %s.' % k) else: match_dtype_shape(action_output_dict[k], reloaded_action_output_dict[k], msg='\nMismatch dict key: %s.' % k) # With non-signature functions, we can check that passing a seed does # the right thing the second time. if seeded: tf.nest.map_structure(self.assertAllClose, action_output_value, reloaded_action_output_value) else: tf.nest.map_structure(match_dtype_shape, action_output_value, reloaded_action_output_value)
def testSaveGetInitialState(self): network = q_rnn_network.QRnnNetwork( input_tensor_spec=self._time_step_spec.observation, action_spec=self._action_spec, lstm_size=(40,)) policy = q_policy.QPolicy( time_step_spec=self._time_step_spec, action_spec=self._action_spec, q_network=network) train_step = common.create_variable('train_step', initial_value=0) saver_nobatch = policy_saver.PolicySaver( policy, train_step=train_step, batch_size=None, use_nest_path_signatures=False) path = os.path.join(self.get_temp_dir(), 'save_model_initial_state_nobatch') self.evaluate(tf.compat.v1.global_variables_initializer()) with self.cached_session(): saver_nobatch.save(path) reloaded_nobatch = tf.compat.v2.saved_model.load(path) self.evaluate( tf.compat.v1.initializers.variables(reloaded_nobatch.model_variables)) self.assertIn('get_initial_state', reloaded_nobatch.signatures) reloaded_get_initial_state = ( reloaded_nobatch.signatures['get_initial_state']) self._compare_input_output_specs( reloaded_get_initial_state, expected_input_specs=(tf.TensorSpec( dtype=tf.int32, shape=(), name='batch_size'),), expected_output_spec=policy.policy_state_spec, batch_input=False, batch_size=None) initial_state = policy.get_initial_state(batch_size=3) initial_state = self.evaluate(initial_state) reloaded_nobatch_initial_state = reloaded_nobatch.get_initial_state( batch_size=3) reloaded_nobatch_initial_state = self.evaluate( reloaded_nobatch_initial_state) tf.nest.map_structure(self.assertAllClose, initial_state, reloaded_nobatch_initial_state) saver_batch = policy_saver.PolicySaver( policy, train_step=train_step, batch_size=3, use_nest_path_signatures=False) path = os.path.join(self.get_temp_dir(), 'save_model_initial_state_batch') with self.cached_session(): saver_batch.save(path) reloaded_batch = tf.compat.v2.saved_model.load(path) self.evaluate( tf.compat.v1.initializers.variables(reloaded_batch.model_variables)) self.assertIn('get_initial_state', reloaded_batch.signatures) reloaded_get_initial_state = reloaded_batch.signatures['get_initial_state'] self._compare_input_output_specs( reloaded_get_initial_state, expected_input_specs=(), expected_output_spec=policy.policy_state_spec, batch_input=False, batch_size=3) reloaded_batch_initial_state = reloaded_batch.get_initial_state() reloaded_batch_initial_state = self.evaluate(reloaded_batch_initial_state) tf.nest.map_structure(self.assertAllClose, initial_state, reloaded_batch_initial_state)
input_tensor_spec = tf_env.observation_spec() time_step_spec = ts.time_step_spec(input_tensor_spec) action_spec = tf_env.action_spec() num_actions = env.size**2 batch_size = 1 observation = tf.cast( [(np.random.randint(env.size - 2, size=(env.size, env.size)) + 1).reshape(25) for _ in range(1)], tf.int32) time_steps = ts.restart(observation, batch_size=batch_size) my_q_network = QNetwork(input_tensor_spec=input_tensor_spec, action_spec=action_spec, num_actions=num_actions) my_q_policy = q_policy.QPolicy(time_step_spec, action_spec, q_network=my_q_network) action_step = my_q_policy.action(time_steps) distribution_step = my_q_policy.distribution(time_steps) print('Action:') print(action_step.action) print('Action distribution:') print(distribution_step.action) num_episodes = tf_metrics.NumberOfEpisodes() env_steps = tf_metrics.EnvironmentSteps() observers = [num_episodes, env_steps] driver = dynamic_episode_driver.DynamicEpisodeDriver(tf_env, my_q_policy,
def __init__( self, time_step_spec, action_spec, q_network, optimizer, epsilon_greedy=0.1, # Params for target network updates target_update_tau=1.0, target_update_period=1, # Params for training. td_errors_loss_fn=None, gamma=1.0, reward_scale_factor=1.0, gradient_clipping=None, # Params for debugging debug_summaries=False, summarize_grads_and_vars=False): """Creates a DQN Agent. Args: time_step_spec: A `TimeStep` spec of the expected time_steps. action_spec: A nest of BoundedTensorSpec representing the actions. q_network: A tf_agents.network.Network to be used by the agent. The network will be called with call(observation, step_type). optimizer: The optimizer to use for training. epsilon_greedy: probability of choosing a random action in the default epsilon-greedy collect policy (used only if a wrapper is not provided to the collect_policy method). target_update_tau: Factor for soft update of the target networks. target_update_period: Period for soft update of the target networks. td_errors_loss_fn: A function for computing the TD errors loss. If None, a default value of element_wise_huber_loss is used. This function takes as input the target and the estimated Q values and returns the loss for each element of the batch. gamma: A discount factor for future rewards. reward_scale_factor: Multiplicative scale for the reward. gradient_clipping: Norm length to clip gradients. debug_summaries: A bool to gather debug summaries. summarize_grads_and_vars: If True, gradient and network variable summaries will be written during training. Raises: ValueError: If the action spec contains more than one action. """ flat_action_spec = nest.flatten(action_spec) self._num_actions = [ spec.maximum - spec.minimum + 1 for spec in flat_action_spec ] # TODO(oars): Get DQN working with more than one dim in the actions. if len(flat_action_spec) > 1 or flat_action_spec[0].shape.ndims > 1: raise ValueError('Only one dimensional actions are supported now.') self._q_network = q_network self._target_q_network = self._q_network.copy(name='TargetQNetwork') self._epsilon_greedy = epsilon_greedy self._target_update_tau = target_update_tau self._target_update_period = target_update_period self._optimizer = optimizer self._td_errors_loss_fn = td_errors_loss_fn or element_wise_huber_loss self._gamma = gamma self._reward_scale_factor = reward_scale_factor self._gradient_clipping = gradient_clipping self._target_update_train_op = None policy = q_policy.QPolicy(time_step_spec, action_spec, q_network=self._q_network) collect_policy = epsilon_greedy_policy.EpsilonGreedyPolicy( policy, epsilon=self._epsilon_greedy) policy = greedy_policy.GreedyPolicy(policy) super(DqnAgent, self).__init__( time_step_spec, action_spec, policy, collect_policy, train_sequence_length=2 if not q_network.state_spec else None, debug_summaries=debug_summaries, summarize_grads_and_vars=summarize_grads_and_vars)
def testUpdateWithCheckpoint(self): if not common.has_eager_been_enabled(): self.skipTest('Only supported in TF2.x.') # Create and saved_model for a q_policy. network = q_network.QNetwork( input_tensor_spec=self._time_step_spec.observation, action_spec=self._action_spec) policy = q_policy.QPolicy(time_step_spec=self._time_step_spec, action_spec=self._action_spec, q_network=network) saver = policy_saver.PolicySaver(policy, batch_size=None) path = os.path.join(self.get_temp_dir(), 'save_model') self.evaluate(tf.compat.v1.global_variables_initializer()) saver.save(path) # Assign -1 to all variables in the policy. Making checkpoint different than # the initial saved_model. self.evaluate( tf.nest.map_structure(lambda v: v.assign(v * 0 + -1), policy.variables())) checkpoint_path = os.path.join(self.get_temp_dir(), 'checkpoint') saver.save_checkpoint(checkpoint_path) # Get an instance of the saved_model. reloaded_policy = tf.compat.v2.saved_model.load(path) self.evaluate( tf.compat.v1.initializers.variables( reloaded_policy.model_variables)) # Verify loaded saved_model variables are different than the current policy. model_variables = self.evaluate(policy.variables()) reloaded_model_variables = self.evaluate( reloaded_policy.model_variables) assert_np_not_equal = lambda a, b: self.assertFalse( np.equal(a, b).any()) tf.nest.map_structure(assert_np_not_equal, model_variables, reloaded_model_variables) # Update from checkpoint. checkpoint = tf.train.Checkpoint(policy=reloaded_policy) manager = tf.train.CheckpointManager(checkpoint, directory=checkpoint_path, max_to_keep=None) checkpoint.restore(manager.latest_checkpoint).expect_partial() self.evaluate( tf.compat.v1.initializers.variables( reloaded_policy.model_variables)) # Verify variables are now equal. model_variables = self.evaluate(policy.variables()) reloaded_model_variables = self.evaluate( reloaded_policy.model_variables) assert_np_all_equal = lambda a, b: self.assertTrue( np.equal(a, b).all()) tf.nest.map_structure(assert_np_all_equal, model_variables, reloaded_model_variables)
def __init__( self, time_step_spec, action_spec, q_network, optimizer, epsilon_greedy=0.1, n_step_update=1, boltzmann_temperature=None, emit_log_probability=False, update_period=None, # Params for target network updates target_update_tau=1.0, target_update_period=1, # Params for training. td_errors_loss_fn=None, gamma=1.0, reward_scale_factor=1.0, gradient_clipping=None, # Params for debugging debug_summaries=False, enable_functions=True, summarize_grads_and_vars=False, train_step_counter=None, name=None): """Creates a DQN Agent. Args: time_step_spec: A `TimeStep` spec of the expected time_steps. action_spec: A nest of BoundedTensorSpec representing the actions. q_network: A tf_agents.network.Network to be used by the agent. The network will be called with call(observation, step_type). optimizer: The optimizer to use for training. epsilon_greedy: probability of choosing a random action in the default epsilon-greedy collect policy (used only if a wrapper is not provided to the collect_policy method). n_step_update: The number of steps to consider when computing TD error and TD loss. Defaults to single-step updates. Note that this requires the user to call train on Trajectory objects with a time dimension of `n_step_update + 1`. However, note that we do not yet support `n_step_update > 1` in the case of RNNs (i.e., non-empty `q_network.state_spec`). boltzmann_temperature: Temperature value to use for Boltzmann sampling of the actions during data collection. The closer to 0.0, the higher the probability of choosing the best action. emit_log_probability: Whether policies emit log probabilities or not. update_period: Update period. target_update_tau: Factor for soft update of the target networks. target_update_period: Period for soft update of the target networks. td_errors_loss_fn: A function for computing the TD errors loss. If None, a default value of element_wise_huber_loss is used. This function takes as input the target and the estimated Q values and returns the loss for each element of the batch. gamma: A discount factor for future rewards. reward_scale_factor: Multiplicative scale for the reward. gradient_clipping: Norm length to clip gradients. debug_summaries: A bool to gather debug summaries. enable_functions: A bool to decide whether or not to enable tf function summarize_grads_and_vars: If True, gradient and network variable summaries will be written during training. train_step_counter: An optional counter to increment every time the train op is run. Defaults to the global_step. name: The name of this agent. All variables in this module will fall under that name. Defaults to the class name. Raises: ValueError: If the action spec contains more than one action or action spec minimum is not equal to 0. NotImplementedError: If `q_network` has non-empty `state_spec` (i.e., an RNN is provided) and `n_step_update > 1`. """ tf.Module.__init__(self, name=name) flat_action_spec = tf.nest.flatten(action_spec) self._num_actions = [ spec.maximum - spec.minimum + 1 for spec in flat_action_spec ] if len(flat_action_spec) > 1 or flat_action_spec[0].shape.ndims > 1: raise ValueError('Only one dimensional actions are supported now.') if not all(spec.minimum == 0 for spec in flat_action_spec): raise ValueError( 'Action specs should have minimum of 0, but saw: {0}'.format( [spec.minimum for spec in flat_action_spec])) if epsilon_greedy is not None and boltzmann_temperature is not None: raise ValueError( 'Configured both epsilon_greedy value {} and temperature {}, ' 'however only one of them can be used for exploration.'.format( epsilon_greedy, boltzmann_temperature)) self._q_network = q_network self._target_q_network = self._q_network.copy(name='TargetQNetwork') self._epsilon_greedy = epsilon_greedy self._n_step_update = n_step_update self._boltzmann_temperature = boltzmann_temperature self._optimizer = optimizer self._td_errors_loss_fn = td_errors_loss_fn or element_wise_huber_loss self._gamma = gamma self._reward_scale_factor = reward_scale_factor self._gradient_clipping = gradient_clipping self._update_target = self._get_target_updater(target_update_tau, target_update_period) policy = q_policy.QPolicy( time_step_spec, action_spec, q_network=self._q_network, emit_log_probability=emit_log_probability) if boltzmann_temperature is not None: collect_policy = boltzmann_policy.BoltzmannPolicy( policy, temperature=self._boltzmann_temperature) else: collect_policy = epsilon_greedy_policy.EpsilonGreedyPolicy( policy, epsilon=self._epsilon_greedy) policy = greedy_policy.GreedyPolicy(policy) if q_network.state_spec and n_step_update != 1: raise NotImplementedError( 'DqnAgent does not currently support n-step updates with stateful ' 'networks (i.e., RNNs), but n_step_update = {}'.format(n_step_update)) train_sequence_length = ( n_step_update + 1 if not q_network.state_spec else None) super(DqnAgent, self).__init__( time_step_spec, action_spec, policy, collect_policy, train_sequence_length=train_sequence_length, update_period=update_period, debug_summaries=debug_summaries, enable_functions=enable_functions, summarize_grads_and_vars=summarize_grads_and_vars, train_step_counter=train_step_counter) tf.compat.v1.summary.scalar( 'epsilon/' + self.name, self._epsilon_greedy, collections=['train_' + self.name])
def __init__( self, time_step_spec, action_spec, q_network, optimizer, epsilon_greedy=0.1, boltzmann_temperature=None, # Params for target network updates target_update_tau=1.0, target_update_period=1, # Params for training. td_errors_loss_fn=None, gamma=1.0, reward_scale_factor=1.0, gradient_clipping=None, # Params for debugging debug_summaries=False, summarize_grads_and_vars=False, train_step_counter=None, name=None): """Creates a DQN Agent. Args: time_step_spec: A `TimeStep` spec of the expected time_steps. action_spec: A nest of BoundedTensorSpec representing the actions. q_network: A tf_agents.network.Network to be used by the agent. The network will be called with call(observation, step_type). optimizer: The optimizer to use for training. epsilon_greedy: probability of choosing a random action in the default epsilon-greedy collect policy (used only if a wrapper is not provided to the collect_policy method). boltzmann_temperature: Temperature value to use for Boltzmann sampling of the actions during data collection. The closer to 0.0, the higher the probability of choosing the best action. target_update_tau: Factor for soft update of the target networks. target_update_period: Period for soft update of the target networks. td_errors_loss_fn: A function for computing the TD errors loss. If None, a default value of element_wise_huber_loss is used. This function takes as input the target and the estimated Q values and returns the loss for each element of the batch. gamma: A discount factor for future rewards. reward_scale_factor: Multiplicative scale for the reward. gradient_clipping: Norm length to clip gradients. debug_summaries: A bool to gather debug summaries. summarize_grads_and_vars: If True, gradient and network variable summaries will be written during training. train_step_counter: An optional counter to increment every time the train op is run. Defaults to the global_step. name: The name of this agent. All variables in this module will fall under that name. Defaults to the class name. Raises: ValueError: If the action spec contains more than one action or action spec minimum is not equal to 0. """ tf.Module.__init__(self, name=name) flat_action_spec = tf.nest.flatten(action_spec) self._num_actions = [ spec.maximum - spec.minimum + 1 for spec in flat_action_spec ] # TODO(oars): Get DQN working with more than one dim in the actions. if len(flat_action_spec) > 1 or flat_action_spec[0].shape.ndims > 1: raise ValueError('Only one dimensional actions are supported now.') if not all(spec.minimum == 0 for spec in flat_action_spec): raise ValueError( 'Action specs should have minimum of 0, but saw: {0}'.format( [spec.minimum for spec in flat_action_spec])) if epsilon_greedy is not None and boltzmann_temperature is not None: raise ValueError( 'Configured both epsilon_greedy value {} and temperature {}, ' 'however only one of them can be used for exploration.'.format( epsilon_greedy, boltzmann_temperature)) self._q_network = q_network self._target_q_network = self._q_network.copy(name='TargetQNetwork') self._epsilon_greedy = epsilon_greedy self._boltzmann_temperature = boltzmann_temperature self._optimizer = optimizer self._td_errors_loss_fn = td_errors_loss_fn or element_wise_huber_loss self._gamma = gamma self._reward_scale_factor = reward_scale_factor self._gradient_clipping = gradient_clipping self._update_target = self._get_target_updater( target_update_tau, target_update_period) policy = q_policy.QPolicy( time_step_spec, action_spec, q_network=self._q_network) if boltzmann_temperature is not None: collect_policy = boltzmann_policy.BoltzmannPolicy( policy, temperature=self._boltzmann_temperature) else: collect_policy = epsilon_greedy_policy.EpsilonGreedyPolicy( policy, epsilon=self._epsilon_greedy) policy = greedy_policy.GreedyPolicy(policy) super(DqnAgent, self).__init__( time_step_spec, action_spec, policy, collect_policy, train_sequence_length=2 if not q_network.state_spec else None, debug_summaries=debug_summaries, summarize_grads_and_vars=summarize_grads_and_vars, train_step_counter=train_step_counter)
def testSaveRestore(self, batch_size): policy_save_path = os.path.join(flags.FLAGS.test_tmpdir, 'policy', str(batch_size)) # Construct a policy to be saved under a tf.Graph instance. policy_saved_graph = tf.Graph() with policy_saved_graph.as_default(): tf_policy = q_policy.QPolicy( self._time_step_spec, self._action_spec, DummyNet(use_constant_initializer=False)) # Parameterized tests reuse temp directories, make no save exists. try: tf.io.gfile.listdir(policy_save_path) tf.io.gfile.rmtree(policy_save_path) except tf.errors.NotFoundError: pass policy_saved = py_tf_policy.PyTFPolicy(tf_policy) policy_saved.session = tf.compat.v1.Session( graph=policy_saved_graph) policy_saved.initialize(batch_size) policy_saved.save(policy_dir=policy_save_path, graph=policy_saved_graph) # Verify that index files were written. There will also be some number of # data files, but this depends on the number of devices. self.assertContainsSubset( set(['checkpoint', 'ckpt-0.index']), set(tf.io.gfile.listdir(policy_save_path))) # Construct a policy to be restored under another tf.Graph instance. policy_restore_graph = tf.Graph() with policy_restore_graph.as_default(): tf_policy = q_policy.QPolicy( self._time_step_spec, self._action_spec, DummyNet(use_constant_initializer=False)) policy_restored = py_tf_policy.PyTFPolicy(tf_policy) policy_restored.session = tf.compat.v1.Session( graph=policy_restore_graph) policy_restored.initialize(batch_size) random_init_vals = policy_restored.session.run( tf_policy.variables()) policy_restored.restore(policy_dir=policy_save_path, graph=policy_restore_graph) restored_vals = policy_restored.session.run(tf_policy.variables()) for random_init_var, restored_var in zip(random_init_vals, restored_vals): self.assertFalse(np.array_equal(random_init_var, restored_var)) # Check that variables in the two policies have identical values. with policy_restore_graph.as_default(): restored_values = policy_restored.session.run( tf.compat.v1.global_variables()) with policy_saved_graph.as_default(): initial_values = policy_saved.session.run( tf.compat.v1.global_variables()) # Networks have two fully connected layers. self.assertLen(initial_values, 4) self.assertLen(restored_values, 4) for initial_var, restored_var in zip(initial_values, restored_values): np.testing.assert_array_equal(initial_var, restored_var)
def testSaveAction(self): if not tf.executing_eagerly(): self.skipTest( 'b/129079730: PolicySaver does not work in TF1.x yet') q_network = q_rnn_network.QRnnNetwork( input_tensor_spec=self._time_step_spec.observation, action_spec=self._action_spec) policy = q_policy.QPolicy(time_step_spec=self._time_step_spec, action_spec=self._action_spec, q_network=q_network) action_seed = 98723 saver = policy_saver.PolicySaver(policy, batch_size=None, seed=action_seed) path = os.path.join(tf.compat.v1.test.get_temp_dir(), 'save_model_action') saver.save(path) reloaded = tf.compat.v2.saved_model.load(path) self.assertIn('action', reloaded.signatures) reloaded_action = reloaded.signatures['action'] self._compare_input_output_specs( reloaded_action, expected_input_specs=(self._time_step_spec, policy.policy_state_spec), expected_output_spec=policy.policy_step_spec, batch_input=True) batch_size = 3 action_inputs = tensor_spec.sample_spec_nest( (self._time_step_spec, policy.policy_state_spec), outer_dims=(batch_size, ), seed=4) function_action_input_dict = dict( (spec.name, value) for (spec, value) in zip( tf.nest.flatten((self._time_step_spec, policy.policy_state_spec )), tf.nest.flatten(action_inputs))) # NOTE(ebrevdo): The graph-level seeds for the policy and the reloaded model # are equal, which in addition to seeding the call to action() and # PolicySaver helps ensure equality of the output of action() in both cases. self.assertEqual(reloaded_action.graph.seed, self._global_seed) action_output = policy.action(*action_inputs, seed=action_seed) # The seed= argument for the SavedModel action call was given at creation of # the PolicySaver. reloaded_action_output_dict = reloaded_action( **function_action_input_dict) action_output_dict = dict( ((spec.name, value) for (spec, value) in zip(tf.nest.flatten(policy.policy_step_spec), tf.nest.flatten(action_output)))) action_output_dict = self.evaluate(action_output_dict) reloaded_action_output_dict = self.evaluate( reloaded_action_output_dict) self.assertAllEqual(action_output_dict.keys(), reloaded_action_output_dict.keys()) for k in action_output_dict: self.assertAllClose(action_output_dict[k], reloaded_action_output_dict[k], msg='\nMismatched dict key: %s.' % k)
def __init__( self, time_step_spec, action_spec, q_network, optimizer, epsilon_greedy=0.1, n_step_update=1, boltzmann_temperature=None, emit_log_probability=False, # Params for target network updates target_q_network=None, target_update_tau=1.0, target_update_period=1, # Params for training. td_errors_loss_fn=None, gamma=1.0, reward_scale_factor=1.0, gradient_clipping=None, # Params for debugging debug_summaries=False, summarize_grads_and_vars=False, train_step_counter=None, name=None): """Creates a DQN Agent. Args: time_step_spec: A `TimeStep` spec of the expected time_steps. action_spec: A nest of BoundedTensorSpec representing the actions. q_network: A `tf_agents.network.Network` to be used by the agent. The network will be called with `call(observation, step_type)` and should emit logits over the action space. optimizer: The optimizer to use for training. epsilon_greedy: probability of choosing a random action in the default epsilon-greedy collect policy (used only if a wrapper is not provided to the collect_policy method). n_step_update: The number of steps to consider when computing TD error and TD loss. Defaults to single-step updates. Note that this requires the user to call train on Trajectory objects with a time dimension of `n_step_update + 1`. However, note that we do not yet support `n_step_update > 1` in the case of RNNs (i.e., non-empty `q_network.state_spec`). boltzmann_temperature: Temperature value to use for Boltzmann sampling of the actions during data collection. The closer to 0.0, the higher the probability of choosing the best action. emit_log_probability: Whether policies emit log probabilities or not. target_q_network: (Optional.) A `tf_agents.network.Network` to be used as the target network during Q learning. Every `target_update_period` train steps, the weights from `q_network` are copied (possibly with smoothing via `target_update_tau`) to `target_q_network`. If `target_q_network` is not provided, it is created by making a copy of `q_network`, which initializes a new network with the same structure and its own layers and weights. Performing a `Network.copy` does not work when the network instance already has trainable parameters (e.g., has already been built, or when the network is sharing layers with another). In these cases, it is up to you to build a copy having weights that are not shared with the original `q_network`, so that this can be used as a target network. If you provide a `target_q_network` that shares any weights with `q_network`, a warning will be logged but no exception is thrown. target_update_tau: Factor for soft update of the target networks. target_update_period: Period for soft update of the target networks. td_errors_loss_fn: A function for computing the TD errors loss. If None, a default value of element_wise_huber_loss is used. This function takes as input the target and the estimated Q values and returns the loss for each element of the batch. gamma: A discount factor for future rewards. reward_scale_factor: Multiplicative scale for the reward. gradient_clipping: Norm length to clip gradients. debug_summaries: A bool to gather debug summaries. summarize_grads_and_vars: If True, gradient and network variable summaries will be written during training. train_step_counter: An optional counter to increment every time the train op is run. Defaults to the global_step. name: The name of this agent. All variables in this module will fall under that name. Defaults to the class name. Raises: ValueError: If the action spec contains more than one action or action spec minimum is not equal to 0. NotImplementedError: If `q_network` has non-empty `state_spec` (i.e., an RNN is provided) and `n_step_update > 1`. """ tf.Module.__init__(self, name=name) self._check_action_spec(action_spec) if epsilon_greedy is not None and boltzmann_temperature is not None: raise ValueError( 'Configured both epsilon_greedy value {} and temperature {}, ' 'however only one of them can be used for exploration.'.format( epsilon_greedy, boltzmann_temperature)) self._q_network = q_network self._target_q_network = common.maybe_copy_target_network_with_checks( self._q_network, target_q_network, 'TargetQNetwork') self._epsilon_greedy = epsilon_greedy self._n_step_update = n_step_update self._boltzmann_temperature = boltzmann_temperature self._optimizer = optimizer self._td_errors_loss_fn = td_errors_loss_fn or common.element_wise_huber_loss self._gamma = gamma self._reward_scale_factor = reward_scale_factor self._gradient_clipping = gradient_clipping self._update_target = self._get_target_updater(target_update_tau, target_update_period) policy, collect_policy = self._setup_policy(time_step_spec, action_spec, boltzmann_temperature, emit_log_probability) self._greedy_policy = policy self._target_policy = q_policy.QPolicy( time_step_spec, action_spec, q_network=self._target_q_network) self._target_greedy_policy = greedy_policy.GreedyPolicy( self._target_policy) if q_network.state_spec and n_step_update != 1: raise NotImplementedError( 'DqnAgent does not currently support n-step updates with stateful ' 'networks (i.e., RNNs), but n_step_update = {}'.format( n_step_update)) train_sequence_length = (n_step_update + 1 if not q_network.state_spec else None) super(DqnAgent, self).__init__(time_step_spec, action_spec, policy, collect_policy, train_sequence_length=train_sequence_length, debug_summaries=debug_summaries, summarize_grads_and_vars=summarize_grads_and_vars, train_step_counter=train_step_counter)
def __init__( self, time_step_spec, action_spec, cloning_network, optimizer, epsilon_greedy=0.1, # Params for training. loss_fn=None, gradient_clipping=None, # Params for debugging debug_summaries=False, summarize_grads_and_vars=False): """Creates an behavioral cloning Agent. Args: time_step_spec: A `TimeStep` spec of the expected time_steps. action_spec: A nest of BoundedTensorSpec representing the actions. cloning_network: A tf_agents.network.Network to be used by the agent. The network will be called as ``` network(observation, step_type, network_state=None) ``` (with `network_state` optional) and must return a 2-tuple with elements `(output, next_network_state)` where `output` will be passed as the first argument to `loss_fn`, and used by a `Policy`. Input tensors will be shaped `[batch, time, ...]` when training, and they will be shaped `[batch, ...]` when the network is called within a `Policy`. If `cloning_network` has an empty network state, then for training `time` will always be `1` (individual examples). optimizer: The optimizer to use for training. epsilon_greedy: probability of choosing a random action in the default epsilon-greedy collect policy (used only if a wrapper is not provided to the collect_policy method). loss_fn: A function for computing the error between the output of the cloning network and the action that was taken. If None, the loss depends on the action dtype. If the dtype is integer, then `loss_fn` is ```python def loss_fn(logits, action): return tf.nn.sparse_softmax_cross_entropy_with_logits( labels=action - action_spec.minimum, logits=logits) ``` If the dtype is floating point, the loss is `tf.math.squared_difference`. `loss_fn` must return a loss value for each element of the batch. gradient_clipping: Norm length to clip gradients. debug_summaries: A bool to gather debug summaries. summarize_grads_and_vars: If True, gradient and network variable summaries will be written during training. Raises: NotImplementedError: If the action spec contains more than one action. """ flat_action_spec = nest.flatten(action_spec) self._num_actions = [ spec.maximum - spec.minimum + 1 for spec in flat_action_spec ] # TODO(oars): Get behavioral cloning working with more than one dim in # the actions. if len(flat_action_spec) > 1: raise NotImplementedError( 'Multi-arity actions are not currently supported.') if flat_action_spec[0].dtype.is_floating: if loss_fn is None: loss_fn = tf.math.squared_difference else: if flat_action_spec[0].shape.ndims > 1: raise NotImplementedError( 'Only scalar and one dimensional integer actions are supported.' ) if loss_fn is None: # TODO(ebrevdo): Maybe move the subtraction of the minimum into a # self._label_fn and rewrite this. def xent_loss_fn(logits, actions): # Subtract the minimum so that we get a proper cross entropy loss on # [0, maximum - minimum). return tf.nn.sparse_softmax_cross_entropy_with_logits( logits=logits, labels=actions - flat_action_spec[0].minimum) loss_fn = xent_loss_fn self._cloning_network = cloning_network self._loss_fn = loss_fn self._epsilon_greedy = epsilon_greedy self._optimizer = optimizer self._gradient_clipping = gradient_clipping policy = q_policy.QPolicy(time_step_spec, action_spec, q_network=self._cloning_network) collect_policy = epsilon_greedy_policy.EpsilonGreedyPolicy( policy, epsilon=self._epsilon_greedy) policy = greedy_policy.GreedyPolicy(policy) super(BehavioralCloningAgent, self).__init__(time_step_spec, action_spec, policy, collect_policy, train_sequence_length=1 if not cloning_network.state_spec else None, debug_summaries=debug_summaries, summarize_grads_and_vars=summarize_grads_and_vars)
def testMultipleActionsRaiseError(self): action_spec = [tensor_spec.BoundedTensorSpec([], tf.int32, 0, 1)] * 2 with self.assertRaisesRegexp(ValueError, 'Only scalar actions'): q_policy.QPolicy(self._time_step_spec, action_spec, q_network=DummyNet())
def testActionSpecsCompatible(self): q_net = DummyNetWithActionSpec(self._action_spec) q_policy.QPolicy(self._time_step_spec, self._action_spec, q_net)
def testBuild(self): policy = q_policy.QPolicy( self._time_step_spec, self._action_spec, q_network=DummyNet()) self.assertEqual(policy.time_step_spec, self._time_step_spec) self.assertEqual(policy.action_spec, self._action_spec)
def testSaveAction(self, seeded, has_state): if not tf.executing_eagerly(): self.skipTest( 'b/129079730: PolicySaver does not work in TF1.x yet') if has_state: network = q_rnn_network.QRnnNetwork( input_tensor_spec=self._time_step_spec.observation, action_spec=self._action_spec) else: network = q_network.QNetwork( input_tensor_spec=self._time_step_spec.observation, action_spec=self._action_spec) policy = q_policy.QPolicy(time_step_spec=self._time_step_spec, action_spec=self._action_spec, q_network=network) action_seed = 98723 saver = policy_saver.PolicySaver(policy, batch_size=None, use_nest_path_signatures=False, seed=action_seed) path = os.path.join(self.get_temp_dir(), 'save_model_action') saver.save(path) reloaded = tf.compat.v2.saved_model.load(path) self.assertIn('action', reloaded.signatures) reloaded_action = reloaded.signatures['action'] self._compare_input_output_specs( reloaded_action, expected_input_specs=(self._time_step_spec, policy.policy_state_spec), expected_output_spec=policy.policy_step_spec, batch_input=True) batch_size = 3 action_inputs = tensor_spec.sample_spec_nest( (self._time_step_spec, policy.policy_state_spec), outer_dims=(batch_size, ), seed=4) function_action_input_dict = dict( (spec.name, value) for (spec, value) in zip( tf.nest.flatten((self._time_step_spec, policy.policy_state_spec )), tf.nest.flatten(action_inputs))) # NOTE(ebrevdo): The graph-level seeds for the policy and the reloaded model # are equal, which in addition to seeding the call to action() and # PolicySaver helps ensure equality of the output of action() in both cases. self.assertEqual(reloaded_action.graph.seed, self._global_seed) action_output = policy.action(*action_inputs, seed=action_seed) # The seed= argument for the SavedModel action call was given at creation of # the PolicySaver. # This is the flat-signature function. reloaded_action_output_dict = reloaded_action( **function_action_input_dict) def match_dtype_shape(x, y, msg=None): self.assertEqual(x.shape, y.shape, msg=msg) self.assertEqual(x.dtype, y.dtype, msg=msg) # This is the non-flat function. if has_state: reloaded_action_output = reloaded.action(*action_inputs) else: # Try both cases: one with an empty policy_state and one with no # policy_state. Compare them. # NOTE(ebrevdo): The first call to .action() must be stored in # reloaded_action_output because this is the version being compared later # against the true action_output and the values will change after the # first call due to randomness. reloaded_action_output = reloaded.action(*action_inputs) reloaded_action_output_no_input_state = reloaded.action( action_inputs[0]) # Even with a seed, multiple calls to action will get different values, # so here we just check the signature matches. tf.nest.map_structure(match_dtype_shape, reloaded_action_output_no_input_state, reloaded_action_output) action_output_dict = dict( ((spec.name, value) for (spec, value) in zip(tf.nest.flatten(policy.policy_step_spec), tf.nest.flatten(action_output)))) # Check output of the flattened signature call. action_output_dict = self.evaluate(action_output_dict) reloaded_action_output_dict = self.evaluate( reloaded_action_output_dict) self.assertAllEqual(action_output_dict.keys(), reloaded_action_output_dict.keys()) for k in action_output_dict: if seeded: self.assertAllClose(action_output_dict[k], reloaded_action_output_dict[k], msg='\nMismatched dict key: %s.' % k) else: match_dtype_shape(action_output_dict[k], reloaded_action_output_dict[k], msg='\nMismatch dict key: %s.' % k) # Check output of the proper structured call. action_output = self.evaluate(action_output) reloaded_action_output = self.evaluate(reloaded_action_output) # With non-signature functions, we can check that passing a seed does the # right thing the second time. if seeded: tf.nest.map_structure(self.assertAllClose, action_output, reloaded_action_output) else: tf.nest.map_structure(match_dtype_shape, action_output, reloaded_action_output)
def __init__(self, time_step_spec, env_action_spec, replay_buffer_action_spec, actor_network, critic_network, actor_optimizer, critic_optimizer, exploration_noise_std=0.1, epsilon_greedy=0.1, boltzmann_temperature=None, emit_log_probability=False, critic_network_2=None, target_actor_network=None, target_critic_network=None, target_critic_network_2=None, target_update_tau=1.0, target_update_period=1, actor_update_period=1, dqda_clipping=None, td_errors_loss_fn=None, gamma=1.0, reward_scale_factor=1.0, target_policy_noise=0.2, target_policy_noise_clip=0.5, gradient_clipping=None, debug_summaries=False, summarize_grads_and_vars=False, train_step_counter=None, name=None): """Creates a Td3Agent Agent. Args: time_step_spec: A `TimeStep` spec of the expected time_steps. env_action_spec: A nest of BoundedTensorSpec representing the environment actions. replay_buffer_action_spec: A nest of BoundedTensorSpec representing the actions serving as input for critic. actor_network: A tf_agents.network.Network to be used by the agent. The network will be called with call(observation, step_type). critic_network: A tf_agents.network.Network to be used by the agent. The network will be called with call(observation, action, step_type). actor_optimizer: The default optimizer to use for the actor network. critic_optimizer: The default optimizer to use for the critic network. exploration_noise_std: Scale factor on exploration policy noise. critic_network_2: (Optional.) A `tf_agents.network.Network` to be used as the second critic network during Q learning. The weights from `critic_network` are copied if this is not provided. target_actor_network: (Optional.) A `tf_agents.network.Network` to be used as the target actor network during Q learning. Every `target_update_period` train steps, the weights from `actor_network` are copied (possibly withsmoothing via `target_update_tau`) to ` target_actor_network`. If `target_actor_network` is not provided, it is created by making a copy of `actor_network`, which initializes a new network with the same structure and its own layers and weights. Performing a `Network.copy` does not work when the network instance already has trainable parameters (e.g., has already been built, or when the network is sharing layers with another). In these cases, it is up to you to build a copy having weights that are not shared with the original `actor_network`, so that this can be used as a target network. If you provide a `target_actor_network` that shares any weights with `actor_network`, a warning will be logged but no exception is thrown. target_critic_network: (Optional.) Similar network as target_actor_network but for the critic_network. See documentation for target_actor_network. target_critic_network_2: (Optional.) Similar network as target_actor_network but for the critic_network_2. See documentation for target_actor_network. Will only be used if 'critic_network_2' is also specified. target_update_tau: Factor for soft update of the target networks. target_update_period: Period for soft update of the target networks. actor_update_period: Period for the optimization step on actor network. dqda_clipping: A scalar or float clips the gradient dqda element-wise between [-dqda_clipping, dqda_clipping]. Default is None representing no clippiing. td_errors_loss_fn: A function for computing the TD errors loss. If None, a default value of elementwise huber_loss is used. gamma: A discount factor for future rewards. reward_scale_factor: Multiplicative scale for the reward. target_policy_noise: Scale factor on target action noise target_policy_noise_clip: Value to clip noise. gradient_clipping: Norm length to clip gradients. debug_summaries: A bool to gather debug summaries. summarize_grads_and_vars: If True, gradient and network variable summaries will be written during training. train_step_counter: An optional counter to increment every time the train op is run. Defaults to the global_step. name: The name of this agent. All variables in this module will fall under that name. Defaults to the class name. """ tf.Module.__init__(self, name=name) self._actor_network = actor_network self._target_actor_network = common.maybe_copy_target_network_with_checks( self._actor_network, target_actor_network, 'TargetActorNetwork') self._critic_network_1 = critic_network self._target_critic_network_1 = ( common.maybe_copy_target_network_with_checks( self._critic_network_1, target_critic_network, 'TargetCriticNetwork1')) if critic_network_2 is not None: self._critic_network_2 = critic_network_2 else: self._critic_network_2 = critic_network.copy(name='CriticNetwork2') # Do not use target_critic_network_2 if critic_network_2 is None. target_critic_network_2 = None self._target_critic_network_2 = ( common.maybe_copy_target_network_with_checks( self._critic_network_2, target_critic_network_2, 'TargetCriticNetwork2')) self._actor_optimizer = actor_optimizer self._critic_optimizer = critic_optimizer self._boltzmann_temperature = boltzmann_temperature self._epsilon_greedy = epsilon_greedy self._replay_buffer_action_spec = replay_buffer_action_spec self._exploration_noise_std = exploration_noise_std self._target_update_tau = target_update_tau self._target_update_period = target_update_period self._actor_update_period = actor_update_period self._dqda_clipping = dqda_clipping self._td_errors_loss_fn = (td_errors_loss_fn or common.element_wise_huber_loss) self._gamma = gamma self._reward_scale_factor = reward_scale_factor self._target_policy_noise = target_policy_noise self._target_policy_noise_clip = target_policy_noise_clip self._gradient_clipping = gradient_clipping self._update_target = self._get_target_updater(target_update_tau, target_update_period) # policy = actor_policy.ActorPolicy( # time_step_spec=time_step_spec, action_spec=action_spec, # actor_network=self._actor_network, clip=True) # collect_policy = actor_policy.ActorPolicy( # time_step_spec=time_step_spec, action_spec=action_spec, # actor_network=self._actor_network, clip=False) # collect_policy = gaussian_policy.GaussianPolicy( # collect_policy, # scale=self._exploration_noise_std, # clip=True) policy = q_policy.QPolicy(time_step_spec, replay_buffer_action_spec, q_network=self._actor_network, emit_log_probability=emit_log_probability) policy._clip = False collect_policy = epsilon_discrete_boltzmann_policy.EpsilonDiscreteBoltzmannPolicy( policy, epsilon=self._epsilon_greedy, env_action_spec=env_action_spec) collect_policy = discrete_boltzmann_policy.DiscreteBoltzmannPolicy( policy, temperature=self._boltzmann_temperature) if boltzmann_temperature is not None: policy = discrete_boltzmann_policy.DiscreteBoltzmannPolicy( policy, temperature=self._boltzmann_temperature / 0.1) # collect_policy = policy # if boltzmann_temperature is not None: # collect_policy = discrete_boltzmann_policy.DiscreteBoltzmannPolicy( # policy, temperature=self._boltzmann_temperature) # else: # collect_policy = epsilon_greedy_policy.EpsilonGreedyPolicy( # policy, epsilon=self._epsilon_greedy) # policy = greedy_policy.GreedyPolicy(policy) # policy = collect_policy super(Td3DiscreteAgent, self).__init__(time_step_spec, env_action_spec, policy, collect_policy, train_sequence_length=2 if not self._actor_network.state_spec else None, debug_summaries=debug_summaries, summarize_grads_and_vars=summarize_grads_and_vars, train_step_counter=train_step_counter)