diff --git a/rllib/BUILD b/rllib/BUILD index b91a903b3..cb0f22b7e 100644 --- a/rllib/BUILD +++ b/rllib/BUILD @@ -1059,6 +1059,13 @@ py_test( srcs = ["policy/tests/test_compute_log_likelihoods.py"] ) +py_test( + name = "policy/tests/test_trajectory_view_api", + tags = ["policy"], + size = "small", + srcs = ["policy/tests/test_trajectory_view_api.py"] +) + # -------------------------------------------------------------------- # Utils: # rllib/utils/ diff --git a/rllib/agents/ppo/ppo_torch_policy.py b/rllib/agents/ppo/ppo_torch_policy.py index de23a26e2..3bab45ac7 100644 --- a/rllib/agents/ppo/ppo_torch_policy.py +++ b/rllib/agents/ppo/ppo_torch_policy.py @@ -10,6 +10,7 @@ from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.policy.torch_policy import EntropyCoeffSchedule, \ LearningRateSchedule from ray.rllib.policy.torch_policy_template import build_torch_policy +from ray.rllib.policy.view_requirement import ViewRequirement from ray.rllib.utils.framework import try_import_torch from ray.rllib.utils.torch_ops import convert_to_torch_tensor, \ explained_variance, sequence_mask @@ -216,6 +217,15 @@ def setup_mixins(policy, obs_space, action_space, config): LearningRateSchedule.__init__(policy, config["lr"], config["lr_schedule"]) +def training_view_requirements_fn(policy): + return { + # Next obs are needed for PPO postprocessing. + SampleBatch.NEXT_OBS: ViewRequirement(SampleBatch.OBS, shift=1), + # VF preds are needed for the loss. + SampleBatch.VF_PREDS: ViewRequirement(shift=0), + } + + PPOTorchPolicy = build_torch_policy( name="PPOTorchPolicy", get_default_config=lambda: ray.rllib.agents.ppo.ppo.DEFAULT_CONFIG, @@ -229,4 +239,6 @@ PPOTorchPolicy = build_torch_policy( mixins=[ LearningRateSchedule, EntropyCoeffSchedule, KLCoeffMixin, ValueNetworkMixin - ]) + ], + training_view_requirements_fn=training_view_requirements_fn, +) diff --git a/rllib/evaluation/sample_batch_builder.py b/rllib/evaluation/sample_batch_builder.py index f5e3b002f..bf1888d07 100644 --- a/rllib/evaluation/sample_batch_builder.py +++ b/rllib/evaluation/sample_batch_builder.py @@ -25,6 +25,8 @@ def to_float_array(v: List[Any]) -> np.ndarray: return arr +# TODO(sven): Remove the following class once we switch to trajectory view API. + @PublicAPI class SampleBatchBuilder: """Util to build a SampleBatch incrementally. @@ -72,6 +74,8 @@ class SampleBatchBuilder: return batch +# TODO(sven): Remove the following class once we switch to trajectory view API. + @DeveloperAPI class MultiAgentSampleBatchBuilder: """Util to build SampleBatches for each policy in a multi-agent env. diff --git a/rllib/evaluation/sample_collector.py b/rllib/evaluation/sample_collector.py new file mode 100644 index 000000000..6880e8bd0 --- /dev/null +++ b/rllib/evaluation/sample_collector.py @@ -0,0 +1,180 @@ +from abc import abstractmethod, ABCMeta +import logging +from typing import Dict, Optional + +from ray.rllib.evaluation.episode import MultiAgentEpisode +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.utils.types import AgentID, EpisodeID, PolicyID, \ + TensorType + +logger = logging.getLogger(__name__) + + +class _SampleCollector(metaclass=ABCMeta): + """Collects samples for all policies and agents from a multi-agent env. + + Note: This is an experimental class only used when + `config._use_trajectory_view_api` = True. + Once `_use_trajectory_view_api` becomes the default in configs: + This class will deprecate the `SampleBatchBuilder` and + `MultiAgentBatchBuilder` classes. + + This API is controlled by RolloutWorker objects to store all data + generated by Environments and Policies/Models during rollout and + postprocessing. It's purposes are to a) make data collection and + SampleBatch/input_dict generation from this data faster, b) to unify + the way we collect samples from environments and model (outputs), thereby + allowing for possible user customizations, c) to allow for more complex + inputs fed into different policies (e.g. multi-agent case with inter-agent + communication channel). + """ + + @abstractmethod + def add_init_obs(self, episode_id: EpisodeID, agent_id: AgentID, + policy_id: PolicyID, init_obs: TensorType) -> None: + """Adds an initial obs (after reset) to this collector. + + Since the very first observation in an environment is collected w/o + additional data (w/o actions, w/o reward) after env.reset() is called, + this method initializes a new trajectory for a given agent. + `add_init_obs()` has to be called first for each agent/episode-ID + combination. After this, only `add_action_reward_next_obs()` must be + called for that same agent/episode-pair. + + Args: + episode_id (EpisodeID): Unique id for the episode we are adding + values for. + agent_id (AgentID): Unique id for the agent we are adding + values for. + policy_id (PolicyID): Unique id for policy controlling the agent. + init_obs (TensorType): Initial observation (after env.reset()). + + Examples: + >>> obs = env.reset() + >>> collector.add_init_obs(12345, 0, "pol0", obs) + >>> obs, r, done, info = env.step(action) + >>> collector.add_action_reward_next_obs(12345, 0, "pol0", { + ... "action": action, "obs": obs, "reward": r, "done": done + ... }) + """ + raise NotImplementedError + + @abstractmethod + def add_action_reward_next_obs( + self, + episode_id: EpisodeID, + agent_id: AgentID, + policy_id: PolicyID, + values: Dict[str, TensorType]) -> None: + """Add the given dictionary (row) of values to this collector. + + The incoming data (`values`) must include action, reward, done, and + next_obs information and may include any other information. + For the initial observation (after Env.reset()) of the given agent/ + episode-ID combination, `add_initial_obs()` must be called instead. + + Args: + episode_id (EpisodeID): Unique id for the episode we are adding + values for. + agent_id (AgentID): Unique id for the agent we are adding + values for. + policy_id (PolicyID): Unique id for policy controlling the agent. + values (Dict[str, TensorType]): Row of values to add for this + agent. This row must contain the keys SampleBatch.ACTION, + REWARD, NEW_OBS, and DONE. + + Examples: + >>> obs = env.reset() + >>> collector.add_init_obs(12345, 0, "pol0", obs) + >>> obs, r, done, info = env.step(action) + >>> collector.add_action_reward_next_obs(12345, 0, "pol0", { + ... "action": action, "obs": obs, "reward": r, "done": done + ... }) + """ + raise NotImplementedError + + @abstractmethod + def total_env_steps(self) -> int: + """Returns total number of steps taken in the env (sum of all agents). + + Returns: + int: The number of steps taken in total in the environment over all + agents. + """ + raise NotImplementedError + + @abstractmethod + def get_inference_input_dict(self, model: ModelV2) -> \ + Dict[str, TensorType]: + """Returns input_dict for an inference forward pass from our data. + + The input_dict can then be used for action computations. + + Args: + model (ModelV2): The ModelV2 object for which to generate the view + (input_dict) from `data`. + + Returns: + Dict[str, TensorType]: The input_dict to be passed into the ModelV2 + for inference/training. + + Examples: + >>> obs, r, done, info = env.step(action) + >>> collector.add_action_reward_next_obs(12345, 0, "pol0", { + ... "action": action, "obs": obs, "reward": r, "done": done + ... }) + >>> input_dict = collector.get_inference_input_dict(policy.model) + >>> action = policy.compute_actions_from_input_dict(input_dict) + >>> # repeat + """ + raise NotImplementedError + + @abstractmethod + def has_non_postprocessed_data(self) -> bool: + """Returns whether there is pending, unprocessed data. + + Returns: + bool: True if there is at least some data that has not been + postprocessed yet. + """ + raise NotImplementedError + + @abstractmethod + def postprocess_trajectories_so_far( + self, episode: Optional[MultiAgentEpisode] = None) -> None: + """Apply postprocessing to unprocessed data (in one or all episodes). + + Generates (single-trajectory) SampleBatches for all Policies/Agents and + calls Policy.postprocess_trajectory on each of these. Postprocessing + may happens in-place, meaning any changes to the viewed data columns + are directly reflected inside this collector's buffers. + Also makes sure that additional (newly created) data columns are + correctly added to the buffers. + + Args: + episode (Optional[MultiAgentEpisode]): The Episode object for which + to post-process data. If not provided, postprocess data for all + episodes. + """ + raise NotImplementedError + + @abstractmethod + def get_multi_agent_batch_and_reset(self): + """Returns the accumulated sample batches for each policy. + + Any unprocessed rows will be first postprocessed with a policy + postprocessor. The internal state of this builder will be reset. + + Args: + episode (Optional[MultiAgentEpisode]): The Episode object that + holds this MultiAgentBatchBuilder object or None. + + Returns: + MultiAgentBatch: Returns the accumulated sample batches for each + policy inside one MultiAgentBatch object. + """ + raise NotImplementedError + + @abstractmethod + def check_missing_dones(self, episode_id: EpisodeID) -> None: + raise NotImplementedError diff --git a/rllib/evaluation/trajectory.py b/rllib/evaluation/trajectory.py deleted file mode 100644 index 6e49fbf0a..000000000 --- a/rllib/evaluation/trajectory.py +++ /dev/null @@ -1,267 +0,0 @@ -import logging -import numpy as np -from typing import Dict, Optional - -from ray.rllib.policy.sample_batch import SampleBatch -from ray.rllib.utils.framework import try_import_tf, try_import_torch -from ray.rllib.utils.types import AgentID, EnvID, PolicyID, TensorType - -tf1, tf, tfv = try_import_tf() -torch, _ = try_import_torch() - -logger = logging.getLogger(__name__) - - -def to_float_array(v): - if torch and isinstance(v[0], torch.Tensor): - arr = torch.stack(v).numpy() # np.array([s.numpy() for s in v]) - else: - arr = np.array(v) - if arr.dtype == np.float64: - return arr.astype(np.float32) # save some memory - return arr - - -class Trajectory: - """A trajectory of a (single) agent throughout one episode. - - Note: This is an experimental class only used when - `config._use_trajectory_view_api` = True. - - Collects all data produced by the environment during stepping of the agent - as well as all model outputs associated with the agent's Policy into - pre-allocated buffers of n timesteps capacity (`self.buffer_size`). - NOTE: A Trajectory object may contain remainders of a previous trajectory, - however, these are only kept for avoiding memory re-allocations. A - convenience cursor and offset-pointers allow for only "viewing" the - currently ongoing trajectory. - Memory re-allocation into larger buffers (`self.buffer_size *= 2`) only - happens if unavoidable (in case the buffer is full AND the currently - ongoing trajectory (episode) takes more than half of the buffer). In all - other cases, the same buffer is used for succeeding episodes/trejactories - (even for different agents). - """ - - # Disambiguate unrolls within a single episode. - _next_unroll_id = 0 - - def __init__(self, buffer_size: Optional[int] = None): - """Initializes a Trajectory object. - - Args: - buffer_size (Optional[int]): The max number of timesteps to - fit into one buffer column. When re-allocating - """ - # The current occupant (agent X in env Y using policy Z) of our - # buffers. - self.env_id: EnvID = None - self.agent_id: AgentID = None - self.policy_id: PolicyID = None - - # Determine the size of the initial buffers. - self.buffer_size = buffer_size or 1000 - # The actual buffer holding dict (by column name (str) -> - # numpy/torch/tf tensors). - self.buffers = {} - - # Holds the initial observation data. - self.initial_obs = None - - # Cursor into the preallocated buffers. This is where all new data - # gets inserted. - self.cursor: int = 0 - # The offset inside our buffer where the current trajectory starts. - self.trajectory_offset: int = 0 - # The offset inside our buffer, from where to build the next - # SampleBatch. - self.sample_batch_offset: int = 0 - - @property - def timestep(self) -> int: - """The timestep in the (currently ongoing) trajectory/episode.""" - return self.cursor - self.trajectory_offset - - def add_init_obs(self, - env_id: EnvID, - agent_id: AgentID, - policy_id: PolicyID, - init_obs: TensorType) -> None: - """Adds a single initial observation (after env.reset()) to the buffer. - - Stores it in self.initial_obs. - - Args: - env_id (EnvID): Unique id for the episode we are adding the initial - observation for. - agent_id (AgentID): Unique id for the agent we are adding the - initial observation for. - policy_id (PolicyID): Unique id for policy controlling the agent. - init_obs (TensorType): Initial observation (after env.reset()). - """ - self.env_id = env_id - self.agent_id = agent_id - self.policy_id = policy_id - self.initial_obs = init_obs - - def add_action_reward_next_obs(self, - env_id: EnvID, - agent_id: AgentID, - policy_id: PolicyID, - values: Dict[str, TensorType]) -> None: - """Add the given dictionary (row) of values to this batch. - - Args: - env_id (EnvID): Unique id for the episode we are adding the initial - observation for. - agent_id (AgentID): Unique id for the agent we are adding the - initial observation for. - policy_id (PolicyID): Unique id for policy controlling the agent. - values (Dict[str, TensorType]): Data dict (interpreted as a single - row) to be added to buffer. Must contain keys: - SampleBatch.ACTIONS, REWARDS, DONES, and OBS. - """ - assert self.initial_obs is not None - assert (SampleBatch.ACTIONS in values and SampleBatch.REWARDS in values - and SampleBatch.NEXT_OBS in values) - assert env_id == self.env_id - assert agent_id == self.agent_id - assert policy_id == self.policy_id - - # Only obs exists so far in buffers: - # Initialize all other columns. - if len(self.buffers) == 0: - self._build_buffers(single_row=values) - - for k, v in values.items(): - self.buffers[k][self.cursor] = v - self.cursor += 1 - - # Extend (re-alloc) buffers if full. - if self.cursor == self.buffer_size: - self._extend_buffers(values) - - def get_sample_batch_and_reset(self) -> SampleBatch: - """Returns a SampleBatch carrying all previously added data. - - If a reset happens and the trajectory is not done yet, we'll keep the - entire ongoing trajectory in memory for Model view requirement purposes - and only actually free the data, once the episode ends. - - Returns: - SampleBatch: The SampleBatch containing this agent's data for the - entire trajectory (so far). The trajectory may not be - terminated yet. This SampleBatch object will contain a - `_last_obs` property, which contains the last observation for - this agent. This should be used by postprocessing functions - instead of the SampleBatch.NEXT_OBS field, which is deprecated. - """ - assert SampleBatch.UNROLL_ID not in self.buffers - - # Convert all our data to numpy arrays, compress float64 to float32, - # and add the last observation data as well (always one more obs than - # all other columns due to the additional obs returned by Env.reset()). - data = {} - for k, v in self.buffers.items(): - data[k] = to_float_array( - v[self.sample_batch_offset:self.cursor]) - - # Add unroll ID column to batch if non-existent. - uid = Trajectory._next_unroll_id - data[SampleBatch.UNROLL_ID] = np.repeat( - uid, self.cursor - self.sample_batch_offset) - - inputs = {uid: {}} - if "t" in self.buffers: - if self.buffers["t"][self.sample_batch_offset] > 0: - for k in self.buffers.keys(): - inputs[uid][k] = \ - self.buffers[k][self.sample_batch_offset - 1] - else: - inputs[uid][SampleBatch.NEXT_OBS] = self.initial_obs - else: - inputs[uid][SampleBatch.NEXT_OBS] = self.initial_obs - - Trajectory._next_unroll_id += 1 - - batch = SampleBatch(data, _initial_inputs=inputs) - - # If done at end -> We can reset our buffers entirely. - if self.buffers[SampleBatch.DONES][self.cursor - 1]: - # Set self.timestep to 0 -> new trajectory w/o re-alloc (not yet, - # only ever re-alloc when necessary). - self.trajectory_offset = self.sample_batch_offset = self.cursor - # No done at end -> leave trajectory_offset as is (trajectory is still - # ongoing), but move the sample_batch offset to cursor. - else: - self.sample_batch_offset = self.cursor - return batch - - def _build_buffers(self, single_row): - """Creates zero-filled pre-allocated numpy buffers for data collection. - - Except for the obs-column, which should already be initialized (done - on call to `self.add_initial_observation()`). - - Args: - single_row (Dict[str,np.ndarray]): Dict of column names (keys) and - sample numpy data (values). Note: Only one of `single_data` or - `data_batch` must be provided. - """ - for col, data in single_row.items(): - # Skip already initialized ones, e.g. 'obs' if used with - # add_initial_observation. - if col in self.buffers: - continue - self.buffers[col] = [None] * self.buffer_size - - def _extend_buffers(self, single_row): - """Extends the buffers (depending on trajectory state/length). - - - Extend all buffer lists (x2) if trajectory starts at 0 (trajectory is - longer than current self.buffer_size). - - Trajectory starts in first half of buffer: Create new buffer lists - (2x buffer sizes) and move Trajectory to beginning of new buffer. - - Trajectory starts in last half of buffer: Leave buffer as is, but - move trajectory to very front (cursor=0). - - Args: - single_row (dict): Data dict example to use in case we have to - re-build buffer. - """ - traj_length = self.cursor - self.trajectory_offset - - # Trajectory starts at 0 (meaning episodes are longer than current - # `self.buffer_size` -> Simply do a resize (enlarge) on each column - # in the buffer. - if self.trajectory_offset == 0: - # Double actual horizon. - for col, data in self.buffers.items(): - self.buffers[col].extend([None] * self.buffer_size) - self.buffer_size *= 2 - - # Trajectory starts in first half of the buffer -> Reallocate a new - # buffer and copy the currently ongoing trajectory into the new buffer. - elif self.trajectory_offset < self.buffer_size / 2: - # Double actual horizon. - self.buffer_size *= 2 - # Store currently ongoing trajectory and build a new buffer. - old_buffers = self.buffers - self.buffers = {} - self._build_buffers(single_row) - # Copy the still ongoing trajectory into the new buffer. - for col, data in old_buffers.items(): - self.buffers[col][:traj_length] = data[self.trajectory_offset: - self.cursor] - - # Do an efficient memory swap: Move current trajectory simply to - # the beginning of the buffer (no reallocation/None-padding necessary). - else: - for col, data in self.buffers.items(): - self.buffers[col][:traj_length] = self.buffers[col][ - self.trajectory_offset:self.cursor] - - # Set all pointers to their correct new values. - self.sample_batch_offset = ( - self.sample_batch_offset - self.trajectory_offset) - self.trajectory_offset = 0 - self.cursor = traj_length diff --git a/rllib/models/modelv2.py b/rllib/models/modelv2.py index f66e29cc1..a8b14bf39 100644 --- a/rllib/models/modelv2.py +++ b/rllib/models/modelv2.py @@ -8,7 +8,7 @@ from ray.rllib.models.preprocessors import get_preprocessor, \ RepeatedValuesPreprocessor from ray.rllib.models.repeated_values import RepeatedValues from ray.rllib.policy.sample_batch import SampleBatch -from ray.rllib.policy.trajectory_view import ViewRequirement +from ray.rllib.policy.view_requirement import ViewRequirement from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI from ray.rllib.utils.framework import try_import_tf, try_import_torch, \ TensorType @@ -246,35 +246,24 @@ class ModelV2: i += 1 return self.__call__(input_dict, states, train_batch.get("seq_lens")) - def get_view_requirements( - self, is_training: bool = False) -> Dict[str, ViewRequirement]: - """Returns a list of ViewRequirements for this Model (or None). + def inference_view_requirements(self) -> Dict[str, ViewRequirement]: + """Returns a dict of ViewRequirements for this Model. Note: This is an experimental API method. - A ViewRequirement object tells the caller of this Model, which - data at which timesteps are needed by this Model. This could be a - sequence of past observations, internal-states, previous rewards, or - other episode data/previous model outputs. - - Args: - is_training (bool): Whether the returned requirements are for - training or inference (default). + The view requirements dict is used to generate input_dicts and + train batches for 1) action computations, 2) postprocessing, and 3) + generating training batches. Returns: - Dict[str, ViewRequirement]: The view requirements as a dict mapping - column names e.g. "obs" to config dicts containing supported - fields. - TODO: (sven) Currently only `timesteps==0` can be setup. + Dict[str, ViewRequirement]: The view requirements dict, mapping + each view key (which will be available in input_dicts) to + an underlying requirement (actual data, timestep shift, etc..). """ # Default implementation for simple RL model: # Single requirement: Pass current obs as input. return { - SampleBatch.CUR_OBS: ViewRequirement(timesteps=0), - SampleBatch.PREV_ACTIONS: ViewRequirement( - SampleBatch.ACTIONS, timesteps=-1), - SampleBatch.PREV_REWARDS: ViewRequirement( - SampleBatch.REWARDS, timesteps=-1), + SampleBatch.OBS: ViewRequirement(shift=0), } def import_from_h5(self, h5_file: str) -> None: diff --git a/rllib/models/torch/recurrent_net.py b/rllib/models/torch/recurrent_net.py index 24c3fad2e..27980e5ac 100644 --- a/rllib/models/torch/recurrent_net.py +++ b/rllib/models/torch/recurrent_net.py @@ -1,10 +1,12 @@ import numpy as np +from typing import Dict from ray.rllib.models.modelv2 import ModelV2 from ray.rllib.models.torch.misc import SlimFC from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 from ray.rllib.policy.rnn_sequencing import add_time_dimension from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.view_requirement import ViewRequirement from ray.rllib.utils.annotations import override, DeveloperAPI from ray.rllib.utils.framework import try_import_torch @@ -169,3 +171,16 @@ class LSTMWrapper(RecurrentNetwork, nn.Module): def value_function(self): assert self._features is not None, "must call forward() first" return torch.reshape(self._value_branch(self._features), [-1]) + + @override(ModelV2) + def inference_view_requirements(self) -> Dict[str, ViewRequirement]: + req = super().inference_view_requirements() + # Optional: prev-actions/rewards for forward pass. + if self.model_config["lstm_use_prev_action_reward"]: + req.update({ + SampleBatch.PREV_REWARDS: ViewRequirement( + SampleBatch.REWARDS, shift=-1), + SampleBatch.PREV_ACTIONS: ViewRequirement( + SampleBatch.ACTIONS, space=self.action_space, shift=-1), + }) + return req diff --git a/rllib/policy/policy.py b/rllib/policy/policy.py index 4fa5df518..580f0bf41 100644 --- a/rllib/policy/policy.py +++ b/rllib/policy/policy.py @@ -214,26 +214,27 @@ class Policy(metaclass=ABCMeta): return single_action, [s[0] for s in state_out], \ {k: v[0] for k, v in info.items()} - def compute_actions_from_trajectories( + def compute_actions_from_input_dict( self, - trajectories: List["Trajectory"], - other_trajectories: Optional[Dict[AgentID, "Trajectory"]] = None, + input_dict: Dict[str, TensorType], explore: bool = None, timestep: Optional[int] = None, **kwargs) -> \ Tuple[TensorType, List[TensorType], Dict[str, TensorType]]: - """Computes actions for the current policy based on . + """Computes actions from collected samples (across multiple-agents). Note: This is an experimental API method. Only used so far by the Sampler iff `_use_trajectory_view_api=True` (also only supported for torch). + Uses the currently "forward-pass-registered" samples from the collector + to construct the input_dict for the Model. Args: - trajectories (List[Trajectory]): A List of Trajectory data used - to create a view for the Model forward call. - other_trajectories (Optional[Dict[AgentID, Trajectory]]): Optional - dict mapping AgentIDs to Trajectory objects. + input_dict (Dict[str, TensorType]): An input dict mapping str + keys to Tensors. `input_dict` already abides to the Policy's + as well as the Model's view requirements and can be passed + to the Model as-is. explore (bool): Whether to pick an exploitation or exploration action (default: None -> use self.config["explore"]). timestep (Optional[int]): The current (sampling) time step. @@ -282,6 +283,25 @@ class Policy(metaclass=ABCMeta): """ raise NotImplementedError + @DeveloperAPI + def training_view_requirements(self): + """Returns a dict of view requirements for operating on this Policy. + + Note: This is an experimental API method. + + The view requirements dict is used to generate input_dicts and + SampleBatches for 1) action computations, 2) postprocessing, and 3) + generating training batches. + The Policy may ask its Model(s) as well for possible additional + requirements (e.g. prev-action/reward in an LSTM). + + Returns: + Dict[str, ViewRequirement]: The view requirements dict, mapping + each view key (which will be available in input_dicts) to + an underlying requirement (actual data, timestep shift, etc..). + """ + return {} + @DeveloperAPI def postprocess_trajectory( self, diff --git a/rllib/policy/sample_batch.py b/rllib/policy/sample_batch.py index 0ef02be92..17cfb1e7f 100644 --- a/rllib/policy/sample_batch.py +++ b/rllib/policy/sample_batch.py @@ -26,6 +26,7 @@ class SampleBatch: """ # Outputs from interacting with the environment + OBS = "obs" CUR_OBS = "obs" NEXT_OBS = "new_obs" ACTIONS = "actions" @@ -40,7 +41,7 @@ class SampleBatch: ACTION_PROB = "action_prob" ACTION_LOGP = "action_logp" - # Uniquely identifies an episode + # Uniquely identifies an episode. EPS_ID = "eps_id" # Uniquely identifies a sample batch. This is important to distinguish RNN @@ -48,10 +49,10 @@ class SampleBatch: # concatenated (fusing sequences across batches can be unsafe). UNROLL_ID = "unroll_id" - # Uniquely identifies an agent within an episode + # Uniquely identifies an agent within an episode. AGENT_INDEX = "agent_index" - # Value function predictions emitted by the behaviour policy + # Value function predictions emitted by the behaviour policy. VF_PREDS = "vf_preds" @PublicAPI diff --git a/rllib/policy/tests/test_trajectory_view_api.py b/rllib/policy/tests/test_trajectory_view_api.py new file mode 100644 index 000000000..9d728f4b7 --- /dev/null +++ b/rllib/policy/tests/test_trajectory_view_api.py @@ -0,0 +1,81 @@ +import unittest + +import ray +import ray.rllib.agents.ppo as ppo +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.test_utils import framework_iterator + + +class TestTrajectoryViewAPI(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + ray.init() + + @classmethod + def tearDownClass(cls) -> None: + ray.shutdown() + + def test_plain(self): + config = ppo.DEFAULT_CONFIG.copy() + for _ in framework_iterator(config, frameworks="torch"): + trainer = ppo.PPOTrainer(config, env="CartPole-v0") + policy = trainer.get_policy() + view_req_model = policy.model.inference_view_requirements() + view_req_policy = policy.training_view_requirements() + assert len(view_req_model) == 1 + assert len(view_req_policy) == 6 + for key in [ + SampleBatch.OBS, SampleBatch.ACTIONS, SampleBatch.REWARDS, + SampleBatch.DONES, SampleBatch.NEXT_OBS, SampleBatch.VF_PREDS + ]: + assert key in view_req_policy + # None of the view cols has a special underlying data_col, + # except next-obs. + if key != SampleBatch.NEXT_OBS: + assert view_req_policy[key].data_col is None + else: + assert view_req_policy[key].data_col == SampleBatch.OBS + assert view_req_policy[key].shift == 1 + trainer.stop() + + def test_lstm_prev_actions_and_rewards(self): + config = ppo.DEFAULT_CONFIG.copy() + config["model"] = config["model"].copy() + # Activate LSTM + prev-action + rewards. + config["model"]["use_lstm"] = True + config["model"]["lstm_use_prev_action_reward"] = True + + for _ in framework_iterator(config, frameworks="torch"): + trainer = ppo.PPOTrainer(config, env="CartPole-v0") + policy = trainer.get_policy() + view_req_model = policy.model.inference_view_requirements() + view_req_policy = policy.training_view_requirements() + assert len(view_req_model) == 3 # obs, prev_a, prev_r + assert len(view_req_policy) == 8 + for key in [ + SampleBatch.OBS, SampleBatch.ACTIONS, SampleBatch.REWARDS, + SampleBatch.DONES, SampleBatch.NEXT_OBS, SampleBatch.VF_PREDS, + SampleBatch.PREV_ACTIONS, SampleBatch.PREV_REWARDS + ]: + assert key in view_req_policy + + if key == SampleBatch.PREV_ACTIONS: + assert view_req_policy[key].data_col == SampleBatch.ACTIONS + assert view_req_policy[key].shift == -1 + elif key == SampleBatch.PREV_REWARDS: + assert view_req_policy[key].data_col == SampleBatch.REWARDS + assert view_req_policy[key].shift == -1 + elif key not in [SampleBatch.NEXT_OBS, + SampleBatch.PREV_ACTIONS, + SampleBatch.PREV_REWARDS]: + assert view_req_policy[key].data_col is None + else: + assert view_req_policy[key].data_col == SampleBatch.OBS + assert view_req_policy[key].shift == 1 + trainer.stop() + + +if __name__ == "__main__": + import pytest + import sys + sys.exit(pytest.main(["-v", __file__])) diff --git a/rllib/policy/tf_policy_template.py b/rllib/policy/tf_policy_template.py index 8f31a1ea5..910ef9c94 100644 --- a/rllib/policy/tf_policy_template.py +++ b/rllib/policy/tf_policy_template.py @@ -271,6 +271,22 @@ def build_tf_policy(name: str, return base.extra_compute_grad_fetches(self) def with_updates(**overrides): + """Allows creating a TFPolicy cls based on settings of another one. + + Keyword Args: + **overrides: The settings (passed into `build_tf_policy`) that + should be different from the class that this method is called + on. + + Returns: + type: A new TFPolicy sub-class. + + Examples: + >> MySpecialDQNPolicyClass = DQNTFPolicy.with_updates( + .. name="MySpecialDQNPolicyClass", + .. loss_function=[some_new_loss_function], + .. ) + """ return build_tf_policy(**dict(original_kwargs, **overrides)) def as_eager(): diff --git a/rllib/policy/torch_policy.py b/rllib/policy/torch_policy.py index de7739724..53002b5fd 100644 --- a/rllib/policy/torch_policy.py +++ b/rllib/policy/torch_policy.py @@ -11,7 +11,7 @@ from ray.rllib.models.torch.torch_action_dist import TorchDistributionWrapper from ray.rllib.policy.policy import Policy, LEARNER_STATS_KEY from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.policy.rnn_sequencing import pad_batch_to_sequences_of_same_size -from ray.rllib.policy.trajectory_view import get_trajectory_view +from ray.rllib.policy.view_requirement import ViewRequirement from ray.rllib.utils import force_list from ray.rllib.utils.annotations import override, DeveloperAPI from ray.rllib.utils.framework import try_import_torch @@ -19,7 +19,7 @@ from ray.rllib.utils.schedules import ConstantSchedule, PiecewiseSchedule from ray.rllib.utils.torch_ops import convert_to_non_torch_type, \ convert_to_torch_tensor from ray.rllib.utils.tracking_dict import UsageTrackingDict -from ray.rllib.utils.types import AgentID, ModelGradients, ModelWeights, \ +from ray.rllib.utils.types import ModelGradients, ModelWeights, \ TensorType, TrainerConfigDict torch, _ = try_import_torch() @@ -103,6 +103,11 @@ class TorchPolicy(Policy): else: self.device = torch.device("cpu") self.model = model.to(self.device) + # Combine view_requirements for Model and Policy. + self.view_requirements = { + **self.model.inference_view_requirements(), + **self.training_view_requirements(), + } self.exploration = self._create_exploration() self.unwrapped_model = model # used to support DistributedDataParallel self._loss = loss @@ -116,9 +121,20 @@ class TorchPolicy(Policy): self.distributed_world_size = None self.max_seq_len = max_seq_len - self.batch_divisibility_req = \ - get_batch_divisibility_req(self) if get_batch_divisibility_req \ - else 1 + self.batch_divisibility_req = get_batch_divisibility_req(self) if \ + callable(get_batch_divisibility_req) else \ + (get_batch_divisibility_req or 1) + + @override(Policy) + def training_view_requirements(self): + if hasattr(self, "view_requirements"): + return self.view_requirements + return { + SampleBatch.ACTIONS: ViewRequirement( + space=self.action_space, shift=0), + SampleBatch.REWARDS: ViewRequirement(shift=0), + SampleBatch.DONES: ViewRequirement(shift=0), + } @override(Policy) @DeveloperAPI @@ -168,10 +184,9 @@ class TorchPolicy(Policy): extra_fetches)) @override(Policy) - def compute_actions_from_trajectories( + def compute_actions_from_input_dict( self, - trajectories: List["Trajectory"], - other_trajectories: Optional[Dict[AgentID, "Trajectory"]] = None, + input_dict: Dict[str, TensorType], explore: bool = None, timestep: Optional[int] = None, **kwargs) -> \ @@ -181,10 +196,8 @@ class TorchPolicy(Policy): timestep = timestep if timestep is not None else self.global_timestep with torch.no_grad(): - # Create a view and pass that to Model as `input_dict`. - input_dict = self._lazy_tensor_dict( - get_trajectory_view( - self.model, trajectories, is_training=False)) + # Pass lazy (torch) tensor dict to Model as `input_dict`. + input_dict = self._lazy_tensor_dict(input_dict) # TODO: (sven) support RNNs w/ fast sampling. state_batches = [] seq_lens = None diff --git a/rllib/policy/torch_policy_template.py b/rllib/policy/torch_policy_template.py index 1133895ed..8e8b44e0b 100644 --- a/rllib/policy/torch_policy_template.py +++ b/rllib/policy/torch_policy_template.py @@ -8,6 +8,7 @@ from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 from ray.rllib.policy.policy import Policy, LEARNER_STATS_KEY from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.policy.torch_policy import TorchPolicy +from ray.rllib.policy.view_requirement import ViewRequirement from ray.rllib.utils import add_mixins from ray.rllib.utils.annotations import override, DeveloperAPI from ray.rllib.utils.framework import try_import_torch @@ -69,6 +70,8 @@ def build_torch_policy(name: str, apply_gradients_fn: Optional[Callable[ [Policy, "torch.optim.Optimizer"], None]] = None, mixins: Optional[List[type]] = None, + training_view_requirements_fn: Optional[Callable[ + [], Dict[str, ViewRequirement]]] = None, get_batch_divisibility_req: Optional[Callable[ [Policy], int]] = None ): @@ -162,6 +165,9 @@ def build_torch_policy(name: str, mixins (Optional[List[type]]): Optional list of any class mixins for the returned policy class. These mixins will be applied in order and will have higher precedence than the TorchPolicy class. + training_view_requirements_fn (Callable[[], + Dict[str, ViewRequirement]]): An optional callable to retrieve + additional train view requirements for this policy. get_batch_divisibility_req (Optional[Callable[[Policy], int]]): Optional callable that returns the divisibility requirement for sample batches. If None, will assume a value of 1. @@ -230,6 +236,13 @@ def build_torch_policy(name: str, if after_init: after_init(self, obs_space, action_space, config) + @override(TorchPolicy) + def training_view_requirements(self): + req = super().training_view_requirements() + if callable(training_view_requirements_fn): + req.update(training_view_requirements_fn(self)) + return req + @override(Policy) def postprocess_trajectory(self, sample_batch, @@ -306,6 +319,22 @@ def build_torch_policy(name: str, return convert_to_non_torch_type(stats_dict) def with_updates(**overrides): + """Allows creating a TorchPolicy cls based on settings of another one. + + Keyword Args: + **overrides: The settings (passed into `build_torch_policy`) that + should be different from the class that this method is called + on. + + Returns: + type: A new TorchPolicy sub-class. + + Examples: + >> MySpecialDQNPolicyClass = DQNTorchPolicy.with_updates( + .. name="MySpecialDQNPolicyClass", + .. loss_function=[some_new_loss_function], + .. ) + """ return build_torch_policy(**dict(original_kwargs, **overrides)) policy_cls.with_updates = staticmethod(with_updates) diff --git a/rllib/policy/trajectory_view.py b/rllib/policy/trajectory_view.py deleted file mode 100644 index 9a45b12a3..000000000 --- a/rllib/policy/trajectory_view.py +++ /dev/null @@ -1,94 +0,0 @@ -import numpy as np -from typing import Dict, Optional, List, TYPE_CHECKING - -from ray.rllib.utils.types import TensorType - -if TYPE_CHECKING: - from ray.rllib.models import ModelV2 - - -class ViewRequirement: - """Single view requirement (for one column in a ModelV2 input_dict). - - Note: This is an experimental class. - - ModelV2 returns a Dict[str, ViewRequirement] upon calling - `ModelV2.get_view_requirements()`, where the str key represents the column - name (C) under which the view is available in the `input_dict` and - ViewRequirement specifies the actual underlying column names (in the - original data buffer), timesteps, and other options to build the view - for N. - - Examples: - >>> # The default ViewRequirement for a Model is: - >>> req = [ModelV2].get_view_requirements(is_training=False) - >>> print(req) - {"obs": ViewRequirement(timesteps=0)} - """ - - def __init__(self, - data_col: Optional[str] = None, - timesteps: int = 0, - fill_mode: str = "zeros", - repeat_mode: str = "all"): - """Initializes a ViewRequirement object. - - Args: - data_col (): The data column name from the SampleBatch (str key). - If None, use the dict key under which this ViewRequirement - resides. - timesteps (Union[List[int], int]): List of relative (or absolute - timesteps) to be present in the input_dict. - fill_mode (str): The fill mode in case t<0 or t>H. - One of "zeros", "tile". - repeat_mode (str): The repeat-mode (one of "all" or "only_first"). - E.g. for training, we only want the first internal state - timestep (the NN will calculate all others again anyways). - """ - self.data_col = data_col - self.timesteps = timesteps - - # Switch on absolute timestep mode. Default: False. - # TODO: (sven) - # "absolute_timesteps", - - self.fill_mode = fill_mode - self.repeat_mode = repeat_mode - - # Provide all data as time major (default: False). - # TODO: (sven) - # "time_major", - - -def get_trajectory_view(model: "ModelV2", - trajectories: List["Trajectory"], - is_training: bool = False) -> Dict[str, TensorType]: - """Returns an input_dict for a Model's forward pass given some data. - - Args: - model (ModelV2): The ModelV2 object for which to generate the view - (input_dict) from `data`. - trajectories (List[Trajectory]): The data from which to generate - an input_dict. - is_training (bool): Whether the view should be generated for training - purposes or inference (default). - - Returns: - Dict[str, TensorType]: The input_dict to be passed into the ModelV2 - for inference/training. - """ - # Get ModelV2's view requirements. - view_reqs = model.get_view_requirements(is_training=is_training) - # Construct the view dict. - view = {} - for view_col, view_req in view_reqs.items(): - # Create the batch of data from the different buffers in `data`. - # TODO: (sven): Here, we actually do create a copy of the data (from a - # list). The only way to avoid this entirely would be to keep a - # single(!) np buffer per column across all currently ongoing - # agents + episodes (which seems very hard to realize). - view[view_col] = np.array([ - t.buffers[view_req.data_col][t.cursor + view_req.timesteps] - for t in trajectories - ]) - return view diff --git a/rllib/policy/view_requirement.py b/rllib/policy/view_requirement.py new file mode 100644 index 000000000..9ecaf7b70 --- /dev/null +++ b/rllib/policy/view_requirement.py @@ -0,0 +1,53 @@ +import gym +from typing import Optional + +from ray.rllib.utils.framework import try_import_torch + +torch, _ = try_import_torch() + + +class ViewRequirement: + """Single view requirement (for one column in an SampleBatch/input_dict). + + Note: This is an experimental class used only if + `_use_trajectory_view_api` in the config is set to True. + + Policies and ModelV2s return a Dict[str, ViewRequirement] upon calling + their `[train|inference]_view_requirements()` methods, where the str key + represents the column name (C) under which the view is available in the + input_dict/SampleBatch and ViewRequirement specifies the actual underlying + column names (in the original data buffer), timestep shifts, and other + options to build the view. + + Examples: + >>> # The default ViewRequirement for a Model is: + >>> req = [ModelV2].inference_view_requirements() + >>> print(req) + {"obs": ViewRequirement(shift=0)} + """ + + def __init__(self, + data_col: Optional[str] = None, + space: gym.Space = None, + shift: int = 0): + """Initializes a ViewRequirement object. + + Args: + data_col (): The data column name from the SampleBatch (str key). + If None, use the dict key under which this ViewRequirement + resides. + space (gym.Space): The gym Space used in case we need to pad data + in inaccessible areas of the trajectory (t<0 or t>H). + Default: Simple box space, e.g. rewards. + shift (Union[List[int], int]): Single shift value of list of + shift values to use relative to the underlying `data_col`. + Example: For a view column "prev_actions", you can set + `data_col="actions"` and `shift=-1`. + Example: For a view column "obs" in an Atari framestacking + fashion, you can set `data_col="obs"` and + `shift=[-3, -2, -1, 0]`. + """ + self.data_col = data_col + self.space = space or gym.spaces.Box( + float("-inf"), float("inf"), shape=()) + self.shift = shift