[rllib] observation function api for multi-agent (#8236)

doc/source/rllib-env.rst

@@ -253,9 +253,9 @@ The most general way of implementing a centralized critic involves modifying the
 
 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**:
+**Strategy 2: Sharing observations through an observation function**:
 
-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>`__.
+Alternatively, you can use an observation function 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 the observation func (i.e., like an env wrapper) and custom model. However, it is a bit less principled in that you have to change the agent observation spaces to include training-time only information. 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
 ~~~~~~~~~~~~~~~

rllib/agents/callbacks.py

@@ -1,7 +1,7 @@
 from typing import Dict
 
 from ray.rllib.env import BaseEnv
-from ray.rllib.policy import Policy
+from ray.rllib.policy import Policy, PolicyID, AgentID
 from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.evaluation import MultiAgentEpisode, RolloutWorker
 from ray.rllib.utils.annotations import PublicAPI
@@ -27,7 +27,7 @@ class DefaultCallbacks:
         self.legacy_callbacks = legacy_callbacks_dict or {}
 
     def on_episode_start(self, worker: RolloutWorker, base_env: BaseEnv,
-                         policies: Dict[str, Policy],
+                         policies: Dict[PolicyID, Policy],
                          episode: MultiAgentEpisode, **kwargs):
         """Callback run on the rollout worker before each episode starts.
 
@@ -73,8 +73,8 @@ class DefaultCallbacks:
             })
 
     def on_episode_end(self, worker: RolloutWorker, base_env: BaseEnv,
-                       policies: Dict[str, Policy], episode: MultiAgentEpisode,
-                       **kwargs):
+                       policies: Dict[PolicyID, Policy],
+                       episode: MultiAgentEpisode, **kwargs):
         """Runs when an episode is done.
 
         Args:
@@ -99,9 +99,9 @@ class DefaultCallbacks:
 
     def on_postprocess_trajectory(
             self, worker: RolloutWorker, episode: MultiAgentEpisode,
-            agent_id: str, policy_id: str, policies: Dict[str, Policy],
-            postprocessed_batch: SampleBatch,
-            original_batches: Dict[str, SampleBatch], **kwargs):
+            agent_id: AgentID, policy_id: PolicyID,
+            policies: Dict[PolicyID, Policy], postprocessed_batch: SampleBatch,
+            original_batches: Dict[AgentID, SampleBatch], **kwargs):
         """Called immediately after a policy's postprocess_fn is called.
 
         You can use this callback to do additional postprocessing for a policy,

rllib/agents/trainer.py

@@ -345,6 +345,10 @@ COMMON_CONFIG = {
         "policy_mapping_fn": None,
         # Optional whitelist of policies to train, or None for all policies.
         "policies_to_train": None,
+        # Optional function that can be used to enhance the local agent
+        # observations to include more state.
+        # See rllib/evaluation/observation_function.py for more info.
+        "observation_fn": None,
     },
 }
 # __sphinx_doc_end__

rllib/evaluation/observation_function.py

@@ -0,0 +1,67 @@
+from typing import Dict
+
+from ray.rllib.env import BaseEnv
+from ray.rllib.policy import Policy, AgentID, PolicyID
+from ray.rllib.evaluation import MultiAgentEpisode, RolloutWorker
+from ray.rllib.utils.framework import TensorType
+
+
+class ObservationFunction:
+    """Interceptor function for rewriting observations from the environment.
+
+    These callbacks can be used for preprocessing of observations, especially
+    in multi-agent scenarios.
+
+    Observations functions can be specified in the multi-agent config by
+    specifying ``{"observation_function": your_obs_func}``. Note that
+    ``your_obs_func`` can be a plain Python function.
+
+    This API is **experimental**.
+    """
+
+    def __call__(self, agent_obs: Dict[AgentID, TensorType],
+                 worker: RolloutWorker, base_env: BaseEnv,
+                 policies: Dict[PolicyID, Policy], episode: MultiAgentEpisode,
+                 **kw) -> Dict[AgentID, TensorType]:
+        """Callback run on each environment step to observe the environment.
+
+        This method takes in the original agent observation dict returned by
+        a MultiAgentEnv, and returns a possibly modified one. It can be
+        thought of as a "wrapper" around the environment.
+
+        TODO(ekl): allow end-to-end differentiation through the observation
+            function and policy losses.
+
+        TODO(ekl): enable batch processing.
+
+        Args:
+            agent_obs (dict): Dictionary of default observations from the
+                environment. The default implementation of observe() simply
+                returns this dict.
+            worker (RolloutWorker): Reference to the current rollout worker.
+            base_env (BaseEnv): BaseEnv running the episode. The underlying
+                env object can be gotten by calling base_env.get_unwrapped().
+            policies (dict): Mapping of policy id to policy objects. In single
+                agent mode there will only be a single "default" policy.
+            episode (MultiAgentEpisode): Episode state object.
+            kwargs: Forward compatibility placeholder.
+
+        Returns:
+            new_agent_obs (dict): copy of agent obs with updates. You can
+                rewrite or drop data from the dict if needed (e.g., the env
+                can have a dummy "global" observation, and the observer can
+                merge the global state into individual observations.
+
+        Examples:
+            >>> # Observer that merges global state into individual obs. It is
+            ... # rewriting the discrete obs into a tuple with global state.
+            >>> example_obs_fn1({"a": 1, "b": 2, "global_state": 101}, ...)
+            {"a": [1, 101], "b": [2, 101]}
+
+            >>> # Observer for e.g., custom centralized critic model. It is
+            ... # rewriting the discrete obs into a dict with more data.
+            >>> example_obs_fn2({"a": 1, "b": 2}, ...)
+            {"a": {"self": 1, "other": 2}, "b": {"self": 2, "other": 1}}
+        """
+
+        return agent_obs

rllib/evaluation/rollout_worker.py

@@ -127,6 +127,7 @@ class RolloutWorker(EvaluatorInterface, ParallelIteratorWorker):
                  sample_async=False,
                  compress_observations=False,
                  num_envs=1,
+                 observation_fn=None,
                  observation_filter="NoFilter",
                  clip_rewards=None,
                  clip_actions=True,
@@ -147,8 +148,8 @@ class RolloutWorker(EvaluatorInterface, ParallelIteratorWorker):
                  soft_horizon=False,
                  no_done_at_end=False,
                  seed=None,
-                 _fake_sampler=False,
-                 extra_python_environs=None):
+                 extra_python_environs=None,
+                 _fake_sampler=False):
         """Initialize a rollout worker.
 
         Arguments:
@@ -194,6 +195,8 @@ class RolloutWorker(EvaluatorInterface, ParallelIteratorWorker):
             num_envs (int): If more than one, will create multiple envs
                 and vectorize the computation of actions. This has no effect if
                 if the env already implements VectorEnv.
+            observation_fn (ObservationFunction): Optional multi-agent
+                observation function.
             observation_filter (str): Name of observation filter to use.
             clip_rewards (bool): Whether to clip rewards to [-1, 1] prior to
                 experience postprocessing. Setting to None means clip for Atari
@@ -240,9 +243,9 @@ class RolloutWorker(EvaluatorInterface, ParallelIteratorWorker):
                 episode and instead record done=False.
             seed (int): Set the seed of both np and tf to this value to
                 to ensure each remote worker has unique exploration behavior.
-            _fake_sampler (bool): Use a fake (inf speed) sampler for testing.
             extra_python_environs (dict): Extra python environments need to
                 be set.
+            _fake_sampler (bool): Use a fake (inf speed) sampler for testing.
         """
         self._original_kwargs = locals().copy()
         del self._original_kwargs["self"]
@@ -463,7 +466,8 @@ class RolloutWorker(EvaluatorInterface, ParallelIteratorWorker):
                 clip_actions=clip_actions,
                 blackhole_outputs="simulation" in input_evaluation,
                 soft_horizon=soft_horizon,
-                no_done_at_end=no_done_at_end)
+                no_done_at_end=no_done_at_end,
+                observation_fn=observation_fn)
             self.sampler.start()
         else:
             self.sampler = SyncSampler(
@@ -481,7 +485,8 @@ class RolloutWorker(EvaluatorInterface, ParallelIteratorWorker):
                 tf_sess=self.tf_sess,
                 clip_actions=clip_actions,
                 soft_horizon=soft_horizon,
-                no_done_at_end=no_done_at_end)
+                no_done_at_end=no_done_at_end,
+                observation_fn=observation_fn)
 
         self.input_reader = input_creator(self.io_context)
         assert isinstance(self.input_reader, InputReader), self.input_reader

rllib/evaluation/sampler.py

@@ -77,7 +77,8 @@ class SyncSampler(SamplerInput):
                  tf_sess=None,
                  clip_actions=True,
                  soft_horizon=False,
-                 no_done_at_end=False):
+                 no_done_at_end=False,
+                 observation_fn=None):
         self.base_env = BaseEnv.to_base_env(env)
         self.rollout_fragment_length = rollout_fragment_length
         self.horizon = horizon
@@ -92,7 +93,7 @@ class SyncSampler(SamplerInput):
             self.policy_mapping_fn, self.rollout_fragment_length, self.horizon,
             self.preprocessors, self.obs_filters, clip_rewards, clip_actions,
             pack, callbacks, tf_sess, self.perf_stats, soft_horizon,
-            no_done_at_end)
+            no_done_at_end, observation_fn)
         self.metrics_queue = queue.Queue()
 
     def get_data(self):
@@ -140,7 +141,8 @@ class AsyncSampler(threading.Thread, SamplerInput):
                  clip_actions=True,
                  blackhole_outputs=False,
                  soft_horizon=False,
-                 no_done_at_end=False):
+                 no_done_at_end=False,
+                 observation_fn=None):
         for _, f in obs_filters.items():
             assert getattr(f, "is_concurrent", False), \
                 "Observation Filter must support concurrent updates."
@@ -167,6 +169,7 @@ class AsyncSampler(threading.Thread, SamplerInput):
         self.no_done_at_end = no_done_at_end
         self.perf_stats = PerfStats()
         self.shutdown = False
+        self.observation_fn = observation_fn
 
     def run(self):
         try:
@@ -188,7 +191,8 @@ class AsyncSampler(threading.Thread, SamplerInput):
             self.policy_mapping_fn, self.rollout_fragment_length, self.horizon,
             self.preprocessors, self.obs_filters, self.clip_rewards,
             self.clip_actions, self.pack, self.callbacks, self.tf_sess,
-            self.perf_stats, self.soft_horizon, self.no_done_at_end)
+            self.perf_stats, self.soft_horizon, self.no_done_at_end,
+            self.observation_fn)
         while not self.shutdown:
             # The timeout variable exists because apparently, if one worker
             # dies, the other workers won't die with it, unless the timeout is
@@ -233,7 +237,8 @@ class AsyncSampler(threading.Thread, SamplerInput):
 def _env_runner(worker, base_env, extra_batch_callback, policies,
                 policy_mapping_fn, rollout_fragment_length, horizon,
                 preprocessors, obs_filters, clip_rewards, clip_actions, pack,
-                callbacks, tf_sess, perf_stats, soft_horizon, no_done_at_end):
+                callbacks, tf_sess, perf_stats, soft_horizon, no_done_at_end,
+                observation_fn):
     """This implements the common experience collection logic.
 
     Args:
@@ -265,6 +270,8 @@ def _env_runner(worker, base_env, extra_batch_callback, policies,
             environment when the horizon is hit.
         no_done_at_end (bool): Ignore the done=True at the end of the episode
             and instead record done=False.
+        observation_fn (ObservationFunction): Optional multi-agent
+            observation func to use for preprocessing observations.
 
     Yields:
         rollout (SampleBatch): Object containing state, action, reward,
@@ -349,7 +356,7 @@ def _env_runner(worker, base_env, extra_batch_callback, policies,
             worker, base_env, policies, batch_builder_pool, active_episodes,
             unfiltered_obs, rewards, dones, infos, off_policy_actions, horizon,
             preprocessors, obs_filters, rollout_fragment_length, pack,
-            callbacks, soft_horizon, no_done_at_end)
+            callbacks, soft_horizon, no_done_at_end, observation_fn)
         perf_stats.processing_time += time.time() - t1
         for o in outputs:
             yield o
@@ -374,11 +381,11 @@ def _env_runner(worker, base_env, extra_batch_callback, policies,
         perf_stats.env_wait_time += time.time() - t4
 
 
-def _process_observations(worker, base_env, policies, batch_builder_pool,
-                          active_episodes, unfiltered_obs, rewards, dones,
-                          infos, off_policy_actions, horizon, preprocessors,
-                          obs_filters, rollout_fragment_length, pack,
-                          callbacks, soft_horizon, no_done_at_end):
+def _process_observations(
+        worker, base_env, policies, batch_builder_pool, active_episodes,
+        unfiltered_obs, rewards, dones, infos, off_policy_actions, horizon,
+        preprocessors, obs_filters, rollout_fragment_length, pack, callbacks,
+        soft_horizon, no_done_at_end, observation_fn):
     """Record new data from the environment and prepare for policy evaluation.
 
     Returns:
@@ -440,8 +447,21 @@ def _process_observations(worker, base_env, policies, batch_builder_pool,
             all_done = False
             active_envs.add(env_id)
 
+        # Custom observation function is applied before preprocessing.
+        if observation_fn:
+            agent_obs = observation_fn(
+                agent_obs=agent_obs,
+                worker=worker,
+                base_env=base_env,
+                policies=policies,
+                episode=episode)
+            if not isinstance(agent_obs, dict):
+                raise ValueError(
+                    "observe() must return a dict of agent observations")
+
         # For each agent in the environment.
         for agent_id, raw_obs in agent_obs.items():
+            assert agent_id != "__all__"
             policy_id = episode.policy_for(agent_id)
             prep_obs = _get_or_raise(preprocessors,
                                      policy_id).transform(raw_obs)
@@ -536,6 +556,13 @@ def _process_observations(worker, base_env, policies, batch_builder_pool,
                 # Creates a new episode if this is not async return
                 # If reset is async, we will get its result in some future poll
                 episode = active_episodes[env_id]
+                if observation_fn:
+                    resetted_obs = observation_fn(
+                        agent_obs=resetted_obs,
+                        worker=worker,
+                        base_env=base_env,
+                        policies=policies,
+                        episode=episode)
                 for agent_id, raw_obs in resetted_obs.items():
                     policy_id = episode.policy_for(agent_id)
                     policy = _get_or_raise(policies, policy_id)

rllib/evaluation/worker_set.py

@@ -257,6 +257,7 @@ class WorkerSet:
             sample_async=config["sample_async"],
             compress_observations=config["compress_observations"],
             num_envs=config["num_envs_per_worker"],
+            observation_fn=config["multiagent"]["observation_fn"],
             observation_filter=config["observation_filter"],
             clip_rewards=config["clip_rewards"],
             clip_actions=config["clip_actions"],

rllib/examples/centralized_critic_2.py

@@ -1,9 +1,9 @@
-"""An example of implementing a centralized critic by modifying the env.
+"""An example of implementing a centralized critic with ObservationFunction.
 
 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.
+change the algorithm at all -- just use callbacks and a custom model.
 However, it is a bit less principled in that you have to change the agent
-observation spaces and the environment.
+observation spaces to include data that is only used at train time.
 
 See also: centralized_critic.py for an alternative approach that instead
 modifies the policy to add a centralized value function.
@@ -14,7 +14,7 @@ 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.agents.callbacks import DefaultCallbacks
 from ray.rllib.examples.twostep_game import TwoStepGame
 from ray.rllib.models import ModelCatalog
 from ray.rllib.models.tf.tf_modelv2 import TFModelV2
@@ -70,62 +70,54 @@ class CentralizedCriticModel(TFModelV2):
         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),
-    })
+class FillInActions(DefaultCallbacks):
+    """Fills in the opponent actions info in the training batches."""
 
-    def __init__(self, env_config):
-        self.env = TwoStepGame(env_config)
+    def on_postprocess_trajectory(self, worker, episode, agent_id, policy_id,
+                                  policies, postprocessed_batch,
+                                  original_batches, **kwargs):
+        to_update = postprocessed_batch[SampleBatch.CUR_OBS]
+        other_id = 1 if agent_id == 0 else 0
+        action_encoder = ModelCatalog.get_preprocessor_for_space(Discrete(2))
 
-    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
-            },
-        }
+        # set the opponent actions into the observation
+        _, opponent_batch = original_batches[other_id]
+        opponent_actions = np.array([
+            action_encoder.transform(a)
+            for a in opponent_batch[SampleBatch.ACTIONS]
+        ])
+        to_update[:, -2:] = opponent_actions
 
 
-def fill_in_actions(info):
-    """Callback that saves opponent actions into the agent obs.
+def central_critic_observer(agent_obs, **kw):
+    """Rewrites the agent obs to include opponent data for training."""
 
-    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
+    new_obs = {
+        0: {
+            "own_obs": agent_obs[0],
+            "opponent_obs": agent_obs[1],
+            "opponent_action": 0,  # filled in by FillInActions
+        },
+        1: {
+            "own_obs": agent_obs[1],
+            "opponent_obs": agent_obs[0],
+            "opponent_action": 0,  # filled in by FillInActions
+        },
+    }
+    return new_obs
 
 
 if __name__ == "__main__":
     args = parser.parse_args()
     ModelCatalog.register_custom_model("cc_model", CentralizedCriticModel)
+    action_space = Discrete(2)
+    observer_space = Dict({
+        "own_obs": Discrete(6),
+        # These two fields are filled in by the CentralCriticObserver, and are
+        # not used for inference, only for training.
+        "opponent_obs": Discrete(6),
+        "opponent_action": Discrete(2),
+    })
     tune.run(
         "PPO",
         stop={
@@ -133,20 +125,17 @@ if __name__ == "__main__":
             "episode_reward_mean": 7.99,
         },
         config={
-            "env": GlobalObsTwoStepGame,
+            "env": TwoStepGame,
             "batch_mode": "complete_episodes",
-            "callbacks": {
-                "on_postprocess_traj": fill_in_actions,
-            },
+            "callbacks": FillInActions,
             "num_workers": 0,
             "multiagent": {
                 "policies": {
-                    "pol1": (None, GlobalObsTwoStepGame.observation_space,
-                             GlobalObsTwoStepGame.action_space, {}),
-                    "pol2": (None, GlobalObsTwoStepGame.observation_space,
-                             GlobalObsTwoStepGame.action_space, {}),
+                    "pol1": (None, observer_space, action_space, {}),
+                    "pol2": (None, observer_space, action_space, {}),
                 },
                 "policy_mapping_fn": lambda x: "pol1" if x == 0 else "pol2",
+                "observation_fn": central_critic_observer,
             },
             "model": {
                 "custom_model": "cc_model",

rllib/policy/__init__.py

@@ -1,10 +1,12 @@
-from ray.rllib.policy.policy import Policy
+from ray.rllib.policy.policy import Policy, PolicyID, AgentID
 from ray.rllib.policy.torch_policy import TorchPolicy
 from ray.rllib.policy.tf_policy import TFPolicy
 from ray.rllib.policy.torch_policy_template import build_torch_policy
 from ray.rllib.policy.tf_policy_template import build_tf_policy
 
 __all__ = [
+    "AgentID",
+    "PolicyID",
     "Policy",
     "TFPolicy",
     "TorchPolicy",

rllib/policy/policy.py

@@ -1,6 +1,7 @@
 from abc import ABCMeta, abstractmethod
 import gym
 import numpy as np
+from typing import Any
 
 from ray.rllib.utils import try_import_tree
 from ray.rllib.utils.annotations import DeveloperAPI
@@ -14,6 +15,12 @@ tree = try_import_tree()
 # `grad_info` dict returned by learn_on_batch() / compute_grads() via this key.
 LEARNER_STATS_KEY = "learner_stats"
 
+# Represents a generic identifier for an agent (e.g., "agent1").
+AgentID = Any
+
+# Represents a generic identifier for a policy (e.g., "pol1").
+PolicyID = str
+
 
 @DeveloperAPI
 class Policy(metaclass=ABCMeta):