[rllib] Centralized critic / PPO example on TwoStepGame (#5392)

ci/jenkins_tests/run_rllib_tests.sh

@@ -428,6 +428,12 @@ docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \
 docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \
     /ray/ci/suppress_output python /ray/rllib/contrib/random_agent/random_agent.py
 
+docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \
+    /ray/ci/suppress_output python /ray/rllib/examples/centralized_critic.py --stop=2000
+
+docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \
+    /ray/ci/suppress_output python /ray/rllib/examples/centralized_critic_2.py --stop=2000
+
 docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \
     /ray/ci/suppress_output python /ray/rllib/examples/twostep_game.py --stop=2000 --run=contrib/MADDPG
 

doc/source/rllib-algorithms.rst

@@ -207,7 +207,7 @@ Policy Gradients
 ----------------
 `[paper] <https://papers.nips.cc/paper/1713-policy-gradient-methods-for-reinforcement-learning-with-function-approximation.pdf>`__ `[implementation] <https://github.com/ray-project/ray/blob/master/rllib/agents/pg/pg.py>`__ We include a vanilla policy gradients implementation as an example algorithm in both TensorFlow and PyTorch. This is usually outperformed by PPO.
 
-.. figure:: ppo-arch.svg
+.. figure:: a2c-arch.svg
 
     Policy gradients architecture (same as A2C)
 

doc/source/rllib-env.rst

@@ -283,9 +283,11 @@ There is a full example of this in the `example training script <https://github.
 Implementing a Centralized Critic
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-Implementing a centralized critic that takes as input the observations and actions of other concurrent agents requires the definition of custom policies. It can be done as follows:
+Here are two ways to implement a centralized critic compatible with the multi-agent API:
 
-1. Querying the critic: this can be done in the ``postprocess_trajectory`` method of a custom policy, which has full access to the policies and observations of concurrent agents via the ``other_agent_batches`` and ``episode`` arguments. The batch of critic predictions can then be added to the postprocessed trajectory. Here's an example:
+**Strategy 1: Sharing experiences in the trajectory preprocessor**:
+
+The most general way of implementing a centralized critic involves modifying the ``postprocess_trajectory`` method of a custom policy, which has full access to the policies and observations of concurrent agents via the ``other_agent_batches`` and ``episode`` arguments. The batch of critic predictions can then be added to the postprocessed trajectory. Here's an example:
 
 .. code-block:: python
 
@@ -300,7 +302,11 @@ Implementing a centralized critic that takes as input the observations and actio
             self.critic_network, feed_dict={"obs": global_obs_batch})
         return sample_batch
 
-2. Updating the critic: the centralized critic loss can be added to the loss of the custom policy, the same as with any other value function. For an example of defining loss inputs, see the `PGPolicy example <https://github.com/ray-project/ray/blob/master/rllib/agents/pg/pg_policy.py>`__.
+To update the critic, you'll also have to modify the loss of the policy. For an end-to-end runnable example, see `examples/centralized_critic.py <https://github.com/ray-project/ray/blob/master/rllib/examples/centralized_critic.py>`__.
+
+**Strategy 2: Sharing observations through the environment**:
+
+Alternatively, the env itself can be modified to share observations between agents. In this strategy, each observation includes all global state, and policies use a custom model to ignore state they aren't supposed to "see" when computing actions. The advantage of this approach is that it's very simple and you don't have to change the algorithm at all -- just use an env wrapper and custom model. However, it is a bit less principled in that you have to change the agent observation spaces and the environment. You can find a runnable example of this strategy at `examples/centralized_critic_2.py <https://github.com/ray-project/ray/blob/master/rllib/examples/centralized_critic_2.py>`__.
 
 Grouping Agents
 ~~~~~~~~~~~~~~~

doc/source/rllib-examples.rst

@@ -59,6 +59,10 @@ Multi-Agent and Hierarchical
    Example of different heuristic and learned policies competing against each other in rock-paper-scissors.
 - `Two-step game <https://github.com/ray-project/ray/blob/master/rllib/examples/twostep_game.py>`__:
    Example of the two-step game from the `QMIX paper <https://arxiv.org/pdf/1803.11485.pdf>`__.
+- `PPO with centralized critic on two-step game <https://github.com/ray-project/ray/blob/master/rllib/examples/centralized_critic.py>`__:
+   Example of customizing PPO to leverage a centralized value function.
+- `Centralized critic in the env <https://github.com/ray-project/ray/blob/master/rllib/examples/centralized_critic_2.py>`__:
+   A simpler method of implementing a centralized critic by augmentating agent observations with global information.
 - `Hand-coded policy <https://github.com/ray-project/ray/blob/master/rllib/examples/multiagent_custom_policy.py>`__:
    Example of running a custom hand-coded policy alongside trainable policies.
 - `Weight sharing between policies <https://github.com/ray-project/ray/blob/master/rllib/examples/multiagent_cartpole.py>`__:
@@ -67,8 +71,6 @@ Multi-Agent and Hierarchical
    Example of alternating training between two DQN and PPO trainers.
 - `Hierarchical training <https://github.com/ray-project/ray/blob/master/rllib/examples/hierarchical_training.py>`__:
    Example of hierarchical training using the multi-agent API.
-- `PPO with centralized value function <https://github.com/ray-project/ray/pull/3642/files>`__:
-   Example of customizing PPO to include a centralized value function, including a runnable script that demonstrates cooperative CartPole.
 
 Community Examples
 ------------------

rllib/agents/a3c/a3c_tf_policy.py

@@ -98,7 +98,7 @@ def stats(policy, batch_tensors):
     }
 
 
-def grad_stats(policy, grads):
+def grad_stats(policy, batch_tensors, grads):
     return {
         "grad_gnorm": tf.global_norm(grads),
         "vf_explained_var": explained_variance(

rllib/agents/impala/vtrace_policy.py

@@ -225,7 +225,7 @@ def stats(policy, batch_tensors):
     }
 
 
-def grad_stats(policy, grads):
+def grad_stats(policy, batch_tensors, grads):
     return {
         "grad_gnorm": tf.global_norm(grads),
     }

rllib/evaluation/metrics.py

@@ -120,8 +120,13 @@ def summarize_episodes(episodes, new_episodes):
     avg_reward = np.mean(episode_rewards)
     avg_length = np.mean(episode_lengths)
 
+    policy_reward_min = {}
+    policy_reward_mean = {}
+    policy_reward_max = {}
     for policy_id, rewards in policy_rewards.copy().items():
-        policy_rewards[policy_id] = np.mean(rewards)
+        policy_reward_min[policy_id] = np.min(rewards)
+        policy_reward_mean[policy_id] = np.mean(rewards)
+        policy_reward_max[policy_id] = np.max(rewards)
 
     for k, v_list in custom_metrics.copy().items():
         custom_metrics[k + "_mean"] = np.mean(v_list)
@@ -153,7 +158,9 @@ def summarize_episodes(episodes, new_episodes):
         episode_reward_mean=avg_reward,
         episode_len_mean=avg_length,
         episodes_this_iter=len(new_episodes),
-        policy_reward_mean=dict(policy_rewards),
+        policy_reward_min=policy_reward_min,
+        policy_reward_max=policy_reward_max,
+        policy_reward_mean=policy_reward_mean,
         custom_metrics=dict(custom_metrics),
         sampler_perf=dict(perf_stats),
         off_policy_estimator=dict(estimators))

rllib/examples/centralized_critic.py

@@ -0,0 +1,222 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+"""An example of customizing PPO to leverage a centralized critic.
+
+Here the model and policy are hard-coded to implement a centralized critic
+for TwoStepGame, but you can adapt this for your own use cases.
+
+Compared to simply running `twostep_game.py --run=PPO`, this centralized
+critic version reaches vf_explained_variance=1.0 more stably since it takes
+into account the opponent actions as well as the policy's. Note that this is
+also using two independent policies instead of weight-sharing with one.
+
+See also: centralized_critic_2.py for a simpler approach that instead
+modifies the environment.
+"""
+
+import argparse
+import numpy as np
+
+from ray import tune
+from ray.rllib.agents.ppo.ppo import PPOTrainer
+from ray.rllib.agents.ppo.ppo_policy import PPOTFPolicy, KLCoeffMixin, \
+    PPOLoss, BEHAVIOUR_LOGITS
+from ray.rllib.evaluation.postprocessing import compute_advantages, \
+    Postprocessing
+from ray.rllib.examples.twostep_game import TwoStepGame
+from ray.rllib.models import ModelCatalog
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.tf_policy import LearningRateSchedule, \
+    EntropyCoeffSchedule
+from ray.rllib.models.tf.tf_modelv2 import TFModelV2
+from ray.rllib.models.tf.fcnet_v2 import FullyConnectedNetwork
+from ray.rllib.utils.explained_variance import explained_variance
+from ray.rllib.utils import try_import_tf
+
+tf = try_import_tf()
+
+OPPONENT_OBS = "opponent_obs"
+OPPONENT_ACTION = "opponent_action"
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--stop", type=int, default=100000)
+
+
+class CentralizedCriticModel(TFModelV2):
+    """Multi-agent model that implements a centralized VF."""
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        super(CentralizedCriticModel, self).__init__(
+            obs_space, action_space, num_outputs, model_config, name)
+        # Base of the model
+        self.model = FullyConnectedNetwork(obs_space, action_space,
+                                           num_outputs, model_config, name)
+        self.register_variables(self.model.variables())
+
+        # Central VF maps (obs, opp_ops, opp_act) -> vf_pred
+        obs = tf.keras.layers.Input(shape=(6, ), name="obs")
+        opp_obs = tf.keras.layers.Input(shape=(6, ), name="opp_obs")
+        opp_act = tf.keras.layers.Input(shape=(2, ), name="opp_act")
+        concat_obs = tf.keras.layers.Concatenate(axis=1)(
+            [obs, opp_obs, opp_act])
+        central_vf_dense = tf.keras.layers.Dense(
+            16, activation=tf.nn.tanh, name="c_vf_dense")(concat_obs)
+        central_vf_out = tf.keras.layers.Dense(
+            1, activation=None, name="c_vf_out")(central_vf_dense)
+        self.central_vf = tf.keras.Model(
+            inputs=[obs, opp_obs, opp_act], outputs=central_vf_out)
+        self.register_variables(self.central_vf.variables)
+
+    def forward(self, input_dict, state, seq_lens):
+        return self.model.forward(input_dict, state, seq_lens)
+
+    def central_value_function(self, obs, opponent_obs, opponent_actions):
+        return tf.reshape(
+            self.central_vf(
+                [obs, opponent_obs,
+                 tf.one_hot(opponent_actions, 2)]), [-1])
+
+
+class CentralizedValueMixin(object):
+    """Add methods to evaluate the central value function from the model."""
+
+    def __init__(self):
+        self.central_value_function = self.model.central_value_function(
+            self.get_placeholder(SampleBatch.CUR_OBS),
+            self.get_placeholder(OPPONENT_OBS),
+            self.get_placeholder(OPPONENT_ACTION))
+
+    def compute_central_vf(self, obs, opponent_obs, opponent_actions):
+        feed_dict = {
+            self.get_placeholder(SampleBatch.CUR_OBS): obs,
+            self.get_placeholder(OPPONENT_OBS): opponent_obs,
+            self.get_placeholder(OPPONENT_ACTION): opponent_actions,
+        }
+        return self.get_session().run(self.central_value_function, feed_dict)
+
+
+# Grabs the opponent obs/act and includes it in the experience batch,
+# and computes GAE using the central vf predictions.
+def centralized_critic_postprocessing(policy,
+                                      sample_batch,
+                                      other_agent_batches=None,
+                                      episode=None):
+    if policy.loss_initialized():
+        assert sample_batch["dones"][-1], \
+            "Not implemented for batch_mode=truncate_episodes"
+        assert other_agent_batches is not None
+        [(_, opponent_batch)] = list(other_agent_batches.values())
+
+        # also record the opponent obs and actions in the trajectory
+        sample_batch[OPPONENT_OBS] = opponent_batch[SampleBatch.CUR_OBS]
+        sample_batch[OPPONENT_ACTION] = opponent_batch[SampleBatch.ACTIONS]
+
+        # overwrite default VF prediction with the central VF
+        sample_batch[SampleBatch.VF_PREDS] = policy.compute_central_vf(
+            sample_batch[SampleBatch.CUR_OBS], sample_batch[OPPONENT_OBS],
+            sample_batch[OPPONENT_ACTION])
+    else:
+        # policy hasn't initialized yet, use zeros
+        sample_batch[OPPONENT_OBS] = np.zeros_like(
+            sample_batch[SampleBatch.CUR_OBS])
+        sample_batch[OPPONENT_ACTION] = np.zeros_like(
+            sample_batch[SampleBatch.ACTIONS])
+        sample_batch[SampleBatch.VF_PREDS] = np.zeros_like(
+            sample_batch[SampleBatch.ACTIONS], dtype=np.float32)
+
+    batch = compute_advantages(
+        sample_batch,
+        0.0,
+        policy.config["gamma"],
+        policy.config["lambda"],
+        use_gae=policy.config["use_gae"])
+    return batch
+
+
+# Copied from PPO but optimizing the central value function
+def loss_with_central_critic(policy, batch_tensors):
+    CentralizedValueMixin.__init__(policy)
+
+    policy.loss_obj = PPOLoss(
+        policy.action_space,
+        batch_tensors[Postprocessing.VALUE_TARGETS],
+        batch_tensors[Postprocessing.ADVANTAGES],
+        batch_tensors[SampleBatch.ACTIONS],
+        batch_tensors[BEHAVIOUR_LOGITS],
+        batch_tensors[SampleBatch.VF_PREDS],
+        policy.action_dist,
+        policy.central_value_function,
+        policy.kl_coeff,
+        tf.ones_like(batch_tensors[Postprocessing.ADVANTAGES], dtype=tf.bool),
+        entropy_coeff=policy.entropy_coeff,
+        clip_param=policy.config["clip_param"],
+        vf_clip_param=policy.config["vf_clip_param"],
+        vf_loss_coeff=policy.config["vf_loss_coeff"],
+        use_gae=policy.config["use_gae"],
+        model_config=policy.config["model"])
+
+    return policy.loss_obj.loss
+
+
+def setup_mixins(policy, obs_space, action_space, config):
+    # copied from PPO
+    KLCoeffMixin.__init__(policy, config)
+    EntropyCoeffSchedule.__init__(policy, config["entropy_coeff"],
+                                  config["entropy_coeff_schedule"])
+    LearningRateSchedule.__init__(policy, config["lr"], config["lr_schedule"])
+    # hack: put in a noop VF so some of the inherited PPO code runs
+    policy.value_function = tf.zeros(
+        tf.shape(policy.get_placeholder(SampleBatch.CUR_OBS))[0])
+
+
+def central_vf_stats(policy, batch_tensors, grads):
+    # Report the explained variance of the central value function.
+    return {
+        "vf_explained_var": explained_variance(
+            batch_tensors[Postprocessing.VALUE_TARGETS],
+            policy.central_value_function),
+    }
+
+
+CCPPO = PPOTFPolicy.with_updates(
+    name="CCPPO",
+    postprocess_fn=centralized_critic_postprocessing,
+    loss_fn=loss_with_central_critic,
+    before_loss_init=setup_mixins,
+    grad_stats_fn=central_vf_stats,
+    mixins=[
+        LearningRateSchedule, EntropyCoeffSchedule, KLCoeffMixin,
+        CentralizedValueMixin
+    ])
+
+CCTrainer = PPOTrainer.with_updates(name="CCPPOTrainer", default_policy=CCPPO)
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    ModelCatalog.register_custom_model("cc_model", CentralizedCriticModel)
+    tune.run(
+        CCTrainer,
+        stop={
+            "timesteps_total": args.stop,
+            "episode_reward_mean": 7.99,
+        },
+        config={
+            "env": TwoStepGame,
+            "batch_mode": "complete_episodes",
+            "num_workers": 0,
+            "multiagent": {
+                "policies": {
+                    "pol1": (None, TwoStepGame.observation_space,
+                             TwoStepGame.action_space, {}),
+                    "pol2": (None, TwoStepGame.observation_space,
+                             TwoStepGame.action_space, {}),
+                },
+                "policy_mapping_fn": tune.function(
+                    lambda x: "pol1" if x == 0 else "pol2"),
+            },
+            "model": {
+                "custom_model": "cc_model",
+            },
+        })

rllib/examples/centralized_critic_2.py

@@ -0,0 +1,158 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+"""An example of implementing a centralized critic by modifying the env.
+
+The advantage of this approach is that it's very simple and you don't have to
+change the algorithm at all -- just use an env wrapper and custom model.
+However, it is a bit less principled in that you have to change the agent
+observation spaces and the environment.
+
+See also: centralized_critic.py for an alternative approach that instead
+modifies the policy to add a centralized value function.
+"""
+
+import numpy as np
+from gym.spaces import Box, Dict, Discrete
+import argparse
+
+from ray import tune
+from ray.rllib.env.multi_agent_env import MultiAgentEnv
+from ray.rllib.examples.twostep_game import TwoStepGame
+from ray.rllib.models import ModelCatalog
+from ray.rllib.models.tf.tf_modelv2 import TFModelV2
+from ray.rllib.models.tf.fcnet_v2 import FullyConnectedNetwork
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.utils import try_import_tf
+
+tf = try_import_tf()
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--stop", type=int, default=100000)
+
+
+class CentralizedCriticModel(TFModelV2):
+    """Multi-agent model that implements a centralized VF.
+
+    It assumes the observation is a dict with 'own_obs' and 'opponent_obs', the
+    former of which can be used for computing actions (i.e., decentralized
+    execution), and the latter for optimization (i.e., centralized learning).
+
+    This model has two parts:
+    - An action model that looks at just 'own_obs' to compute actions
+    - A value model that also looks at the 'opponent_obs' / 'opponent_action'
+      to compute the value (it does this by using the 'obs_flat' tensor).
+    """
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        super(CentralizedCriticModel, self).__init__(
+            obs_space, action_space, num_outputs, model_config, name)
+
+        self.action_model = FullyConnectedNetwork(
+            Box(low=0, high=1, shape=(6, )),  # one-hot encoded Discrete(6)
+            action_space,
+            num_outputs,
+            model_config,
+            name + "_action")
+        self.register_variables(self.action_model.variables())
+
+        self.value_model = FullyConnectedNetwork(obs_space, action_space, 1,
+                                                 model_config, name + "_vf")
+        self.register_variables(self.value_model.variables())
+
+    def forward(self, input_dict, state, seq_lens):
+        self._value_out, _ = self.value_model({
+            "obs": input_dict["obs_flat"]
+        }, state, seq_lens)
+        return self.action_model({
+            "obs": input_dict["obs"]["own_obs"]
+        }, state, seq_lens)
+
+    def value_function(self):
+        return tf.reshape(self._value_out, [-1])
+
+
+class GlobalObsTwoStepGame(MultiAgentEnv):
+    action_space = Discrete(2)
+    observation_space = Dict({
+        "own_obs": Discrete(6),
+        "opponent_obs": Discrete(6),
+        "opponent_action": Discrete(2),
+    })
+
+    def __init__(self, env_config):
+        self.env = TwoStepGame(env_config)
+
+    def reset(self):
+        obs_dict = self.env.reset()
+        return self.to_global_obs(obs_dict)
+
+    def step(self, action_dict):
+        obs_dict, rewards, dones, infos = self.env.step(action_dict)
+        return self.to_global_obs(obs_dict), rewards, dones, infos
+
+    def to_global_obs(self, obs_dict):
+        return {
+            self.env.agent_1: {
+                "own_obs": obs_dict[self.env.agent_1],
+                "opponent_obs": obs_dict[self.env.agent_2],
+                "opponent_action": 0,  # populated by fill_in_actions
+            },
+            self.env.agent_2: {
+                "own_obs": obs_dict[self.env.agent_2],
+                "opponent_obs": obs_dict[self.env.agent_1],
+                "opponent_action": 0,  # populated by fill_in_actions
+            },
+        }
+
+
+def fill_in_actions(info):
+    """Callback that saves opponent actions into the agent obs.
+
+    If you don't care about opponent actions you can leave this out."""
+
+    to_update = info["post_batch"][SampleBatch.CUR_OBS]
+    my_id = info["agent_id"]
+    other_id = 1 if my_id == 0 else 0
+    action_encoder = ModelCatalog.get_preprocessor_for_space(Discrete(2))
+
+    # set the opponent actions into the observation
+    _, opponent_batch = info["all_pre_batches"][other_id]
+    opponent_actions = np.array([
+        action_encoder.transform(a)
+        for a in opponent_batch[SampleBatch.ACTIONS]
+    ])
+    to_update[:, -2:] = opponent_actions
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    ModelCatalog.register_custom_model("cc_model", CentralizedCriticModel)
+    tune.run(
+        "PPO",
+        stop={
+            "timesteps_total": args.stop,
+            "episode_reward_mean": 7.99,
+        },
+        config={
+            "env": GlobalObsTwoStepGame,
+            "batch_mode": "complete_episodes",
+            "callbacks": {
+                "on_postprocess_traj": tune.function(fill_in_actions),
+            },
+            "num_workers": 0,
+            "multiagent": {
+                "policies": {
+                    "pol1": (None, GlobalObsTwoStepGame.observation_space,
+                             GlobalObsTwoStepGame.action_space, {}),
+                    "pol2": (None, GlobalObsTwoStepGame.observation_space,
+                             GlobalObsTwoStepGame.action_space, {}),
+                },
+                "policy_mapping_fn": tune.function(
+                    lambda x: "pol1" if x == 0 else "pol2"),
+            },
+            "model": {
+                "custom_model": "cc_model",
+            },
+        })

rllib/examples/twostep_game.py

@@ -1,4 +1,13 @@
-"""The two-step game from QMIX: https://arxiv.org/pdf/1803.11485.pdf"""
+"""The two-step game from QMIX: https://arxiv.org/pdf/1803.11485.pdf
+
+Configurations you can try:
+    - normal policy gradients (PG)
+    - contrib/MADDPG
+    - QMIX
+    - APEX_QMIX
+
+See also: centralized_critic.py for centralized critic PPO on this game.
+"""
 
 from __future__ import absolute_import
 from __future__ import division

rllib/models/tf/fcnet_v2.py

@@ -80,7 +80,7 @@ class FullyConnectedNetwork(TFModelV2):
         self.register_variables(self.base_model.variables)
 
     def forward(self, input_dict, state, seq_lens):
-        model_out, self._value_out = self.base_model(input_dict["obs"])
+        model_out, self._value_out = self.base_model(input_dict["obs_flat"])
         return model_out, state
 
     def value_function(self):

rllib/policy/dynamic_tf_policy.py

@@ -242,6 +242,7 @@ class DynamicTFPolicy(TFPolicy):
             existing_inputs=input_dict,
             existing_model=self.model)
 
+        instance._loss_input_dict = input_dict
         loss = instance._do_loss_init(input_dict)
         loss_inputs = [(k, existing_inputs[i])
                        for i, (k, _) in enumerate(self._loss_inputs)]
@@ -249,7 +250,7 @@ class DynamicTFPolicy(TFPolicy):
         TFPolicy._initialize_loss(instance, loss, loss_inputs)
         if instance._grad_stats_fn:
             instance._stats_fetches.update(
-                instance._grad_stats_fn(instance, instance._grads))
+                instance._grad_stats_fn(instance, input_dict, instance._grads))
         return instance
 
     @override(Policy)
@@ -326,13 +327,15 @@ class DynamicTFPolicy(TFPolicy):
                 "Initializing loss function with dummy input:\n\n{}\n".format(
                     summarize(batch_tensors)))
 
+        self._loss_input_dict = batch_tensors
         loss = self._do_loss_init(batch_tensors)
         for k in sorted(batch_tensors.accessed_keys):
             loss_inputs.append((k, batch_tensors[k]))
 
         TFPolicy._initialize_loss(self, loss, loss_inputs)
         if self._grad_stats_fn:
-            self._stats_fetches.update(self._grad_stats_fn(self, self._grads))
+            self._stats_fetches.update(
+                self._grad_stats_fn(self, batch_tensors, self._grads))
         self._sess.run(tf.global_variables_initializer())
 
     def _do_loss_init(self, batch_tensors):

rllib/policy/policy.py

@@ -152,7 +152,7 @@ class Policy(object):
                 which will contain at most one episode trajectory.
             other_agent_batches (dict): In a multi-agent env, this contains a
                 mapping of agent ids to (policy, agent_batch) tuples
-                containing the policy and experiences of the other agent.
+                containing the policy and experiences of the other agents.
             episode (MultiAgentEpisode): this provides access to all of the
                 internal episode state, which may be useful for model-based or
                 multi-agent algorithms.

rllib/policy/tf_policy.py

@@ -121,6 +121,7 @@ class TFPolicy(Policy):
         self._batch_divisibility_req = batch_divisibility_req
         self._update_ops = update_ops
         self._stats_fetches = {}
+        self._loss_input_dict = None
 
         if loss is not None:
             self._initialize_loss(loss, loss_inputs)
@@ -155,14 +156,8 @@ class TFPolicy(Policy):
         if name in obs_inputs:
             return obs_inputs[name]
 
-        if not self.loss_initialized():
-            raise RuntimeError(
-                "You cannot call policy.get_placeholder() for non-obs inputs "
-                "before the loss has been initialized. To avoid this, use "
-                "policy.loss_initialized() to check whether this is the "
-                "case, or move the call to later (e.g., from stats_fn to "
-                "grad_stats_fn).")
-
+        assert self._loss_input_dict is not None, \
+            "Should have set this before get_placeholder can be called"
         return self._loss_input_dict[name]
 
     def get_session(self):

rllib/policy/tf_policy_template.py

@@ -67,7 +67,7 @@ def build_tf_policy(name,
         apply_gradients_fn (func): optional function that returns an apply
             gradients op given (policy, optimizer, grads_and_vars)
         grad_stats_fn (func): optional function that returns a dict of
-            TF fetches given the policy and loss gradient tensors
+            TF fetches given the policy, batch input, and gradient tensors
         extra_action_fetches_fn (func): optional function that returns
             a dict of TF fetches given the policy object
         extra_action_feed_fn (func): optional function that returns a feed dict