wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-24 13:20:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

[RLlib] Trajectory view API - 02 actual API scaffold (#9753)

Browse Source
This commit is contained in:
Sven Mika
2020-08-06 10:54:20 +02:00
committed by GitHub
parent 19d785b947
commit 57690a3a9f
15 changed files with 465 additions and 406 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rllib/BUILD
+7
View File
@@ -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/
rllib/agents/ppo/ppo_torch_policy.py
+13 -1
View File
@@ -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,
)
rllib/evaluation/sample_batch_builder.py
+4
View File
@@ -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.
rllib/evaluation/sample_collector.py
+180
View File
@@ -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
rllib/evaluation/trajectory.py
-267
View File
@@ -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
rllib/models/modelv2.py
+10 -21
View File
@@ -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:
rllib/models/torch/recurrent_net.py
+15
View File
@@ -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
rllib/policy/policy.py
+28 -8
View File
@@ -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,
rllib/policy/sample_batch.py
+4 -3
View File
@@ -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
rllib/policy/tests/test_trajectory_view_api.py
+81
View File
@@ -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__]))
rllib/policy/tf_policy_template.py
+16
View File
@@ -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():
rllib/policy/torch_policy.py
+25 -12
View File
@@ -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
rllib/policy/torch_policy_template.py
+29
View File
@@ -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)
rllib/policy/trajectory_view.py
-94
View File
@@ -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
rllib/policy/view_requirement.py
+53
View File
@@ -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