From 592f3132106820f4613ab4fbd6ce5c221eaa78b6 Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Thu, 8 Aug 2019 14:03:28 -0700 Subject: [PATCH] [rllib] Centralized critic / PPO example on TwoStepGame (#5392) --- ci/jenkins_tests/run_rllib_tests.sh | 6 + doc/source/rllib-algorithms.rst | 2 +- doc/source/rllib-env.rst | 12 +- doc/source/rllib-examples.rst | 6 +- rllib/agents/a3c/a3c_tf_policy.py | 2 +- rllib/agents/impala/vtrace_policy.py | 2 +- rllib/evaluation/metrics.py | 11 +- rllib/examples/centralized_critic.py | 222 +++++++++++++++++++++++++ rllib/examples/centralized_critic_2.py | 158 ++++++++++++++++++ rllib/examples/twostep_game.py | 11 +- rllib/models/tf/fcnet_v2.py | 2 +- rllib/policy/dynamic_tf_policy.py | 7 +- rllib/policy/policy.py | 2 +- rllib/policy/tf_policy.py | 11 +- rllib/policy/tf_policy_template.py | 2 +- 15 files changed, 432 insertions(+), 24 deletions(-) create mode 100644 rllib/examples/centralized_critic.py create mode 100644 rllib/examples/centralized_critic_2.py diff --git a/ci/jenkins_tests/run_rllib_tests.sh b/ci/jenkins_tests/run_rllib_tests.sh index 9ad5db1f2..1fbb46e82 100644 --- a/ci/jenkins_tests/run_rllib_tests.sh +++ b/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 diff --git a/doc/source/rllib-algorithms.rst b/doc/source/rllib-algorithms.rst index d00e59b5b..244929b9c 100644 --- a/doc/source/rllib-algorithms.rst +++ b/doc/source/rllib-algorithms.rst @@ -207,7 +207,7 @@ Policy Gradients ---------------- `[paper] `__ `[implementation] `__ 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) diff --git a/doc/source/rllib-env.rst b/doc/source/rllib-env.rst index 618e7a221..1aa9aba54 100644 --- a/doc/source/rllib-env.rst +++ b/doc/source/rllib-env.rst @@ -283,9 +283,11 @@ There is a full example of this in the `example training script `__. +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 `__. + +**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 `__. Grouping Agents ~~~~~~~~~~~~~~~ diff --git a/doc/source/rllib-examples.rst b/doc/source/rllib-examples.rst index 8c6c61cf7..036b6f8ee 100644 --- a/doc/source/rllib-examples.rst +++ b/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 `__: Example of the two-step game from the `QMIX paper `__. +- `PPO with centralized critic on two-step game `__: + Example of customizing PPO to leverage a centralized value function. +- `Centralized critic in the env `__: + A simpler method of implementing a centralized critic by augmentating agent observations with global information. - `Hand-coded policy `__: Example of running a custom hand-coded policy alongside trainable policies. - `Weight sharing between policies `__: @@ -67,8 +71,6 @@ Multi-Agent and Hierarchical Example of alternating training between two DQN and PPO trainers. - `Hierarchical training `__: Example of hierarchical training using the multi-agent API. -- `PPO with centralized value function `__: - Example of customizing PPO to include a centralized value function, including a runnable script that demonstrates cooperative CartPole. Community Examples ------------------ diff --git a/rllib/agents/a3c/a3c_tf_policy.py b/rllib/agents/a3c/a3c_tf_policy.py index a614dc2d3..64974b0c7 100644 --- a/rllib/agents/a3c/a3c_tf_policy.py +++ b/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( diff --git a/rllib/agents/impala/vtrace_policy.py b/rllib/agents/impala/vtrace_policy.py index a288fb5d1..45228ea71 100644 --- a/rllib/agents/impala/vtrace_policy.py +++ b/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), } diff --git a/rllib/evaluation/metrics.py b/rllib/evaluation/metrics.py index 75318e929..6f9eda1f2 100644 --- a/rllib/evaluation/metrics.py +++ b/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)) diff --git a/rllib/examples/centralized_critic.py b/rllib/examples/centralized_critic.py new file mode 100644 index 000000000..7e9495204 --- /dev/null +++ b/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", + }, + }) diff --git a/rllib/examples/centralized_critic_2.py b/rllib/examples/centralized_critic_2.py new file mode 100644 index 000000000..419b5fd74 --- /dev/null +++ b/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", + }, + }) diff --git a/rllib/examples/twostep_game.py b/rllib/examples/twostep_game.py index b0f530919..f47f64730 100644 --- a/rllib/examples/twostep_game.py +++ b/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 diff --git a/rllib/models/tf/fcnet_v2.py b/rllib/models/tf/fcnet_v2.py index 2a8d576d1..2231b45a9 100644 --- a/rllib/models/tf/fcnet_v2.py +++ b/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): diff --git a/rllib/policy/dynamic_tf_policy.py b/rllib/policy/dynamic_tf_policy.py index 75c3a7d91..be80ec6b3 100644 --- a/rllib/policy/dynamic_tf_policy.py +++ b/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): diff --git a/rllib/policy/policy.py b/rllib/policy/policy.py index 4fecb15f5..dcf19db14 100644 --- a/rllib/policy/policy.py +++ b/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. diff --git a/rllib/policy/tf_policy.py b/rllib/policy/tf_policy.py index a85b56347..560504488 100644 --- a/rllib/policy/tf_policy.py +++ b/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): diff --git a/rllib/policy/tf_policy_template.py b/rllib/policy/tf_policy_template.py index e44a6dac5..94db613fb 100644 --- a/rllib/policy/tf_policy_template.py +++ b/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