[RLlib] Curiosity exploration module: tf/tf2.x/tf-eager support. (#11945)

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
 
 

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(

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."""

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

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

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)

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:

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

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:

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