Sample Collections and Trajectory Views
Contents

Sample Collections and Trajectory Views#
The SampleCollector Class is Used to Store and Retrieve Temporary Data#
RLlib’s RolloutWorkers,
when running against a live environment, use the SamplerInput
class to interact
with that environment and produce batches of experiences.
The two implemented sub-classes of SamplerInput
are SyncSampler
and AsyncSampler
(residing under the RolloutWorker.sampler
property).
In case the “_use_trajectory_view_api” top-level config key is set to True
(by default since version >=1.1.0), every such sampler object will use the
SampleCollector
API to store and retrieve temporary environment-, model-, and other data
during rollouts (see figure below).
Sample collection process implemented by RLlib:
The Policy’s model tells the Sampler and its SampleCollector object, which data to store and
how to present it back to the dependent methods (e.g. Model.compute_actions()
).
This is done using a dict that maps strings (column names) to ViewRequirement
objects (details see below).
The exact behavior for a single such rollout and the number of environment transitions therein
are determined by the following AlgorithmConfig.rollout(..)
args:
- batch_mode [truncate_episodes|complete_episodes]:
- truncated_episodes (default value):
Rollouts are performed over exactly
rollout_fragment_length
(see below) number of steps. Thereby, steps are counted as either environment steps or as individual agent steps (seecount_steps_as
below). It does not matter, whether one or more episodes end within this rollout or whether the rollout starts in the middle of an already ongoing episode.- complete_episodes:
Each rollout always only contains full episodes (from beginning to terminal), never any episode fragments. The number of episodes in the rollout is 1 or larger. The
rollout_fragment_length
setting defines the minimum number of timesteps that will be covered in the rollout. For example, ifrollout_fragment_length=100
and your episodes are always 98 timesteps long, then rollouts will happen over two complete episodes and always be 196 timesteps long: 98 < 100 -> too short, keep rollout going; 98+98 >= 100 -> good, stop rollout after 2 episodes (196 timesteps). Note that you have to be careful when choosingcomplete_episodes
as batch_mode: If your environment does not terminate easily, this setting could lead to enormous batch sizes.
- rollout_fragment_length [int]:
The exact number of environment- or agent steps to be performed per rollout, if the
batch_mode
setting (see above) is “truncate_episodes”. Ifbatch_mode
is “complete_episodes”,rollout_fragment_length
is ignored, The unit to count fragments in is set viamultiagent.count_steps_by=[env_steps|agent_steps]
(within themultiagent
config dict).
Above: The two supported batch modes in RLlib. For “truncated_episodes”,
batches can a) span over more than one episode, b) end in the middle of an episode, and
c) start in the middle of an episode. Also, Policy.postprocess_trajectory()
is always
called at the end of a rollout-fragment (red lines on right side) as well as at the end
of each episode (arrow heads). This way, RLlib makes sure that the
Policy.postprocess_trajectory()
method never sees data from more than one episode.
… as well as AlgorithmConfig.multi_agent(count_steps_by=..)
:
- count_steps_by [env_steps|agent_steps]:
Within the Algorithm’s
multiagent
config dict, you can set the unit, by which RLlib will count a) rollout fragment lengths as well as b) the size of the final train_batch (see below). The two supported values are:- env_steps (default):
Each call to
[Env].step()
is counted as one. It does not matter, how many individual agents are stepping simultaneously in this very call (not all existing agents in the environment may step at the same time).- agent_steps:
In a multi-agent environment, count each individual agent’s step as one. For example, if N agents are in an environment and all these N agents always step at the same time, a single env step corresponds to N agent steps. Note that in the single-agent case,
env_steps
andagent_steps
are the same thing.
To trigger a single rollout, RLlib calls RolloutWorker.sample()
, which returns
a SampleBatch or MultiAgentBatch object representing all the data collected during that
rollout. These batches are then usually further concatenated (from the num_workers
parallelized RolloutWorkers) to form a final train batch. The size of that train batch is determined
by the train_batch_size
config parameter. Train batches are usually sent to the Policy’s
learn_on_batch
method, which handles loss- and gradient calculations, and optimizer stepping.
RLlib’s default SampleCollector
class is the SimpleListCollector
, which appends single timestep data (e.g. actions)
to lists, then builds SampleBatches from these and sends them to the downstream processing functions.
It thereby tries to avoid collecting duplicate data separately (OBS and NEXT_OBS use the same underlying list).
If you want to implement your own collection logic and data structures, you can sub-class SampleCollector
and specify that new class under the Algorithm’s “sample_collector” config key.
Let’s now look at how the Policy’s Model lets the RolloutWorker and its SampleCollector
know, what data in the ongoing episode/trajectory to use for the different required method calls
during rollouts. These method calls in particular are:
Policy.compute_actions_from_input_dict()
to compute actions to be taken in an episode.
Policy.postprocess_trajectory()
, which is called after an episode ends or a rollout hit its
rollout_fragment_length
limit (in batch_mode=truncated_episodes
), and Policy.learn_on_batch()
,
which is called with a “train_batch” to improve the policy.
Trajectory View API#
The trajectory view API allows custom models to define what parts of the trajectory they require in order to execute the forward pass. For example, in the simplest case, a model might only look at the latest observation. However, an RNN- or attention based model could look at previous states emitted by the model, concatenate previously seen rewards with the current observation, or require the entire range of the n most recent observations.
The trajectory view API lets models define these requirements and lets RLlib gather the required data for the forward pass in an efficient way.
Since the following methods all call into the model class, they are all indirectly using the trajectory view API. It is important to note that the API is only accessible to the user via the model classes (see below on how to setup trajectory view requirements for a custom model).
In particular, the methods receiving inputs that depend on a Model’s trajectory view rules are:
Policy.compute_actions_from_input_dict()
Policy.postprocess_trajectory()
andPolicy.learn_on_batch()
(and consecutively: the Policy’s loss function).
The input data to these methods can stem from either the environment (observations, rewards, and env infos), the model itself (previously computed actions, internal state outputs, action-probs, etc..) or the Sampler (e.g. agent index, env ID, episode ID, timestep, etc..). All data has an associated time axis, which is 0-based, meaning that the first action taken, the first reward received in an episode, and the first observation (directly after a reset) all have t=0.
The idea is to allow more flexibility and standardization in how a model defines required “views” on the ongoing trajectory (during action computations/inference), past episodes (training on a batch), or even trajectories of other agents in the same episode, some of which may even use a different policy.
Such a “view requirements” formalism is helpful when having to support more complex model setups like RNNs, attention nets, observation image framestacking (e.g. for Atari), and building multi-agent communication channels.
The way to define a set of rules used for making the Model see certain
data is through a “view requirements dict”, residing in the Policy.model.view_requirements
property.
View requirements dicts map strings (column names), such as “obs” or “actions” to
a ViewRequirement
object, which defines the exact conditions by which this column
should be populated with data.
View Requirement Dictionaries#
View requirements are stored within the view_requirements
property of the ModelV2
class.
You can acccess it like this:
my_simple_model = ModelV2(...)
print(my_simple_model.view_requirements)
>>>{"obs": ViewRequirement(shift=0, space=[observation space])}
my_lstm_model = LSTMModel(...)
print(my_lstm_model.view_requirements)
>>>{
>>> "obs": ViewRequirement(shift=0, space=[observation space]),
>>> "prev_actions": ViewRequirement(shift=-1, data_col="actions", space=[action space]),
>>> "prev_rewards": ViewRequirement(shift=-1, data_col="rewards"),
>>>}
The view_requirements
property holds a dictionary mapping
string keys (e.g. “actions”, “rewards”, “next_obs”, etc..)
to a ViewRequirement
object. This ViewRequirement
object determines what exact data to
provide under the given key in case a SampleBatch or a single-timestep (action computing) “input dict”
needs to be build and fed into one of the above ModelV2- or Policy methods.
Above: An example ViewRequirements
dict that causes the current observation
and the previous action to be available in each compute_action call, as
well as for the Policy’s postprocess_trajectory()
function (and train batch).
A similar setup is often used by LSTM/RNN-based models.
The ViewRequirement class#
Here is a description of the constructor-settable properties of a ViewRequirement object and what each of these properties controls.
- data_col:
An optional string key referencing the underlying data to use to create the view. If not provided, assumes that there is data under the dict-key under which this ViewRequirement resides.
Examples:
ModelV2.view_requirements = {"rewards": ViewRequirements(shift=0)}
# -> implies that the underlying data to use are the collected rewards
# from the environment.
ModelV2.view_requirements = {"prev_rewards": ViewRequirements(data_col="rewards", shift=-1)}
# -> means that the actual data used to create the "prev_rewards" column
# is the "rewards" data from the environment (shifted by 1 timestep).
- space:
An optional gym.Space used as a hint for the SampleCollector to know, how to fill timesteps before the episode actually started (e.g. if shift=-2, we need dummy data at timesteps -2 and -1).
- shift [int]:
An int, a list of ints, or a range string (e.g. “-50:-1”) to indicate which time offsets or ranges of the underlying data to use for the view.
Examples:
shift=0 # -> Use the data under ``data_col`` as is.
shift=1 # -> Use the data under ``data_col``, but shifted by +1 timestep
# (used by e.g. next_obs views).
shift=-1 # -> Use the data under ``data_col``, but shifted by -1 timestep
# (used by e.g. prev_actions views).
shift=[-2, -1] # -> Use the data under ``data_col``, but always provide 2 values
# at each timestep (the two previous ones).
# Could be used e.g. to feed the last two actions or rewards into an LSTM.
shift="-50:-1" # -> Use the data under ``data_col``, but always provide a range of
# the last 50 timesteps (used by our attention nets).
- used_for_training [bool]:
True by default. If False, the column will not be available inside the train batch (arriving in the Policy’s loss function). RLlib will automatically switch this to False for a given column, if it detects during Policy initialization that that column is not accessed inside the loss function (see below).
How does RLlib determine, which Views are required?#
When initializing a Policy, it automatically determines how to later build batches
for postprocessing, loss function calls, and action computations, based on
the Model’s view_requirements
dict. It does so by sending generic dummy batches
through its compute_actions_from_input_dict
, postprocess_trajectory
, and loss functions
and then checks, which fields in these dummy batches get accessed, overwritten, deleted or added.
Based on these test passes, the Policy then throws out those ViewRequirements from an initial
very broad list, that it deems unnecessary. This procedure saves a lot of data copying
during later rollouts, batch transfers (via ray) and loss calculations and makes things like
manually deleting columns from a SampleBatch (e.g. PPO used to delete the “next_obs” column
inside the postprocessing function) unnecessary.
Note that the “rewards” and “dones” columns are never discarded and thus should always
arrive in your loss function’s SampleBatch (train_batch
arg).
Setting ViewRequirements manually in your Model#
If you need to specify special view requirements for your model, you can add
columns to the Model’s view_requirements
dict in the
Model’s constructor.
For example, our auto-LSTM wrapper classes (tf and torch) have these additional lines in their constructors (torch version shown here):
# Add prev-a/r to this model's view, if required.
if model_config["lstm_use_prev_action"]:
self.view_requirements[SampleBatch.PREV_ACTIONS] = ViewRequirement(
SampleBatch.ACTIONS, space=self.action_space, shift=-1
)
if model_config["lstm_use_prev_reward"]:
self.view_requirements[SampleBatch.PREV_REWARDS] = ViewRequirement(
SampleBatch.REWARDS, shift=-1
)
This makes sure that, if the users requires this via the model config, previous rewards
and/or previous actions are added properly to the compute_actions
input-dicts and SampleBatches
used for postprocessing and training.
Another example are our attention nets, which make sure the last n (memory) model outputs are always fed back into the model on the next time step (tf version shown here).
# Setup trajectory views (`memory-inference` x past memory outs).
for i in range(self.num_transformer_units):
space = Box(-1.0, 1.0, shape=(self.attention_dim,))
self.view_requirements["state_in_{}".format(i)] = ViewRequirement(
"state_out_{}".format(i),
shift="-{}:-1".format(self.memory_inference),
# Repeat the incoming state every max-seq-len times.
batch_repeat_value=self.max_seq_len,
space=space,
)
self.view_requirements["state_out_{}".format(i)] = ViewRequirement(
space=space, used_for_training=False
)
Setting ViewRequirements manually after Policy construction#
Here is a simple example, of how you can modify and add to the ViewRequirements dict even after policy (or RolloutWorker) creation. However, note that it’s better to make these modifications to your batches in your postprocessing function:
# Modify view_requirements in the Policy object.
action_space = Discrete(2)
rollout_worker = RolloutWorker(
env_creator=lambda _: gym.make("CartPole-v1"),
policy_config=ppo.DEFAULT_CONFIG,
policy_spec=ppo.PPOTorchPolicy,
)
policy = rollout_worker.policy_map["default_policy"]
# Add the next action to the view reqs of the policy.
# This should be visible then in postprocessing and train batches.
policy.view_requirements["next_actions"] = ViewRequirement(
SampleBatch.ACTIONS, shift=1, space=action_space)
# Check, whether a sampled batch has the requested `next_actions` view.
batch = rollout_worker.sample()
self.assertTrue("next_actions" in batch.data)
# Doing the same in a custom postprocessing callback function:
class MyCallback(DefaultCallbacks):
# ...
@override(DefaultCallbacks)
def on_postprocess_trajectory(self, worker, episode, agent_id, policy_id,
policies, postprocessed_batch, original_batches,
**kwargs):
postprocessed_batch["next_actions"] = np.concatenate(
[postprocessed_batch["actions"][1:],
np.zeros_like([policies[policy_id].action_space.sample()])])
The above two examples add a “next_action” view to the postprocessed SampleBatche needed
used by the Policy for training. It will not feed the “next_action”
to the Model’s compute_action
calls (it can’t b/c the next action is of course not known
at that point).