[RLlib] Curiosity minor fixes, do-overs, and testing. (#10143)

python/requirements.txt

@@ -46,6 +46,7 @@ dataclasses
 dask[complete]
 feather-format
 gym
+gym-minigrid
 kubernetes
 lxml
 mypy

rllib/BUILD

@@ -1073,6 +1073,13 @@ py_test(
 # Tag: utils
 # --------------------------------------------------------------------
 
+py_test(
+    name = "test_curiosity",
+    tags = ["utils"],
+    size = "large",
+    srcs = ["utils/exploration/tests/test_curiosity.py"]
+)
+
 py_test(
     name = "test_explorations",
     tags = ["utils"],

rllib/agents/ppo/ppo_torch_policy.py

@@ -74,9 +74,7 @@ class PPOLoss:
                 return torch.sum(t[valid_mask]) / num_valid
 
         else:
-
-            def reduce_mean_valid(t):
-                return torch.mean(t)
+            reduce_mean_valid = torch.mean
 
         prev_dist = dist_class(prev_logits, model)
         # Make loss functions.

rllib/agents/trainer.py

@@ -573,7 +573,8 @@ class Trainer(Trainable):
             # Try gym.
             else:
                 import gym  # soft dependency
-                self.env_creator = lambda env_config: gym.make(env)
+                self.env_creator = \
+                    lambda env_config: gym.make(env, **env_config)
         else:
             self.env_creator = lambda env_config: None
 

rllib/evaluation/sample_batch_builder.py

@@ -196,13 +196,14 @@ class MultiAgentSampleBatchBuilder:
                 raise ValueError(
                     "Batches sent to postprocessing must only contain steps "
                     "from a single trajectory.", pre_batch)
-            post_batches[agent_id] = policy.postprocess_trajectory(
-                pre_batch, other_batches, episode)
             # Call the Policy's Exploration's postprocess method.
+            post_batches[agent_id] = pre_batch
             if getattr(policy, "exploration", None) is not None:
                 policy.exploration.postprocess_trajectory(
                     policy, post_batches[agent_id],
                     getattr(policy, "_sess", None))
+            post_batches[agent_id] = policy.postprocess_trajectory(
+                post_batches[agent_id], other_batches, episode)
 
         if log_once("after_post"):
             logger.info(

rllib/models/catalog.py

@@ -37,14 +37,14 @@ logger = logging.getLogger(__name__)
 # __sphinx_doc_begin__
 MODEL_DEFAULTS: ModelConfigDict = {
     # === Built-in options ===
+    # Number of hidden layers for fully connected net
+    "fcnet_hiddens": [256, 256],
+    # Nonlinearity for fully connected net (tanh, relu)
+    "fcnet_activation": "tanh",
     # Filter config. List of [out_channels, kernel, stride] for each filter
     "conv_filters": None,
     # Nonlinearity for built-in convnet
     "conv_activation": "relu",
-    # Nonlinearity for fully connected net (tanh, relu)
-    "fcnet_activation": "tanh",
-    # Number of hidden layers for fully connected net
-    "fcnet_hiddens": [256, 256],
     # For DiagGaussian action distributions, make the second half of the model
     # outputs floating bias variables instead of state-dependent. This only
     # has an effect is using the default fully connected net.

rllib/models/torch/fcnet.py

@@ -5,7 +5,7 @@ from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
 from ray.rllib.models.torch.misc import SlimFC, AppendBiasLayer, \
     normc_initializer
 from ray.rllib.utils.annotations import override
-from ray.rllib.utils.framework import get_activation_fn, try_import_torch
+from ray.rllib.utils.framework import try_import_torch
 
 torch, nn = try_import_torch()
 
@@ -21,8 +21,7 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module):
                               model_config, name)
         nn.Module.__init__(self)
 
-        activation = get_activation_fn(
-            model_config.get("fcnet_activation"), framework="torch")
+        activation = model_config.get("fcnet_activation")
         hiddens = model_config.get("fcnet_hiddens")
         no_final_linear = model_config.get("no_final_linear")
         self.vf_share_layers = model_config.get("vf_share_layers")

rllib/models/torch/misc.py

@@ -68,20 +68,26 @@ class SlimConv2d(nn.Module):
             bias_init=0):
         super(SlimConv2d, self).__init__()
         layers = []
+        # Padding layer.
         if padding:
             layers.append(nn.ZeroPad2d(padding))
+        # Actual Conv2D layer (including correct initialization logic).
         conv = nn.Conv2d(in_channels, out_channels, kernel, stride)
         if initializer:
             if initializer == "default":
                 initializer = nn.init.xavier_uniform_
             initializer(conv.weight)
         nn.init.constant_(conv.bias, bias_init)
-
         layers.append(conv)
-        if activation_fn:
+        # Activation function (if any; default=ReLu).
+        if isinstance(activation_fn, str):
             if activation_fn == "default":
                 activation_fn = nn.ReLU
+            else:
+                activation_fn = get_activation_fn(activation_fn, "torch")
+        if activation_fn is not None:
             layers.append(activation_fn())
+        # Put everything in sequence.
         self._model = nn.Sequential(*layers)
 
     def forward(self, x):
@@ -100,16 +106,19 @@ class SlimFC(nn.Module):
                  bias_init=0.0):
         super(SlimFC, self).__init__()
         layers = []
+        # Actual Conv2D layer (including correct initialization logic).
         linear = nn.Linear(in_size, out_size, bias=use_bias)
         if initializer:
             initializer(linear.weight)
         if use_bias is True:
             nn.init.constant_(linear.bias, bias_init)
         layers.append(linear)
+        # Activation function (if any; default=None (linear)).
         if isinstance(activation_fn, str):
             activation_fn = get_activation_fn(activation_fn, "torch")
         if activation_fn is not None:
             layers.append(activation_fn())
+        # Put everything in sequence.
         self._model = nn.Sequential(*layers)
 
     def forward(self, x):

rllib/models/torch/recurrent_net.py

@@ -139,7 +139,7 @@ class LSTMWrapper(RecurrentNetwork, nn.Module):
                     wrapped_out,
                     torch.reshape(input_dict[SampleBatch.PREV_ACTIONS].float(),
                                   [-1, self.action_dim]),
-                    torch.reshape(input_dict[SampleBatch.PREV_REWARDS],
+                    torch.reshape(input_dict[SampleBatch.PREV_REWARDS].float(),
                                   [-1, 1]),
                 ],
                 dim=1)

rllib/models/torch/visionnet.py

@@ -5,7 +5,7 @@ from ray.rllib.models.torch.misc import normc_initializer, same_padding, \
     SlimConv2d, SlimFC
 from ray.rllib.models.tf.visionnet_v1 import _get_filter_config
 from ray.rllib.utils.annotations import override
-from ray.rllib.utils.framework import get_activation_fn, try_import_torch
+from ray.rllib.utils.framework import try_import_torch
 
 _, nn = try_import_torch()
 
@@ -22,8 +22,7 @@ class VisionNetwork(TorchModelV2, nn.Module):
                               model_config, name)
         nn.Module.__init__(self)
 
-        activation = get_activation_fn(
-            self.model_config.get("conv_activation"), framework="torch")
+        activation = self.model_config.get("conv_activation")
         filters = self.model_config["conv_filters"]
         no_final_linear = self.model_config.get("no_final_linear")
         vf_share_layers = self.model_config.get("vf_share_layers")
@@ -101,7 +100,10 @@ class VisionNetwork(TorchModelV2, nn.Module):
         self._value_branch_separate = self._value_branch = None
         if vf_share_layers:
             self._value_branch = SlimFC(
-                out_channels, 1, initializer=normc_initializer(0.01))
+                out_channels,
+                1,
+                initializer=normc_initializer(0.01),
+                activation_fn=None)
         else:
             vf_layers = []
             (w, h, in_channels) = obs_space.shape
@@ -136,7 +138,8 @@ class VisionNetwork(TorchModelV2, nn.Module):
                     out_channels=1,
                     kernel=1,
                     stride=1,
-                    padding=None))
+                    padding=None,
+                    activation_fn=None))
             self._value_branch_separate = nn.Sequential(*vf_layers)
 
         # Holds the current "base" output (before logits layer).

rllib/policy/torch_policy.py

@@ -2,7 +2,7 @@ import functools
 import gym
 import numpy as np
 import time
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable, Dict, List, Optional, Tuple, Type, Union
 
 import ray
 from ray.rllib.models.modelv2 import ModelV2
@@ -47,15 +47,21 @@ class TorchPolicy(Policy):
             config: TrainerConfigDict,
             *,
             model: ModelV2,
-            loss: Callable[[Policy, ModelV2, type, SampleBatch], TensorType],
-            action_distribution_class: TorchDistributionWrapper,
-            action_sampler_fn: Callable[[TensorType, List[TensorType]], Tuple[
-                TensorType, TensorType]] = None,
+            loss: Callable[[
+                Policy, ModelV2, Type[TorchDistributionWrapper], SampleBatch
+            ], Union[TensorType, List[TensorType]]],
+            action_distribution_class: Type[TorchDistributionWrapper],
+            action_sampler_fn: Optional[Callable[[
+                TensorType, List[TensorType]
+            ], Tuple[TensorType, TensorType]]] = None,
             action_distribution_fn: Optional[Callable[[
                 Policy, ModelV2, TensorType, TensorType, TensorType
-            ], Tuple[TensorType, type, List[TensorType]]]] = None,
+            ], Tuple[TensorType, Type[TorchDistributionWrapper], List[
+                TensorType]]]] = None,
             max_seq_len: int = 20,
-            get_batch_divisibility_req: Optional[int] = None):
+            get_batch_divisibility_req: Optional[Callable[[Policy],
+                                                          int]] = None,
+    ):
         """Build a policy from policy and loss torch modules.
 
         Note that model will be placed on GPU device if CUDA_VISIBLE_DEVICES
@@ -69,11 +75,11 @@ class TorchPolicy(Policy):
             model (ModelV2): PyTorch policy module. Given observations as
                 input, this module must return a list of outputs where the
                 first item is action logits, and the rest can be any value.
-            loss (Callable[[Policy, ModelV2, type, SampleBatch], TensorType]):
-                Function that takes (policy, model, dist_class, train_batch)
-                and returns a single scalar loss.
-            action_distribution_class (TorchDistributionWrapper): Class for
-                a torch action distribution.
+            loss (Callable[[Policy, ModelV2, Type[TorchDistributionWrapper],
+                SampleBatch], Union[TensorType, List[TensorType]]]): Callable
+                that returns a single scalar loss or a list of loss terms.
+            action_distribution_class (Type[TorchDistributionWrapper]): Class
+                for a torch action distribution.
             action_sampler_fn (Callable[[TensorType, List[TensorType]],
                 Tuple[TensorType, TensorType]]): A callable returning a
                 sampled action and its log-likelihood given Policy, ModelV2,
@@ -337,15 +343,21 @@ class TorchPolicy(Policy):
             batch_divisibility_req=self.batch_divisibility_req)
 
         train_batch = self._lazy_tensor_dict(postprocessed_batch)
+
+        # Calculate the actual policy loss.
         loss_out = force_list(
             self._loss(self, self.model, self.dist_class, train_batch))
+
         # Call Model's custom-loss with Policy loss outputs and train_batch.
         if self.model:
             loss_out = self.model.custom_loss(loss_out, train_batch)
-        # Modifies the loss as specified by the Exploration strategy.
+
+        # Give Exploration component that chance to modify the loss (or add
+        # its own terms).
         if hasattr(self, "exploration"):
             loss_out = self.exploration.get_exploration_loss(
                 loss_out, train_batch)
+
         assert len(loss_out) == len(self._optimizers)
         # assert not any(torch.isnan(l) for l in loss_out)
         fetches = self.extra_compute_grad_fetches()

rllib/policy/torch_policy_template.py

@@ -298,9 +298,8 @@ def build_torch_policy(
                 optimizers = TorchPolicy.optimizer(self)
             optimizers = force_list(optimizers)
             if hasattr(self, "exploration"):
-                exploration_optimizers = force_list(
-                    self.exploration.get_exploration_optimizer(self.config))
-                optimizers.extend(exploration_optimizers)
+                optimizers = self.exploration.get_exploration_optimizer(
+                    optimizers)
             return optimizers
 
         @override(TorchPolicy)

rllib/utils/exploration/curiosity.py

@@ -1,268 +1,286 @@
-"""
-Curiosity-driven Exploration by Self-supervised Prediction - Pathak, Agrawal,
-Efros, and Darrell - UC Berkeley - ICML 2017.
-
-This implements the curiosty-based loss function from
-https://arxiv.org/pdf/1705.05363.pdf. We learn a simplified model of the
-environment based on three networks:
-    1) embedding states into latent space (the "features" network)
-    2) predicting the next embedded state, given a state and action (the
-        "forwards" network)
-    3) predicting the action, given two consecutive embedded state (the
-        "inverse" network)
-
-If the agent was unable to successfully predict the state-action-next_state
-sequence, we modify the standard reward with a penalty. Therefore, if a state
-transition was unexpected, the agent becomes "curious" and further explores
-this transition.
-
-This is tailored for sparse reward environments, as it generates an intrinsic
-reward.
-"""
-from gym.spaces import Space
-from typing import Union, Optional
+from gym.spaces import Discrete, Space
+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.torch.misc import SlimFC
+from ray.rllib.models.torch.torch_action_dist import TorchCategorical
+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.from_config import from_config
-from ray.rllib.utils.typing import SampleBatchType, TrainerConfigDict
+from ray.rllib.utils.typing import FromConfigSpec, ModelConfigDict, \
+    SampleBatchType
 
 torch, nn = try_import_torch()
-
-# TODO: (tanay) how to test if action space is discrete
-"""
-Example Configuration
-
-config = ppo.DEFAULT_CONFIG
-env = "CartPole-v0"
-config["framework"] = "torch"
-config["exploration_config"] = {
-    "type": "ray.rllib.utils.exploration.curiosity_exploration.Curiosity",
-    "forward_net_hiddens": [64],
-    "inverse_net_hiddens": [32,4],
-    "feature_net_hiddens": [16,8],
-    "feature_dim": 8,
-    "forward_activation": "relu",
-    "inverse_activation": "relu",
-    "feature_activation": "relu",
-    "submodule": "EpsilonGreedy",
-}
-trainer = ppo.PPOTrainer(config=config, env=env)
-trainer.train()
-"""
+F = None
+if nn is not None:
+    F = nn.functional
 
 
 class Curiosity(Exploration):
-    def __init__(self, action_space: Space, *, framework: str, **kwargs):
-        """
+    """Implementation of:
+    [1] Curiosity-driven Exploration by Self-supervised Prediction
+    Pathak, Agrawal, Efros, and Darrell - UC Berkeley - ICML 2017.
+    https://arxiv.org/pdf/1705.05363.pdf
+
+    Learns a simplified model of the environment based on three networks:
+    1) Embedding observations into latent space ("feature" network).
+    2) Predicting the action, given two consecutive embedded observations
+        ("inverse" network).
+    3) Predicting the next embedded obs, given an obs and action
+        ("forward" network).
+
+    The less the agent is able to predict the actually observed next feature
+    vector, given obs and action (through the forwards network), the larger the
+    "intrinsic reward", which will be added to the extrinsic reward.
+    Therefore, if a state transition was unexpected, the agent becomes
+    "curious" and will further explore this transition leading to better
+    exploration in sparse rewards environments.
+    """
+
+    def __init__(self,
+                 action_space: Space,
+                 *,
+                 framework: str,
+                 model: ModelV2,
+                 feature_dim: int = 288,
+                 feature_net_config: Optional[ModelConfigDict] = None,
+                 inverse_net_hiddens: Tuple[int] = (256, ),
+                 inverse_net_activation: str = "relu",
+                 forward_net_hiddens: Tuple[int] = (256, ),
+                 forward_net_activation: str = "relu",
+                 beta: float = 0.2,
+                 eta: float = 1.0,
+                 lr: float = 1e-3,
+                 sub_exploration: Optional[FromConfigSpec] = None,
+                 **kwargs):
+        """Initializes a Curiosity object.
+
+        Uses as defaults the hyperparameters described in [1].
+
         Args:
-            action_space (Space): The action space in which to explore.
-            framework (str): One of "tf" or "torch". Currently only torch is
-                supported.
+             feature_dim (int): The dimensionality of the feature (phi)
+                vectors.
+             feature_net_config (Optional[ModelConfigDict]): Optional model
+                configuration for the feature network, producing feature
+                vectors (phi) from observations. This can be used to configure
+                fcnet- or conv_net setups to properly process any observation
+                space.
+             inverse_net_hiddens (Tuple[int]): Tuple of the layer sizes of the
+                inverse (action predicting) NN head (on top of the feature
+                outputs for phi and phi').
+             inverse_net_activation (str): Activation specifier for the inverse
+                net.
+             forward_net_hiddens (Tuple[int]): Tuple of the layer sizes of the
+                forward (phi' predicting) NN head.
+             forward_net_activation (str): Activation specifier for the forward
+                net.
+             beta (float): Weight for the forward loss (over the inverse loss,
+                which gets weight=1.0-beta) in the common loss term.
+             eta (float): Weight for intrinsic rewards before being added to
+                extrinsic ones.
+             lr (float): The learning rate for the curiosity-specific
+                optimizer, optimizing feature-, inverse-, and forward nets.
+             sub_exploration (Optional[FromConfigSpec]): The config dict for
+                the underlying Exploration to use (e.g. epsilon-greedy for
+                DQN). If None, uses the FromSpecDict provided in the Policy's
+                default config.
         """
         if framework != "torch":
-            raise NotImplementedError("only torch is currently supported for "
-                                      "curiosity")
-
-        # Parse the curiosity-specific arguments
-        # If it was not specified in the config, assign the given default
-        def extract_from_kwargs(key, default):
-            if key in kwargs:
-                temp = kwargs[key]
-                del kwargs[key]
-                return temp
-            else:
-                return default
-
-        # Casts a single int to a list, else leaves it unchanged
-        def cast_to_list(l):
-            if type(l) == int:
-                return [l]
-            else:
-                return l
-
-        submodule_type = extract_from_kwargs("submodule", "StochasticSampling")
-        self.feature_dim = extract_from_kwargs("feature_dim", 32)
-
-        forward_activation = extract_from_kwargs("forward_activation", nn.ReLU)
-        inverse_activation = extract_from_kwargs("inverse_activation", nn.ReLU)
-        feature_activation = extract_from_kwargs("feature_activation", nn.ReLU)
-
-        feature_net_hiddens = cast_to_list(
-            extract_from_kwargs("feature_net_hiddens", [64]))
-        inverse_net_hiddens = cast_to_list(
-            extract_from_kwargs("inverse_net_hiddens", [64]))
-        forward_net_hiddens = cast_to_list(
-            extract_from_kwargs("forward_net_hiddens", [64]))
+            raise ValueError("Only torch is currently supported for Curiosity")
+        elif not isinstance(action_space, Discrete):
+            raise ValueError(
+                "Only Discrete action spaces supported for Curiosity so far.")
 
         super().__init__(
-            action_space=action_space, framework=framework, **kwargs)
+            action_space, model=model, framework=framework, **kwargs)
 
-        # TODO: what should this look like for multidimensional obs spaces
-        self.obs_space_dim = kwargs["model"].obs_space.shape[0]
-        # TODO can we always assume 1
-        self.action_space_dim = 1
+        self.feature_dim = feature_dim
+        if feature_net_config is None:
+            feature_net_config = self.policy_config["model"].copy()
+        self.feature_net_config = feature_net_config
+        self.inverse_net_hiddens = inverse_net_hiddens
+        self.inverse_net_activation = inverse_net_activation
+        self.forward_net_hiddens = forward_net_hiddens
+        self.forward_net_activation = forward_net_activation
 
-        # Given a list of layer dimensions, create a FC ReLU net.
-        # If layer_dims is [4,8,6] we'll have a two layer net: 4->8 and 8->6
-        def create_fc_net(layer_dims, activation):
-            layers = []
-            for i in range(len(layer_dims) - 1):
-                layers.append(
-                    SlimFC(
-                        in_size=layer_dims[i],
-                        out_size=layer_dims[i + 1],
-                        use_bias=False,
-                        activation_fn=activation))
-            return nn.Sequential(*layers)
+        self.beta = beta
+        self.eta = eta
+        self.lr = lr
+        # TODO: (sven) if sub_exploration is None, use Trainer's default
+        #  Exploration config.
+        if sub_exploration is None:
+            raise NotImplementedError
+        self.sub_exploration = sub_exploration
 
-        # List of dimension of each layer. Appends the hidden dims.
-        feature_dims = [self.obs_space_dim
-                        ] + feature_net_hiddens + [self.feature_dim]
-        inverse_dims = [2 * self.feature_dim
-                        ] + inverse_net_hiddens + [self.action_space_dim]
-        forward_dims = [self.feature_dim + self.action_space_dim] + \
-            forward_net_hiddens + [self.feature_dim]
+        # Creates modules/layers inside the actual ModelV2.
+        self._curiosity_feature_net = ModelCatalog.get_model_v2(
+            self.model.obs_space,
+            self.action_space,
+            self.feature_dim,
+            model_config=self.feature_net_config,
+            framework=self.framework,
+            name="feature_net",
+        )
 
-        # Creates actual models
-        self.feature_model = create_fc_net(feature_dims, feature_activation)
-        self.inverse_model = create_fc_net(inverse_dims, inverse_activation)
-        self.forward_model = create_fc_net(forward_dims, forward_activation)
+        self._curiosity_inverse_fcnet = self._create_fc_net(
+            [2 * self.feature_dim] + list(self.inverse_net_hiddens) +
+            [self.action_space.n], self.inverse_net_activation)
 
-        # Convenient reductions
-        self.criterion = torch.nn.MSELoss(reduction="none")
-        self.criterion_reduced = torch.nn.MSELoss(reduction="sum")
+        self._curiosity_forward_fcnet = self._create_fc_net(
+            [self.feature_dim + self.action_space.n
+             ] + list(forward_net_hiddens) + [self.feature_dim],
+            self.forward_net_activation)
 
         # This is only used to select the correct action
         self.exploration_submodule = from_config(
             cls=Exploration,
-            config={
-                "type": submodule_type,
-                "action_space": action_space,
-                "framework": framework,
-                "policy_config": self.policy_config,
-                "model": self.model,
-                "num_workers": self.num_workers,
-                "worker_index": self.worker_index
-            })
+            config=self.sub_exploration,
+            action_space=self.action_space,
+            framework=self.framework,
+            policy_config=self.policy_config,
+            model=self.model,
+            num_workers=self.num_workers,
+            worker_index=self.worker_index,
+        )
 
+    @override(Exploration)
     def get_exploration_action(self,
                                *,
                                action_distribution: ActionDistribution,
                                timestep: Union[int, TensorType],
                                explore: bool = True):
+        # Simply delegate to sub-Exploration module.
+        return self.exploration_submodule.get_exploration_action(
+            action_distribution=action_distribution,
+            timestep=timestep,
+            explore=explore)
+
+    @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)
+
+        # Add the Adam for curiosity NN updating to the Policy's optimizers.
+        return optimizers + [
+            torch.optim.Adam(
+                forward_params + inverse_params + feature_params, lr=self.lr)
+        ]
+
+    @override(Exploration)
+    def postprocess_trajectory(self, policy, sample_batch, tf_sess=None):
+        """Calculates phi values (obs, obs', and predicted obs') and ri.
+
+        Stores calculated phi, phi' and predicted phi' as well as the intrinsic
+        rewards in the batch for loss processing by the policy.
         """
-        Returns the action to take next
+        batch_size = sample_batch[SampleBatch.OBS].shape[0]
+        phis, _ = self.model._curiosity_feature_net({
+            SampleBatch.OBS: torch.cat([
+                torch.from_numpy(sample_batch[SampleBatch.OBS]),
+                torch.from_numpy(sample_batch[SampleBatch.NEXT_OBS])
+            ])
+        })
+        phi, next_phi = phis[:batch_size], phis[batch_size:]
+
+        # Detach phi from graph (should not backpropagate through feature net
+        # for forward-loss).
+        predicted_next_phi = self.model._curiosity_forward_fcnet(
+            torch.cat(
+                [
+                    phi.detach(),
+                    F.one_hot(
+                        torch.from_numpy(
+                            sample_batch[SampleBatch.ACTIONS]).long(),
+                        num_classes=self.action_space.n).float()
+                ],
+                dim=-1))
+
+        # Forward loss term (predicted phi', given phi and action vs actually
+        # observed phi').
+        forward_l2_norm_sqared = 0.5 * torch.sum(
+            torch.pow(predicted_next_phi - next_phi, 2.0), dim=-1)
+        # Scale forward loss by eta hyper-parameter.
+        sample_batch[SampleBatch.REWARDS] = \
+            sample_batch[SampleBatch.REWARDS] + \
+            self.eta * forward_l2_norm_sqared.detach().cpu().numpy()
+        return sample_batch
+
+    @override(Exploration)
+    def get_exploration_loss(self, policy_loss, train_batch: SampleBatchType):
+        """Adds the loss for the inverse and forward models to policy_loss.
+        """
+        batch_size = train_batch[SampleBatch.OBS].shape[0]
+        phis, _ = self.model._curiosity_feature_net({
+            SampleBatch.OBS: torch.cat(
+                [
+                    train_batch[SampleBatch.OBS],
+                    train_batch[SampleBatch.NEXT_OBS]
+                ],
+                dim=0)
+        })
+        phi, next_phi = phis[:batch_size], phis[batch_size:]
+        # Inverse loss term (prediced action that led from phi to phi' vs
+        # actual action taken).
+        phi_next_phi = torch.cat([phi, next_phi], dim=-1)
+        dist_inputs = self.model._curiosity_inverse_fcnet(phi_next_phi)
+        action_dist = TorchCategorical(dist_inputs, self.model)
+        # Neg log(p); p=probability of observed action given the inverse-NN
+        # predicted action distribution.
+        inverse_loss = -action_dist.logp(train_batch[SampleBatch.ACTIONS])
+        inverse_loss = torch.mean(inverse_loss)
+
+        # Forward loss term has already been calculated during train batch pre-
+        # processing (just have to weight with beta here).
+        predicted_next_phi = self.model._curiosity_forward_fcnet(
+            torch.cat(
+                [
+                    phi,
+                    F.one_hot(
+                        train_batch[SampleBatch.ACTIONS].long(),
+                        num_classes=self.action_space.n).float()
+                ],
+                dim=-1))
+        forward_loss = torch.mean(0.5 * torch.sum(
+            torch.pow(predicted_next_phi - next_phi, 2.0), dim=-1))
+
+        # Append our loss to the policy loss(es).
+        return policy_loss + [
+            (1.0 - self.beta) * inverse_loss + self.beta * forward_loss
+        ]
+
+    def _create_fc_net(self, layer_dims, activation):
+        """Given a list of layer dimensions (incl. input-dim), creates FC-net.
 
         Args:
-            action_distribution (ActionDistribution): The probabilistic
-                distribution we sample actions from
-            timestep (Union[int, TensorType]):
-            explore (bool): If true, uses the submodule strategy to select the
-                next action
+            layer_dims (Tuple[int]): Tuple of layer dims, including the input
+                dimension.
+            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.
         """
-        return self.exploration_submodule.get_exploration_action(
-            action_distribution=action_distribution, timestep=timestep)
-
-    def get_exploration_loss(self, policy_loss, sample_batch: SampleBatchType):
-        """
-        Returns the intrinsic reward associated to the explorations strategy
-            policy_loss (TensorType): The loss from the policy, not associated
-                to the exploration strategy, which we will modify
-            sample_batch (SampleBatchType): The SampleBatch of observations, to
-                which we will associate an intrinsic loss.
-        """
-
-        # Cast to torch tensors, to be fed into the model
-        obs_list = sample_batch["obs"].float()
-        next_obs_list = sample_batch["new_obs"].float()
-        emb_next_obs_list = self._get_latent_vector(next_obs_list).float()
-        actions_list = sample_batch["actions"].float()
-
-        actions_pred = self._predict_action(obs_list, next_obs_list)
-        embedding_pred = self._predict_next_obs(obs_list, actions_list)
-
-        # L2 losses for predicted action and next state
-        embedding_loss = self.criterion_reduced(emb_next_obs_list,
-                                                embedding_pred)
-        actions_loss = self.criterion_reduced(
-            actions_pred.squeeze(1), actions_list)
-        return policy_loss + [embedding_loss + actions_loss]
-
-    def _get_latent_vector(self, obs: TensorType) -> TensorType:
-        """
-        Returns the embedded vector phi(state)
-            obs (TensorType): a batch of states
-        """
-        return self.feature_model(obs)
-
-    def get_exploration_optimizers(self, config: TrainerConfigDict):
-        """Returns optimizer (or list) for environmental dynamics networks.
-        """
-        forward_params = list(self.forward_model.parameters())
-        inverse_params = list(self.inverse_model.parameters())
-        feature_params = list(self.feature_model.parameters())
-
-        return torch.optim.Adam(
-            forward_params + inverse_params + feature_params, lr=1e-3)
-
-    def postprocess_trajectory(self,
-                               policy,
-                               sample_batch: SampleBatchType,
-                               tf_sess: Optional["tf.Session"] = None):
-        """Calculates intrinsic rewards and adds them to "rewards" in batch.
-
-        Calculations are based on difference between predicted and actually
-        observed next observations.
-        """
-
-        # Extract the relevant data from the SampleBatch, and cast to Tensors
-        obs_list = torch.from_numpy(sample_batch["obs"]).float()
-        next_obs_list = torch.from_numpy(sample_batch["new_obs"]).float()
-        emb_next_obs_list = self._get_latent_vector(next_obs_list).float()
-        actions_list = torch.from_numpy(sample_batch["actions"]).float()
-
-        # Equation (2) in paper.
-        actions_pred = self._predict_action(obs_list, next_obs_list)
-        embedding_pred = self._predict_next_obs(obs_list, actions_list)
-
-        # A vector of L2 losses corresponding to each observation,
-        # Equation (7) in paper.
-        embedding_loss = torch.sum(
-            self.criterion(emb_next_obs_list, embedding_pred), dim=-1)
-
-        # Equation (3) in paper. TODO discrete action space
-        actions_loss = self.criterion(actions_pred.squeeze(1), actions_list)
-
-        # Modifies environment rewards by subtracting intrinsic rewards
-        sample_batch["rewards"] = sample_batch["rewards"] - \
-            embedding_loss.clone().detach().numpy() - \
-            actions_loss.clone().detach().numpy()
-
-    def _predict_action(self, obs: TensorType, next_obs: TensorType):
-        """
-        Returns the predicted action, given two states. This is the inverse
-        dynamics model.
-
-        obs (TensorType): Observed state at time t.
-        next_obs (TensorType): Observed state at time t+1
-        """
-        return self.inverse_model(
-            torch.cat(
-                (self._get_latent_vector(obs),
-                 self._get_latent_vector(next_obs)),
-                axis=-1))
-
-    # raw obs (not embedded)
-    def _predict_next_obs(self, obs: TensorType, action: TensorType):
-        """
-        Returns the predicted next state, given an action and state.
-
-        obs (TensorType): Observed state at time t.
-        action (TensorType): Action taken at time t
-        """
-        return self.forward_model(
-            torch.cat(
-                (self._get_latent_vector(obs), action.unsqueeze(1)), dim=-1))
+        layers = []
+        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],
+                    activation_fn=act))
+        return nn.Sequential(*layers)

rllib/utils/exploration/exploration.py

@@ -1,12 +1,17 @@
 from gym.spaces import Space
-from typing import Union
+from typing import List, Optional, Union, TYPE_CHECKING
 
 from ray.rllib.models.action_dist import ActionDistribution
 from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.utils.annotations import DeveloperAPI
 from ray.rllib.utils.framework import try_import_torch, TensorType
+from ray.rllib.utils.typing import LocalOptimizer, TrainerConfigDict
 
-torch, nn = try_import_torch()
+if TYPE_CHECKING:
+    from ray.rllib.policy.policy import Policy
+
+_, nn = try_import_torch()
 
 
 @DeveloperAPI
@@ -19,13 +24,13 @@ class Exploration:
     """
 
     def __init__(self, action_space: Space, *, framework: str,
-                 policy_config: dict, model: ModelV2, num_workers: int,
-                 worker_index: int):
+                 policy_config: TrainerConfigDict, model: ModelV2,
+                 num_workers: int, worker_index: int):
         """
         Args:
             action_space (Space): The action space in which to explore.
             framework (str): One of "tf" or "torch".
-            policy_config (dict): The Policy's config dict.
+            policy_config (TrainerConfigDict): The Policy's config dict.
             model (ModelV2): The Policy's model.
             num_workers (int): The overall number of workers used.
             worker_index (int): The index of the worker using this class.
@@ -45,17 +50,20 @@ class Exploration:
                 self.device = params[0].device
 
     @DeveloperAPI
-    def before_compute_actions(self,
-                               *,
-                               timestep=None,
-                               explore=None,
-                               tf_sess=None,
-                               **kwargs):
+    def before_compute_actions(
+            self,
+            *,
+            timestep: Optional[Union[TensorType, int]] = None,
+            explore: Optional[Union[TensorType, bool]] = None,
+            tf_sess: Optional["tf.Session"] = None,
+            **kwargs):
         """Hook for preparations before policy.compute_actions() is called.
 
         Args:
-            timestep (Optional[TensorType]): An optional timestep tensor.
-            explore (Optional[TensorType]): An optional explore boolean flag.
+            timestep (Optional[Union[TensorType, int]]): An optional timestep
+                tensor.
+            explore (Optional[Union[TensorType, bool]]): An optional explore
+                boolean flag.
             tf_sess (Optional[tf.Session]): The tf-session object to use.
             **kwargs: Forward compatibility kwargs.
         """
@@ -65,7 +73,7 @@ class Exploration:
     def get_exploration_action(self,
                                *,
                                action_distribution: ActionDistribution,
-                               timestep: Union[int, TensorType],
+                               timestep: Union[TensorType, int],
                                explore: bool = True):
         """Returns a (possibly) exploratory action and its log-likelihood.
 
@@ -76,11 +84,11 @@ class Exploration:
             action_distribution (ActionDistribution): The instantiated
                 ActionDistribution object to work with when creating
                 exploration actions.
-            timestep (int|TensorType): The current sampling time step. It can
-                be a tensor for TF graph mode, otherwise an integer.
-            explore (bool): True: "Normal" exploration behavior.
-                False: Suppress all exploratory behavior and return
-                    a deterministic action.
+            timestep (Union[TensorType, int]): The current sampling time step.
+                It can be a tensor for TF graph mode, otherwise an integer.
+            explore (Union[TensorType, bool]): True: "Normal" exploration
+                behavior. False: Suppress all exploratory behavior and return
+                a deterministic action.
 
         Returns:
             Tuple:
@@ -90,28 +98,6 @@ class Exploration:
         """
         pass
 
-    @DeveloperAPI
-    def get_exploration_loss(self, policy_loss, sample_batch):
-        """Modifies the policy loss with a loss associated to the exploration
-        strategy.
-
-        Args:
-            policy_loss (TODO): Loss from the Policy
-            sample_batch (SampleBatch): The SampleBatch object to post-process.
-        """
-        return policy_loss
-
-    @DeveloperAPI
-    def get_exploration_optimizer(self, config=None):
-        """
-        Returns: an optimizer for the loss from get_exploration_loss (in case
-        the exploration strategy has trainable components)
-
-        Args:
-            config: configuration for an optimizer
-        """
-        return []
-
     @DeveloperAPI
     def on_episode_start(self,
                          policy,
@@ -147,7 +133,10 @@ class Exploration:
         pass
 
     @DeveloperAPI
-    def postprocess_trajectory(self, policy, sample_batch, tf_sess=None):
+    def postprocess_trajectory(self,
+                               policy: "Policy",
+                               sample_batch,
+                               tf_sess=None):
         """Handles post-processing of done episode trajectories.
 
         Changes the given batch in place. This callback is invoked by the
@@ -160,6 +149,43 @@ class Exploration:
         """
         return sample_batch
 
+    @DeveloperAPI
+    def get_exploration_optimizer(self, optimizers: List[LocalOptimizer]):
+        """May add optimizer(s) to the Policy's own `optimizers`.
+
+        The number of optimizers (Policy's plus Exploration's optimizers) must
+        match the number of loss terms produced by the Policy's loss function
+        and the Exploration component's loss terms.
+
+        Args:
+            optimizers (List[LocalOptimizer]): The list of the Policy's
+                local optimizers.
+
+        Returns:
+            List[LocalOptimizer]: The updated list of local optimizers to use
+                on the different loss terms.
+        """
+        return optimizers
+
+    @DeveloperAPI
+    def get_exploration_loss(self, policy_loss: List[TensorType],
+                             train_batch: SampleBatch):
+        """May add loss term(s) to the Policy's own loss(es).
+
+        Args:
+            policy_loss (List[TensorType]): Loss(es) already calculated by the
+                Policy's own loss function and maybe the Model's custom loss.
+            train_batch (SampleBatch): The training data to calculate the
+                loss(es) for. This train data has already gone through
+                this Exploration's `preprocess_train_batch()` method.
+
+        Returns:
+            List[TensorType]: The updated list of loss terms.
+                This may be the original Policy loss(es), altered, and/or new
+                loss terms added to it.
+        """
+        return policy_loss
+
     @DeveloperAPI
     def get_info(self, sess=None):
         """Returns a description of the current exploration state.

rllib/utils/exploration/tests/test_curiosity.py

@@ -1,67 +1,158 @@
+import gym
+import gym_minigrid
 import numpy as np
 import ray
 import sys
 import unittest
 
-from ray.rllib.utils import check
 import ray.rllib.agents.ppo as ppo
+from ray.rllib.utils.test_utils import framework_iterator
+from ray.rllib.utils.numpy import one_hot
+from ray.tune import register_env
+
+
+class OneHotWrapper(gym.core.ObservationWrapper):
+    def __init__(self, env):
+        super().__init__(env)
+        self.observation_space = gym.spaces.Box(
+            # 11=objects; 6=colors; 3=states
+            # +4: direction
+            0.0,
+            1.0,
+            shape=(49 * (11 + 6 + 3) + 4, ),
+            dtype=np.float32)
+        self.init_x = None
+        self.init_y = None
+        self.x_positions = []
+        self.y_positions = []
+
+    def observation(self, obs):
+        # Debug output: max-x/y positions to watch exploration progress.
+        if self.step_count == 0:
+            if self.x_positions:
+                # max_diff = max(
+                #   np.sqrt((np.array(self.x_positions) - self.init_x) ** 2 + (
+                #            np.array(self.y_positions) - self.init_y) ** 2))
+                # print("After reset: max delta-x/y={}".format(max_diff))
+                self.x_positions = []
+                self.y_positions = []
+            self.init_x = self.agent_pos[0]
+            self.init_y = self.agent_pos[1]
+
+        # Are we carrying the key?
+        if self.carrying is not None:
+            print("Carrying KEY!!")
+
+        self.x_positions.append(self.agent_pos[0])
+        self.y_positions.append(self.agent_pos[1])
+
+        # One-hot the last dim into 11, 6, 3 one-hot vectors, then flatten.
+        objects = one_hot(obs[:, :, 0], depth=11)
+        colors = one_hot(obs[:, :, 1], depth=6)
+        states = one_hot(obs[:, :, 2], depth=3)
+        # Is the door we see open?
+        for x in range(7):
+            for y in range(7):
+                if objects[x, y, 4] == 1.0 and states[x, y, 0] == 1.0:
+                    print("Door OPEN!!")
+
+        all_ = np.concatenate([objects, colors, states], -1)
+        ret = np.reshape(all_, (-1, ))
+        direction = one_hot(
+            np.array(self.agent_dir), depth=4).astype(np.float32)
+        return np.concatenate([ret, direction])
+
+
+def env_maker(config):
+    name = config.get("name", "MiniGrid-Empty-5x5-v0")
+    env = gym.make(name)
+    # Only use image portion of observation (discard goal and direction).
+    env = gym_minigrid.wrappers.ImgObsWrapper(env)
+    env = OneHotWrapper(env)
+    return env
+
+
+register_env("mini-grid", env_maker)
+CONV_FILTERS = [[16, [11, 11], 3], [32, [9, 9], 3], [64, [5, 5], 3]]
 
 
 class TestCuriosity(unittest.TestCase):
-
-    # Sets up a single ray environment for every test.
-
     @classmethod
     def setUpClass(cls):
-        ray.init(local_mode=True)
+        ray.init()
 
     @classmethod
     def tearDownClass(cls):
         ray.shutdown()
 
-    def test_no_curiosity(self):
-        config = ppo.DEFAULT_CONFIG
-        env = "CartPole-v0"
-        dummy_obs = np.array([0.0, 0.1, 0.0, 0.0])
-        prev_a = np.array(0)
-        config["framework"] = "torch"
-        config["exploration_config"] = {"type": "ParameterNoise"}
-
-        trainer = ppo.PPOTrainer(config=config, env=env)
-        trainer.train()
-
-        # Make sure all actions drawn are the same, given same
-        # observations. Tests the explorations API.
-
-        actions = []
-        for _ in range(5):
-            actions.append(
-                trainer.compute_action(
-                    observation=dummy_obs,
-                    explore=False,
-                    prev_action=prev_a,
-                    prev_reward=1.0 if prev_a is not None else None))
-            check(actions[-1], actions[0])
-        print(actions)
-
-    def test_curiosity(self):
-        config = ppo.DEFAULT_CONFIG
-
-        env = "CartPole-v0"
-        config["framework"] = "torch"
-        config["exploration_config"] = {
-            "type": "ray.rllib.utils.exploration.curiosity.Curiosity",
-            "forward_net_hiddens": [64],
-            "inverse_net_hiddens": [32, 4],
-            "feature_net_hiddens": [16, 8],
-            "feature_dim": 8,
-            "forward_activation": "relu",
-            "inverse_activation": "relu",
-            "feature_activation": "relu",
-            "submodule": "EpsilonGreedy",
+    def test_curiosity_on_large_frozen_lake(self):
+        config = ppo.DEFAULT_CONFIG.copy()
+        # 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": [
+                "SFFFFFFFFFFFFFFF",
+                "FFFFFFFFFFFFFFFF",
+                "FFFFFFFFFFFFFFFF",
+                "FFFFFFFFFFFFFFFF",
+                "FFFFFFFFFFFFFFFF",
+                "FFFFFFFFFFFFFFFF",
+                "FFFFFFFFFFFFFFFF",
+                "FFFFFFFFFFFFFFFF",
+                "FFFFFFFFFFFFFFFF",
+                "FFFFFFFFFFFFFFFF",
+                "FFFFFFFFFFFFFFFF",
+                "FFFFFFFFFFFFFFFF",
+                "FFFFFFFFFFFFFFFF",
+                "FFFFFFFFFFFFFFFF",
+                "FFFFFFFFFFFFFFFF",
+                "FFFFFFFFFFFFFFFG",
+            ],
+            "is_slippery": False
         }
-        trainer = ppo.PPOTrainer(config=config, env=env)
-        trainer.train()
+        # Limit horizon to make it really hard for non-curious agent to reach
+        # the goal state.
+        config["horizon"] = 40
+        config["num_workers"] = 0  # local only
+        config["train_batch_size"] = 512
+        config["num_sgd_iter"] = 10
+
+        num_iterations = 30
+        for _ in framework_iterator(config, frameworks="torch"):
+            # W/ Curiosity.
+            config["exploration_config"] = {
+                "type": "Curiosity",
+                "feature_dim": 128,
+                "eta": 0.05,
+                "sub_exploration": {
+                    "type": "StochasticSampling",
+                }
+            }
+            trainer = ppo.PPOTrainer(config=config)
+            rewards_w = 0.0
+            for _ in range(num_iterations):
+                result = trainer.train()
+                rewards_w += result["episode_reward_mean"]
+                print(result)
+            rewards_w /= num_iterations
+            trainer.stop()
+
+            # W/o Curiosity.
+            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)
+            rewards_wo /= num_iterations
+            trainer.stop()
+
+            self.assertTrue(rewards_wo == 0.0)
+            self.assertGreater(rewards_w, 0.1)
 
 
 if __name__ == "__main__":

rllib/utils/spaces/space_utils.py

@@ -67,7 +67,7 @@ def flatten_to_single_ndarray(input_):
     """Returns a single np.ndarray given a list/tuple of np.ndarrays.
 
     Args:
-        input_ (Union[List[np.ndarray],np.ndarray]): The list of ndarrays or
+        input_ (Union[List[np.ndarray], np.ndarray]): The list of ndarrays or
             a single ndarray.
 
     Returns:

rllib/utils/typing.py

@@ -19,6 +19,10 @@ EnvConfigDict = dict
 # the model catalog.
 ModelConfigDict = dict
 
+# Objects that can be created through the `from_config()` util method
+# need a config dict with a "type" key, a class path (str), or a type directly.
+FromConfigSpec = Union[Dict[str, Any], type, str]
+
 # Represents a BaseEnv, MultiAgentEnv, ExternalEnv, ExternalMultiAgentEnv,
 # VectorEnv, or gym.Env.
 EnvType = Any
@@ -61,6 +65,10 @@ FileType = Any
 # Represents the result dict returned by Trainer.train().
 ResultDict = dict
 
+# A tf or torch local optimizer object.
+LocalOptimizer = Union["tf.keras.optimizers.Optimizer",
+                       "torch.optim.Optimizer"]
+
 # Dict of tensors returned by compute gradients on the policy, e.g.,
 # {"td_error": [...], "learner_stats": {"vf_loss": ..., ...}}, for multi-agent,
 # {"policy1": {"learner_stats": ..., }, "policy2": ...}.