diff --git a/rllib/agents/ddpg/ddpg_tf_policy.py b/rllib/agents/ddpg/ddpg_tf_policy.py index 6b092ed38..414910cc3 100644 --- a/rllib/agents/ddpg/ddpg_tf_policy.py +++ b/rllib/agents/ddpg/ddpg_tf_policy.py @@ -245,6 +245,8 @@ def make_ddpg_optimizers(policy, config): learning_rate=config["actor_lr"]) policy._critic_optimizer = tf1.train.AdamOptimizer( learning_rate=config["critic_lr"]) + # TODO: (sven) make this function return both optimizers and + # TFPolicy handle optimizers vs loss terms correctly (like torch). return None diff --git a/rllib/execution/multi_gpu_impl.py b/rllib/execution/multi_gpu_impl.py index 010ee477c..340fa8a84 100644 --- a/rllib/execution/multi_gpu_impl.py +++ b/rllib/execution/multi_gpu_impl.py @@ -62,7 +62,7 @@ class LocalSyncParallelOptimizer: self.loss_inputs = input_placeholders + rnn_inputs self.build_graph = build_graph - # First initialize the shared loss network + # First initialize the shared loss network. with tf1.name_scope(TOWER_SCOPE_NAME): self._shared_loss = build_graph(self.loss_inputs) shared_ops = tf1.get_collection( diff --git a/rllib/policy/eager_tf_policy.py b/rllib/policy/eager_tf_policy.py index 7d92f6500..dda13e563 100644 --- a/rllib/policy/eager_tf_policy.py +++ b/rllib/policy/eager_tf_policy.py @@ -11,7 +11,7 @@ from ray.rllib.models.repeated_values import RepeatedValues from ray.rllib.policy.policy import Policy, LEARNER_STATS_KEY from ray.rllib.policy.rnn_sequencing import pad_batch_to_sequences_of_same_size from ray.rllib.policy.sample_batch import SampleBatch -from ray.rllib.utils import add_mixins +from ray.rllib.utils import add_mixins, force_list from ray.rllib.utils.annotations import override from ray.rllib.utils.framework import try_import_tf from ray.rllib.utils.tf_ops import convert_to_non_tf_type @@ -278,9 +278,16 @@ def build_eager_tf_policy(name, self._loss_initialized = True if optimizer_fn: - self._optimizer = optimizer_fn(self, config) + optimizers = optimizer_fn(self, config) else: - self._optimizer = tf.keras.optimizers.Adam(config["lr"]) + optimizers = tf.keras.optimizers.Adam(config["lr"]) + optimizers = force_list(optimizers) + if getattr(self, "exploration", None): + optimizers = self.exploration.get_exploration_optimizer( + optimizers) + # TODO: (sven) Allow tf policy to have more than 1 optimizer. + # Just like torch Policy does. + self._optimizer = optimizers[0] if optimizers else None if after_init: after_init(self, observation_space, action_space, config) @@ -577,7 +584,8 @@ def build_eager_tf_policy(name, if apply_gradients_fn: apply_gradients_fn(self, self._optimizer, grads_and_vars) else: - self._optimizer.apply_gradients(grads_and_vars) + self._optimizer.apply_gradients( + [(g, v) for g, v in grads_and_vars if g is not None]) def _compute_gradients(self, samples): """Computes and returns grads as eager tensors.""" diff --git a/rllib/policy/tf_policy_template.py b/rllib/policy/tf_policy_template.py index 9de69a6ca..11f2680bc 100644 --- a/rllib/policy/tf_policy_template.py +++ b/rllib/policy/tf_policy_template.py @@ -8,11 +8,14 @@ from ray.rllib.policy import eager_tf_policy from ray.rllib.policy.policy import Policy, LEARNER_STATS_KEY from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.policy.tf_policy import TFPolicy -from ray.rllib.utils import add_mixins +from ray.rllib.utils import add_mixins, force_list from ray.rllib.utils.annotations import override, DeveloperAPI +from ray.rllib.utils.framework import try_import_tf from ray.rllib.utils.typing import AgentID, ModelGradients, TensorType, \ TrainerConfigDict +tf1, tf, tfv = try_import_tf() + @DeveloperAPI def build_tf_policy( @@ -244,9 +247,16 @@ def build_tf_policy( @override(TFPolicy) def optimizer(self): if optimizer_fn: - return optimizer_fn(self, self.config) + optimizers = optimizer_fn(self, self.config) else: - return base.optimizer(self) + optimizers = base.optimizer(self) + optimizers = force_list(optimizers) + if getattr(self, "exploration", None): + optimizers = self.exploration.get_exploration_optimizer( + optimizers) + # TODO: (sven) Allow tf-eager policy to have more than 1 optimizer. + # Just like torch Policy does. + return optimizers[0] if optimizers else None @override(TFPolicy) def gradients(self, optimizer, loss): diff --git a/rllib/policy/torch_policy_template.py b/rllib/policy/torch_policy_template.py index 22feb0e10..c1320b9d4 100644 --- a/rllib/policy/torch_policy_template.py +++ b/rllib/policy/torch_policy_template.py @@ -320,7 +320,7 @@ def build_torch_policy( else: optimizers = TorchPolicy.optimizer(self) optimizers = force_list(optimizers) - if hasattr(self, "exploration"): + if getattr(self, "exploration", None): optimizers = self.exploration.get_exploration_optimizer( optimizers) return optimizers diff --git a/rllib/utils/exploration/curiosity.py b/rllib/utils/exploration/curiosity.py index 87f77093f..175ce8019 100644 --- a/rllib/utils/exploration/curiosity.py +++ b/rllib/utils/exploration/curiosity.py @@ -4,18 +4,22 @@ from typing import Optional, Tuple, Union from ray.rllib.models.action_dist import ActionDistribution from ray.rllib.models.catalog import ModelCatalog -from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.modelv2 import ModelV2, NullContextManager +from ray.rllib.models.tf.tf_action_dist import Categorical, MultiCategorical from ray.rllib.models.torch.misc import SlimFC from ray.rllib.models.torch.torch_action_dist import TorchCategorical, \ TorchMultiCategorical from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.utils.annotations import override from ray.rllib.utils.exploration.exploration import Exploration -from ray.rllib.utils.framework import try_import_torch, TensorType +from ray.rllib.utils.framework import get_activation_fn, try_import_tf, \ + try_import_torch from ray.rllib.utils.from_config import from_config +from ray.rllib.utils.tf_ops import get_placeholder, one_hot as tf_one_hot from ray.rllib.utils.torch_ops import one_hot -from ray.rllib.utils.typing import FromConfigSpec, ModelConfigDict +from ray.rllib.utils.typing import FromConfigSpec, ModelConfigDict, TensorType +tf1, tf, tfv = try_import_tf() torch, nn = try_import_torch() F = None if nn is not None: @@ -91,9 +95,7 @@ class Curiosity(Exploration): DQN). If None, uses the FromSpecDict provided in the Policy's default config. """ - if framework != "torch": - raise ValueError("Only torch is currently supported for Curiosity") - elif not isinstance(action_space, (Discrete, MultiDiscrete)): + if not isinstance(action_space, (Discrete, MultiDiscrete)): raise ValueError( "Only (Multi)Discrete action spaces supported for Curiosity " "so far!") @@ -139,12 +141,15 @@ class Curiosity(Exploration): self._curiosity_inverse_fcnet = self._create_fc_net( [2 * self.feature_dim] + list(self.inverse_net_hiddens) + - [self.action_dim], self.inverse_net_activation) + [self.action_dim], + self.inverse_net_activation, + name="inverse_net") self._curiosity_forward_fcnet = self._create_fc_net( [self.feature_dim + self.action_dim] + list( self.forward_net_hiddens) + [self.feature_dim], - self.forward_net_activation) + self.forward_net_activation, + name="forward_net") # This is only used to select the correct action self.exploration_submodule = from_config( @@ -172,26 +177,48 @@ class Curiosity(Exploration): @override(Exploration) def get_exploration_optimizer(self, optimizers): - feature_params = list(self._curiosity_feature_net.parameters()) - inverse_params = list(self._curiosity_inverse_fcnet.parameters()) - forward_params = list(self._curiosity_forward_fcnet.parameters()) - - # Now that the Policy's own optimizer(s) have been created (from - # the Model parameters (IMPORTANT: w/o(!) the curiosity params), - # we can add our curiosity sub-modules to the Policy's Model. - self.model._curiosity_feature_net = \ - self._curiosity_feature_net.to(self.device) - self.model._curiosity_inverse_fcnet = \ - self._curiosity_inverse_fcnet.to(self.device) - self.model._curiosity_forward_fcnet = \ - self._curiosity_forward_fcnet.to(self.device) - # Create, but don't add Adam for curiosity NN updating to the policy. # If we added and returned it here, it would be used in the policy's # update loop, which we don't want (curiosity updating happens inside # `postprocess_trajectory`). - self._optimizer = torch.optim.Adam( - forward_params + inverse_params + feature_params, lr=self.lr) + if self.framework == "torch": + feature_params = list(self._curiosity_feature_net.parameters()) + inverse_params = list(self._curiosity_inverse_fcnet.parameters()) + forward_params = list(self._curiosity_forward_fcnet.parameters()) + + # Now that the Policy's own optimizer(s) have been created (from + # the Model parameters (IMPORTANT: w/o(!) the curiosity params), + # we can add our curiosity sub-modules to the Policy's Model. + self.model._curiosity_feature_net = \ + self._curiosity_feature_net.to(self.device) + self.model._curiosity_inverse_fcnet = \ + self._curiosity_inverse_fcnet.to(self.device) + self.model._curiosity_forward_fcnet = \ + self._curiosity_forward_fcnet.to(self.device) + self._optimizer = torch.optim.Adam( + forward_params + inverse_params + feature_params, lr=self.lr) + else: + self.model._curiosity_feature_net = self._curiosity_feature_net + self.model._curiosity_inverse_fcnet = self._curiosity_inverse_fcnet + self.model._curiosity_forward_fcnet = self._curiosity_forward_fcnet + # Feature net is a RLlib ModelV2, the other 2 are keras Models. + self._optimizer_var_list = \ + self._curiosity_feature_net.base_model.variables + \ + self._curiosity_inverse_fcnet.variables + \ + self._curiosity_forward_fcnet.variables + self.model.register_variables(self._optimizer_var_list) + self._optimizer = tf1.train.AdamOptimizer(learning_rate=self.lr) + # Create placeholders and initialize the loss. + if self.framework == "tf": + self._obs_ph = get_placeholder( + space=self.model.obs_space, name="_curiosity_obs") + self._next_obs_ph = get_placeholder( + space=self.model.obs_space, name="_curiosity_next_obs") + self._action_ph = get_placeholder( + space=self.model.action_space, name="_curiosity_action") + self._forward_l2_norm_sqared, self._update_op = \ + self._postprocess_helper_tf( + self._obs_ph, self._next_obs_ph, self._action_ph) return optimizers @@ -202,6 +229,87 @@ class Curiosity(Exploration): Also calculates forward and inverse losses and updates the curiosity module on the provided batch using our optimizer. """ + if self.framework != "torch": + self._postprocess_tf(policy, sample_batch, tf_sess) + else: + self._postprocess_torch(policy, sample_batch) + + def _postprocess_tf(self, policy, sample_batch, tf_sess): + # tf1 static-graph: Perform session call on our loss and update ops. + if self.framework == "tf": + forward_l2_norm_sqared, _ = tf_sess.run( + [self._forward_l2_norm_sqared, self._update_op], + feed_dict={ + self._obs_ph: sample_batch[SampleBatch.OBS], + self._next_obs_ph: sample_batch[SampleBatch.NEXT_OBS], + self._action_ph: sample_batch[SampleBatch.ACTIONS], + }) + # tf-eager: Perform model calls, loss calculations, and optimizer + # stepping on the fly. + else: + forward_l2_norm_sqared, _ = self._postprocess_helper_tf( + sample_batch[SampleBatch.OBS], + sample_batch[SampleBatch.NEXT_OBS], + sample_batch[SampleBatch.ACTIONS], + ) + # Scale intrinsic reward by eta hyper-parameter. + sample_batch[SampleBatch.REWARDS] = \ + sample_batch[SampleBatch.REWARDS] + \ + self.eta * forward_l2_norm_sqared + + return sample_batch + + def _postprocess_helper_tf(self, obs, next_obs, actions): + with (tf.GradientTape() + if self.framework != "tf" else NullContextManager()) as tape: + # Push both observations through feature net to get both phis. + phis, _ = self.model._curiosity_feature_net({ + SampleBatch.OBS: tf.concat([obs, next_obs], axis=0) + }) + phi, next_phi = tf.split(phis, 2) + + # Predict next phi with forward model. + predicted_next_phi = self.model._curiosity_forward_fcnet( + tf.concat( + [phi, tf_one_hot(actions, self.action_space)], axis=-1)) + + # Forward loss term (predicted phi', given phi and action vs + # actually observed phi'). + forward_l2_norm_sqared = 0.5 * tf.reduce_sum( + tf.square(predicted_next_phi - next_phi), axis=-1) + forward_loss = tf.reduce_mean(forward_l2_norm_sqared) + + # Inverse loss term (prediced action that led from phi to phi' vs + # actual action taken). + phi_cat_next_phi = tf.concat([phi, next_phi], axis=-1) + dist_inputs = self.model._curiosity_inverse_fcnet(phi_cat_next_phi) + action_dist = Categorical(dist_inputs, self.model) if \ + isinstance(self.action_space, Discrete) else \ + MultiCategorical( + dist_inputs, self.model, self.action_space.nvec) + # Neg log(p); p=probability of observed action given the inverse-NN + # predicted action distribution. + inverse_loss = -action_dist.logp(actions) + inverse_loss = tf.reduce_mean(inverse_loss) + + # Calculate the ICM loss. + loss = (1.0 - self.beta) * inverse_loss + self.beta * forward_loss + + # Step the optimizer. + if self.framework != "tf": + grads = tape.gradient(loss, self._optimizer_var_list) + grads_and_vars = [(g, v) + for g, v in zip(grads, self._optimizer_var_list) + if g is not None] + update_op = self._optimizer.apply_gradients(grads_and_vars) + else: + update_op = self._optimizer.minimize( + loss, var_list=self._optimizer_var_list) + + # Return the squared l2 norm and the optimizer update op. + return forward_l2_norm_sqared, update_op + + def _postprocess_torch(self, policy, sample_batch): # Push both observations through feature net to get both phis. phis, _ = self.model._curiosity_feature_net({ SampleBatch.OBS: torch.cat([ @@ -253,7 +361,7 @@ class Curiosity(Exploration): # Return the postprocessed sample batch (with the corrected rewards). return sample_batch - def _create_fc_net(self, layer_dims, activation): + def _create_fc_net(self, layer_dims, activation, name=None): """Given a list of layer dimensions (incl. input-dim), creates FC-net. Args: @@ -262,16 +370,32 @@ class Curiosity(Exploration): activation (str): An activation specifier string (e.g. "relu"). Examples: - If layer_dims is [4,8,6] we'll have a two layer net: 4->8 and 8->6, - where the second layer does not have an activation anymore. + If layer_dims is [4,8,6] we'll have a two layer net: 4->8 (8 nodes) + and 8->6 (6 nodes), where the second layer (6 nodes) does not have + an activation anymore. 4 is the input dimension. """ - layers = [] + layers = [ + tf.keras.layers.Input( + shape=(layer_dims[0], ), name="{}_in".format(name)) + ] if self.framework != "torch" else [] + for i in range(len(layer_dims) - 1): act = activation if i < len(layer_dims) - 2 else None - layers.append( - SlimFC( - in_size=layer_dims[i], - out_size=layer_dims[i + 1], - initializer=torch.nn.init.xavier_uniform_, - activation_fn=act)) - return nn.Sequential(*layers) + if self.framework == "torch": + layers.append( + SlimFC( + in_size=layer_dims[i], + out_size=layer_dims[i + 1], + initializer=torch.nn.init.xavier_uniform_, + activation_fn=act)) + else: + layers.append( + tf.keras.layers.Dense( + units=layer_dims[i + 1], + activation=get_activation_fn(act), + name="{}_{}".format(name, i))) + + if self.framework == "torch": + return nn.Sequential(*layers) + else: + return tf.keras.Sequential(layers) diff --git a/rllib/utils/exploration/exploration.py b/rllib/utils/exploration/exploration.py index 4b88fbf97..ab457ece8 100644 --- a/rllib/utils/exploration/exploration.py +++ b/rllib/utils/exploration/exploration.py @@ -157,7 +157,8 @@ class Exploration: return sample_batch @DeveloperAPI - def get_exploration_optimizer(self, optimizers: List[LocalOptimizer]): + def get_exploration_optimizer(self, optimizers: List[LocalOptimizer]) -> \ + List[LocalOptimizer]: """May add optimizer(s) to the Policy's own `optimizers`. The number of optimizers (Policy's plus Exploration's optimizers) must @@ -176,7 +177,7 @@ class Exploration: @DeveloperAPI def get_exploration_loss(self, policy_loss: List[TensorType], - train_batch: SampleBatch): + train_batch: SampleBatch) -> List[TensorType]: """May add loss term(s) to the Policy's own loss(es). Args: diff --git a/rllib/utils/exploration/tests/test_curiosity.py b/rllib/utils/exploration/tests/test_curiosity.py index 8a2d4da18..ca7ab9011 100644 --- a/rllib/utils/exploration/tests/test_curiosity.py +++ b/rllib/utils/exploration/tests/test_curiosity.py @@ -24,7 +24,7 @@ class MyCallBack(DefaultCallbacks): policy_id, policies, postprocessed_batch, original_batches, **kwargs): pos = np.argmax(postprocessed_batch["obs"], -1) - x, y = pos % 10, pos // 10 + x, y = pos % 8, pos // 8 self.deltas.extend((x**2 + y**2)**0.5) def on_sample_end(self, *, worker, samples, **kwargs): @@ -125,23 +125,21 @@ class TestCuriosity(unittest.TestCase): def tearDownClass(cls): ray.shutdown() - def test_curiosity_on_large_frozen_lake(self): + def test_curiosity_on_frozen_lake(self): config = ppo.DEFAULT_CONFIG.copy() - # A very large frozen-lake that's hard for a random policy to solve + # A very large frozen-lake that's hard for a random policy to solve # due to 0.0 feedback. config["env"] = "FrozenLake-v0" config["env_config"] = { "desc": [ - "SFFFFFFFFF", - "FFFFFFFFFF", - "FFFFFFFFFF", - "FFFFFFFFFF", - "FFFFFFFFFF", - "FFFFFFFFFF", - "FFFFFFFFFF", - "FFFFFFFFFF", - "FFFFFFFFFF", - "FFFFFFFFFG", + "SFFFFFFF", + "FFFFFFFF", + "FFFFFFFF", + "FFFFFFFF", + "FFFFFFFF", + "FFFFFFFF", + "FFFFFFFF", + "FFFFFFFG", ], "is_slippery": False } @@ -149,13 +147,13 @@ class TestCuriosity(unittest.TestCase): config["callbacks"] = MyCallBack # Limit horizon to make it really hard for non-curious agent to reach # the goal state. - config["horizon"] = 23 + config["horizon"] = 16 # Local only. config["num_workers"] = 0 config["lr"] = 0.001 num_iterations = 10 - for _ in framework_iterator(config, frameworks="torch"): + for fw in framework_iterator(config): # W/ Curiosity. Expect to learn something. config["exploration_config"] = { "type": "Curiosity", @@ -182,18 +180,21 @@ class TestCuriosity(unittest.TestCase): trainer.stop() self.assertTrue(learnt) - # W/o Curiosity. Expect to learn nothing. - config["exploration_config"] = { - "type": "StochasticSampling", - } - trainer = ppo.PPOTrainer(config=config) - rewards_wo = 0.0 - for _ in range(num_iterations): - result = trainer.train() - rewards_wo += result["episode_reward_mean"] - print(result) - trainer.stop() - self.assertTrue(rewards_wo == 0.0) + if fw == "tf": + # W/o Curiosity. Expect to learn nothing. + print("Trying w/o curiosity (not expected to learn).") + config["exploration_config"] = { + "type": "StochasticSampling", + } + trainer = ppo.PPOTrainer(config=config) + rewards_wo = 0.0 + for _ in range(num_iterations): + result = trainer.train() + rewards_wo += result["episode_reward_mean"] + print(result) + trainer.stop() + self.assertTrue(rewards_wo == 0.0) + print("Did not reach goal w/o curiosity!") def test_curiosity_on_partially_observable_domain(self): config = ppo.DEFAULT_CONFIG.copy() diff --git a/rllib/utils/framework.py b/rllib/utils/framework.py index 3595399d6..33e8bbffc 100644 --- a/rllib/utils/framework.py +++ b/rllib/utils/framework.py @@ -225,11 +225,12 @@ def get_variable(value, # TODO: (sven) move to models/utils.py -def get_activation_fn(name, framework="tf"): +def get_activation_fn(name: Optional[str] = None, framework: str = "tf"): """Returns a framework specific activation function, given a name string. Args: - name (str): One of "relu" (default), "tanh", "swish", or "linear". + name (Optional[str]): One of "relu" (default), "tanh", "swish", or + "linear" or None. framework (str): One of "tf" or "torch". Returns: diff --git a/rllib/utils/tf_ops.py b/rllib/utils/tf_ops.py index d0344d123..f74926aaa 100644 --- a/rllib/utils/tf_ops.py +++ b/rllib/utils/tf_ops.py @@ -132,7 +132,10 @@ def make_tf_callable(session_or_none, dynamic_shape=False): assert session_or_none is not None def make_wrapper(fn): - if session_or_none: + # Static-graph mode: Create placeholders and make a session call each + # time the wrapped function is called. Return this session call's + # outputs. + if session_or_none is not None: args_placeholders = [] kwargs_placeholders = {} symbolic_out = [None] @@ -183,6 +186,7 @@ def make_tf_callable(session_or_none, dynamic_shape=False): return ret return call + # Eager mode (call function as is). else: return fn