wehub/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
ray-project--ray/rllib/policy/sample_batch.py
T
wehub-resource-sync f1825c8ceb chore: import upstream snapshot with attribution
2026-07-13 13:17:40 +08:00

1844 lines
67 KiB
Python
Raw Permalink Blame History

import collections
import itertools
import sys
from functools import partial
from numbers import Number
from typing import Dict, Iterator, List, Optional, Set, Union
import numpy as np
import tree # pip install dm_tree
from ray._common.deprecation import Deprecated, deprecation_warning
from ray.rllib.core.columns import Columns
from ray.rllib.utils.annotations import DeveloperAPI, ExperimentalAPI, PublicAPI
from ray.rllib.utils.compression import is_compressed, pack, unpack
from ray.rllib.utils.framework import try_import_tf, try_import_torch
from ray.rllib.utils.torch_utils import convert_to_torch_tensor
from ray.rllib.utils.typing import (
ModuleID,
PolicyID,
SampleBatchType,
TensorType,
ViewRequirementsDict,
)
from ray.util import log_once
tf1, tf, tfv = try_import_tf()
torch, _ = try_import_torch()
# Default policy id for single agent environments
DEFAULT_POLICY_ID = "default_policy"
@DeveloperAPI
def attempt_count_timesteps(tensor_dict: dict):
"""Attempt to count timesteps based on dimensions of individual elements.
Returns the first successfully counted number of timesteps.
We do not attempt to count on INFOS or any state_in_* and state_out_* keys. The
number of timesteps we count in cases where we are unable to count is zero.
Args:
tensor_dict: A SampleBatch or another dict.
Returns:
count: The inferred number of timesteps >= 0.
"""
# Try to infer the "length" of the SampleBatch by finding the first
# value that is actually a ndarray/tensor.
# Skip manual counting routine if we can directly infer count from sequence lengths
seq_lens = tensor_dict.get(SampleBatch.SEQ_LENS)
if (
seq_lens is not None
and not (tf and tf.is_tensor(seq_lens) and not hasattr(seq_lens, "numpy"))
and len(seq_lens) > 0
):
if torch and torch.is_tensor(seq_lens):
return int(seq_lens.sum().item())
else:
return int(sum(seq_lens))
for k, v in tensor_dict.items():
if k == SampleBatch.SEQ_LENS:
continue
assert isinstance(k, str), tensor_dict
if (
k == SampleBatch.INFOS
or k.startswith("state_in_")
or k.startswith("state_out_")
):
# Don't attempt to count on infos since we make no assumptions
# about its content
# Don't attempt to count on state since nesting can potentially mess
# things up
continue
# If this is a nested dict (for example a nested observation),
# try to flatten it, assert that all elements have the same length (batch
# dimension)
v_list = tree.flatten(v) if isinstance(v, (dict, tuple)) else [v]
# TODO: Drop support for lists and Numbers as values.
# If v_list contains lists or Numbers, convert them to arrays, too.
v_list = [
np.array(_v) if isinstance(_v, (Number, list)) else _v for _v in v_list
]
try:
# Add one of the elements' length, since they are all the same
_len = len(v_list[0])
if _len:
return _len
except Exception:
pass
# Return zero if we are unable to count
return 0
@PublicAPI
class SampleBatch(dict):
"""Wrapper around a dictionary with string keys and array-like values.
For example, {"obs": [1, 2, 3], "reward": [0, -1, 1]} is a batch of three
samples, each with an "obs" and "reward" attribute.
"""
# On rows in SampleBatch:
# Each comment signifies how values relate to each other within a given row.
# A row generally signifies one timestep. Most importantly, at t=0, SampleBatch.OBS
# will usually be the reset-observation, while SampleBatch.ACTIONS will be the
# action based on the reset-observation and so on. This scheme is derived from
# RLlib's sampling logic.
# The following fields have all been moved to `Columns` and are only left here
# for backward compatibility.
OBS = Columns.OBS
ACTIONS = Columns.ACTIONS
REWARDS = Columns.REWARDS
TERMINATEDS = Columns.TERMINATEDS
TRUNCATEDS = Columns.TRUNCATEDS
INFOS = Columns.INFOS
SEQ_LENS = Columns.SEQ_LENS
T = Columns.T
ACTION_DIST_INPUTS = Columns.ACTION_DIST_INPUTS
ACTION_PROB = Columns.ACTION_PROB
ACTION_LOGP = Columns.ACTION_LOGP
VF_PREDS = Columns.VF_PREDS
VALUES_BOOTSTRAPPED = Columns.VALUES_BOOTSTRAPPED
EPS_ID = Columns.EPS_ID
NEXT_OBS = Columns.NEXT_OBS
# Action distribution object.
ACTION_DIST = "action_dist"
# Action chosen before SampleBatch.ACTIONS.
PREV_ACTIONS = "prev_actions"
# Reward received before SampleBatch.REWARDS.
PREV_REWARDS = "prev_rewards"
ENV_ID = "env_id" # An env ID (e.g. the index for a vectorized sub-env).
AGENT_INDEX = "agent_index" # Uniquely identifies an agent within an episode.
# Uniquely identifies a sample batch. This is important to distinguish RNN
# sequences from the same episode when multiple sample batches are
# concatenated (fusing sequences across batches can be unsafe).
UNROLL_ID = "unroll_id"
# RE 3
# This is only computed and used when RE3 exploration strategy is enabled.
OBS_EMBEDS = "obs_embeds"
# Decision Transformer
RETURNS_TO_GO = "returns_to_go"
ATTENTION_MASKS = "attention_masks"
# Do not set this key directly. Instead, the values under this key are
# auto-computed via the values of the TERMINATEDS and TRUNCATEDS keys.
DONES = "dones"
# Use SampleBatch.OBS instead.
CUR_OBS = "obs"
@PublicAPI
def __init__(self, *args, **kwargs):
"""Constructs a sample batch (same params as dict constructor).
Note: All args and those kwargs not listed below will be passed
as-is to the parent dict constructor.
Args:
_time_major: Whether data in this sample batch
is time-major. This is False by default and only relevant
if the data contains sequences.
_max_seq_len: The max sequence chunk length
if the data contains sequences.
_zero_padded: Whether the data in this batch
contains sequences AND these sequences are right-zero-padded
according to the `_max_seq_len` setting.
_is_training: Whether this batch is used for
training. If False, batch may be used for e.g. action
computations (inference).
"""
if SampleBatch.DONES in kwargs:
raise KeyError(
"SampleBatch cannot be constructed anymore with a `DONES` key! "
"Instead, set the new TERMINATEDS and TRUNCATEDS keys. The values under"
" DONES will then be automatically computed using terminated|truncated."
)
# Possible seq_lens (TxB or BxT) setup.
self.time_major = kwargs.pop("_time_major", None)
# Maximum seq len value.
self.max_seq_len = kwargs.pop("_max_seq_len", None)
# Is alredy right-zero-padded?
self.zero_padded = kwargs.pop("_zero_padded", False)
# Whether this batch is used for training (vs inference).
self._is_training = kwargs.pop("_is_training", None)
# Weighted average number of grad updates that have been performed on the
# policy/ies that were used to collect this batch.
# E.g.: Two rollout workers collect samples of 50ts each
# (rollout_fragment_length=50). One of them has a policy that has undergone
# 2 updates thus far, the other worker uses a policy that has undergone 3
# updates thus far. The train batch size is 100, so we concatenate these 2
# batches to a new one that's 100ts long. This new 100ts batch will have its
# `num_gradient_updates` property set to 2.5 as it's the weighted average
# (both original batches contribute 50%).
self.num_grad_updates: Optional[float] = kwargs.pop("_num_grad_updates", None)
# Call super constructor. This will make the actual data accessible
# by column name (str) via e.g. self["some-col"].
dict.__init__(self, *args, **kwargs)
# Indicates whether, for this batch, sequence lengths should be slices by
# their index in the batch or by their index as a sequence.
# This is useful if a batch contains tensors of shape (B, T, ...), where each
# index of B indicates one sequence. In this case, when slicing the batch,
# we want one sequence to be slices out per index in B (
# `_slice_seq_lens_by_batch_index=True`. However, if the padded batch
# contains tensors of shape (B*T, ...), where each index of B*T indicates
# one timestep, we want one sequence to be sliced per T steps in B*T (
# `self._slice_seq_lens_in_B=False`).
# ._slice_seq_lens_in_B = True is only meant to be used for batches that we
# feed into Learner._update(), all other places in RLlib are not expected to
# need this.
self._slice_seq_lens_in_B = False
self.accessed_keys = set()
self.added_keys = set()
self.deleted_keys = set()
self.intercepted_values = {}
self.get_interceptor = None
# Clear out None seq-lens.
seq_lens_ = self.get(SampleBatch.SEQ_LENS)
if seq_lens_ is None or (isinstance(seq_lens_, list) and len(seq_lens_) == 0):
self.pop(SampleBatch.SEQ_LENS, None)
# Numpyfy seq_lens if list.
elif isinstance(seq_lens_, list):
self[SampleBatch.SEQ_LENS] = seq_lens_ = np.array(seq_lens_, dtype=np.int32)
elif (torch and torch.is_tensor(seq_lens_)) or (tf and tf.is_tensor(seq_lens_)):
self[SampleBatch.SEQ_LENS] = seq_lens_
if (
self.max_seq_len is None
and seq_lens_ is not None
and not (tf and tf.is_tensor(seq_lens_))
and len(seq_lens_) > 0
):
if torch and torch.is_tensor(seq_lens_):
self.max_seq_len = seq_lens_.max().item()
else:
self.max_seq_len = max(seq_lens_)
if self._is_training is None:
self._is_training = self.pop("is_training", False)
for k, v in self.items():
# TODO: Drop support for lists and Numbers as values.
# Convert lists of int|float into numpy arrays make sure all data
# has same length.
if isinstance(v, (Number, list)) and not k == SampleBatch.INFOS:
self[k] = np.array(v)
self.count = attempt_count_timesteps(self)
# A convenience map for slicing this batch into sub-batches along
# the time axis. This helps reduce repeated iterations through the
# batch's seq_lens array to find good slicing points. Built lazily
# when needed.
self._slice_map = []
@PublicAPI
def __len__(self) -> int:
"""Returns the amount of samples in the sample batch."""
return self.count
@PublicAPI
def agent_steps(self) -> int:
"""Returns the same as len(self) (number of steps in this batch).
To make this compatible with `MultiAgentBatch.agent_steps()`.
"""
return len(self)
@PublicAPI
def env_steps(self) -> int:
"""Returns the same as len(self) (number of steps in this batch).
To make this compatible with `MultiAgentBatch.env_steps()`.
"""
return len(self)
@DeveloperAPI
def enable_slicing_by_batch_id(self):
self._slice_seq_lens_in_B = True
@DeveloperAPI
def disable_slicing_by_batch_id(self):
self._slice_seq_lens_in_B = False
@ExperimentalAPI
def is_terminated_or_truncated(self) -> bool:
"""Returns True if `self` is either terminated or truncated at idx -1."""
return self[SampleBatch.TERMINATEDS][-1] or (
SampleBatch.TRUNCATEDS in self and self[SampleBatch.TRUNCATEDS][-1]
)
@ExperimentalAPI
def is_single_trajectory(self) -> bool:
"""Returns True if this SampleBatch only contains one trajectory.
This is determined by checking all timesteps (except for the last) for being
not terminated AND (if applicable) not truncated.
"""
return not any(self[SampleBatch.TERMINATEDS][:-1]) and (
SampleBatch.TRUNCATEDS not in self
or not any(self[SampleBatch.TRUNCATEDS][:-1])
)
@staticmethod
@PublicAPI
@Deprecated(new="concat_samples() from rllib.policy.sample_batch", error=True)
def concat_samples(samples):
pass
@PublicAPI
def concat(self, other: "SampleBatch") -> "SampleBatch":
"""Concatenates `other` to this one and returns a new SampleBatch.
Args:
other: The other SampleBatch object to concat to this one.
Returns:
The new SampleBatch, resulting from concating `other` to `self`.
.. testcode::
:skipif: True
import numpy as np
from ray.rllib.policy.sample_batch import SampleBatch
b1 = SampleBatch({"a": np.array([1, 2])})
b2 = SampleBatch({"a": np.array([3, 4, 5])})
print(b1.concat(b2))
.. testoutput::
{"a": np.array([1, 2, 3, 4, 5])}
"""
return concat_samples([self, other])
@PublicAPI
def copy(self, shallow: bool = False) -> "SampleBatch":
"""Creates a deep or shallow copy of this SampleBatch and returns it.
Args:
shallow: Whether the copying should be done shallowly.
Returns:
A deep or shallow copy of this SampleBatch object.
"""
copy_ = dict(self)
data = tree.map_structure(
lambda v: (
np.array(v, copy=not shallow) if isinstance(v, np.ndarray) else v
),
copy_,
)
copy_ = SampleBatch(
data,
_time_major=self.time_major,
_zero_padded=self.zero_padded,
_max_seq_len=self.max_seq_len,
_num_grad_updates=self.num_grad_updates,
)
copy_.set_get_interceptor(self.get_interceptor)
copy_.added_keys = self.added_keys
copy_.deleted_keys = self.deleted_keys
copy_.accessed_keys = self.accessed_keys
return copy_
@PublicAPI
def rows(self) -> Iterator[Dict[str, TensorType]]:
"""Returns an iterator over data rows, i.e. dicts with column values.
Note that if `seq_lens` is set in self, we set it to 1 in the rows.
Yields:
The column values of the row in this iteration.
.. testcode::
:skipif: True
from ray.rllib.policy.sample_batch import SampleBatch
batch = SampleBatch({
"a": [1, 2, 3],
"b": [4, 5, 6],
"seq_lens": [1, 2]
})
for row in batch.rows():
print(row)
.. testoutput::
{"a": 1, "b": 4, "seq_lens": 1}
{"a": 2, "b": 5, "seq_lens": 1}
{"a": 3, "b": 6, "seq_lens": 1}
"""
seq_lens = None if self.get(SampleBatch.SEQ_LENS, 1) is None else 1
self_as_dict = dict(self)
for i in range(self.count):
yield tree.map_structure_with_path(
lambda p, v, i=i: v[i] if p[0] != self.SEQ_LENS else seq_lens,
self_as_dict,
)
@PublicAPI
def columns(self, keys: List[str]) -> List[any]:
"""Returns a list of the batch-data in the specified columns.
Args:
keys: List of column names fo which to return the data.
Returns:
The list of data items ordered by the order of column
names in `keys`.
.. testcode::
:skipif: True
from ray.rllib.policy.sample_batch import SampleBatch
batch = SampleBatch({"a": [1], "b": [2], "c": [3]})
print(batch.columns(["a", "b"]))
.. testoutput::
[[1], [2]]
"""
# TODO: (sven) Make this work for nested data as well.
out = []
for k in keys:
out.append(self[k])
return out
@PublicAPI
def shuffle(self) -> "SampleBatch":
"""Shuffles the rows of this batch in-place.
Returns:
This very (now shuffled) SampleBatch.
Raises:
ValueError: If self[SampleBatch.SEQ_LENS] is defined.
.. testcode::
:skipif: True
from ray.rllib.policy.sample_batch import SampleBatch
batch = SampleBatch({"a": [1, 2, 3, 4]})
print(batch.shuffle())
.. testoutput::
{"a": [4, 1, 3, 2]}
"""
has_time_rank = self.get(SampleBatch.SEQ_LENS) is not None
# Shuffling the data when we have `seq_lens` defined is probably
# a bad idea!
if has_time_rank and not self.zero_padded:
raise ValueError(
"SampleBatch.shuffle not possible when your data has "
"`seq_lens` defined AND is not zero-padded yet!"
)
# Get a permutation over the single items once and use the same
# permutation for all the data (otherwise, data would become
# meaningless).
# - Shuffle by individual item.
if not has_time_rank:
permutation = np.random.permutation(self.count)
# - Shuffle along batch axis (leave axis=1/time-axis as-is).
else:
permutation = np.random.permutation(len(self[SampleBatch.SEQ_LENS]))
self_as_dict = dict(self)
infos = self_as_dict.pop(Columns.INFOS, None)
shuffled = tree.map_structure(lambda v: v[permutation], self_as_dict)
if infos is not None:
self_as_dict[Columns.INFOS] = [infos[i] for i in permutation]
self.update(shuffled)
# Flush cache such that intercepted values are recalculated after the
# shuffling.
self.intercepted_values = {}
return self
@PublicAPI
def split_by_episode(self, key: Optional[str] = None) -> List["SampleBatch"]:
"""Splits by `eps_id` column and returns list of new batches.
If `eps_id` is not present, splits by `dones` instead.
Args:
key: If specified, overwrite default and use key to split.
Returns:
List of batches, one per distinct episode.
Raises:
KeyError: If the `eps_id` AND `dones` columns are not present.
.. testcode::
:skipif: True
from ray.rllib.policy.sample_batch import SampleBatch
# "eps_id" is present
batch = SampleBatch(
{"a": [1, 2, 3], "eps_id": [0, 0, 1]})
print(batch.split_by_episode())
# "eps_id" not present, split by "dones" instead
batch = SampleBatch(
{"a": [1, 2, 3, 4, 5], "dones": [0, 0, 1, 0, 1]})
print(batch.split_by_episode())
# The last episode is appended even if it does not end with done
batch = SampleBatch(
{"a": [1, 2, 3, 4, 5], "dones": [0, 0, 1, 0, 0]})
print(batch.split_by_episode())
batch = SampleBatch(
{"a": [1, 2, 3, 4, 5], "dones": [0, 0, 0, 0, 0]})
print(batch.split_by_episode())
.. testoutput::
[{"a": [1, 2], "eps_id": [0, 0]}, {"a": [3], "eps_id": [1]}]
[{"a": [1, 2, 3], "dones": [0, 0, 1]}, {"a": [4, 5], "dones": [0, 1]}]
[{"a": [1, 2, 3], "dones": [0, 0, 1]}, {"a": [4, 5], "dones": [0, 0]}]
[{"a": [1, 2, 3, 4, 5], "dones": [0, 0, 0, 0, 0]}]
"""
assert key is None or key in [SampleBatch.EPS_ID, SampleBatch.DONES], (
f"`SampleBatch.split_by_episode(key={key})` invalid! "
f"Must be [None|'dones'|'eps_id']."
)
def slice_by_eps_id():
slices = []
# Produce a new slice whenever we find a new episode ID.
cur_eps_id = self[SampleBatch.EPS_ID][0]
offset = 0
for i in range(self.count):
next_eps_id = self[SampleBatch.EPS_ID][i]
if next_eps_id != cur_eps_id:
slices.append(self[offset:i])
offset = i
cur_eps_id = next_eps_id
# Add final slice.
slices.append(self[offset : self.count])
return slices
def slice_by_terminateds_or_truncateds():
slices = []
offset = 0
for i in range(self.count):
if self[SampleBatch.TERMINATEDS][i] or (
SampleBatch.TRUNCATEDS in self and self[SampleBatch.TRUNCATEDS][i]
):
# Since self[i] is the last timestep of the episode,
# append it to the batch, then set offset to the start
# of the next batch
slices.append(self[offset : i + 1])
offset = i + 1
# Add final slice.
if offset != self.count:
slices.append(self[offset:])
return slices
key_to_method = {
SampleBatch.EPS_ID: slice_by_eps_id,
SampleBatch.DONES: slice_by_terminateds_or_truncateds,
}
# If key not specified, default to this order.
key_resolve_order = [SampleBatch.EPS_ID, SampleBatch.DONES]
slices = None
if key is not None:
# If key specified, directly use it.
if key == SampleBatch.EPS_ID and key not in self:
raise KeyError(f"{self} does not have key `{key}`!")
slices = key_to_method[key]()
else:
# If key not specified, go in order.
for key in key_resolve_order:
if key == SampleBatch.DONES or key in self:
slices = key_to_method[key]()
break
if slices is None:
raise KeyError(f"{self} does not have keys {key_resolve_order}!")
assert (
sum(s.count for s in slices) == self.count
), f"Calling split_by_episode on {self} returns {slices}"
f"which should in total have {self.count} timesteps!"
return slices
def slice(
self, start: int, end: int, state_start=None, state_end=None
) -> "SampleBatch":
"""Returns a slice of the row data of this batch (w/o copying).
Args:
start: Starting index. If < 0, will left-zero-pad.
end: Ending index.
Returns:
A new SampleBatch, which has a slice of this batch's data.
"""
if (
self.get(SampleBatch.SEQ_LENS) is not None
and len(self[SampleBatch.SEQ_LENS]) > 0
):
if start < 0:
data = {
k: np.concatenate(
[
np.zeros(shape=(-start,) + v.shape[1:], dtype=v.dtype),
v[0:end],
]
)
for k, v in self.items()
if k != SampleBatch.SEQ_LENS and not k.startswith("state_in_")
}
else:
data = {
k: tree.map_structure(lambda s: s[start:end], v)
for k, v in self.items()
if k != SampleBatch.SEQ_LENS and not k.startswith("state_in_")
}
if state_start is not None:
assert state_end is not None
state_idx = 0
state_key = "state_in_{}".format(state_idx)
while state_key in self:
data[state_key] = self[state_key][state_start:state_end]
state_idx += 1
state_key = "state_in_{}".format(state_idx)
seq_lens = list(self[SampleBatch.SEQ_LENS][state_start:state_end])
# Adjust seq_lens if necessary.
data_len = len(data[next(iter(data))])
if sum(seq_lens) != data_len:
assert sum(seq_lens) > data_len
seq_lens[-1] = data_len - sum(seq_lens[:-1])
else:
# Fix state_in_x data.
count = 0
state_start = None
seq_lens = None
for i, seq_len in enumerate(self[SampleBatch.SEQ_LENS]):
count += seq_len
if count >= end:
state_idx = 0
state_key = "state_in_{}".format(state_idx)
if state_start is None:
state_start = i
while state_key in self:
data[state_key] = self[state_key][state_start : i + 1]
state_idx += 1
state_key = "state_in_{}".format(state_idx)
seq_lens = list(self[SampleBatch.SEQ_LENS][state_start:i]) + [
seq_len - (count - end)
]
if start < 0:
seq_lens[0] += -start
diff = sum(seq_lens) - (end - start)
if diff > 0:
seq_lens[0] -= diff
assert sum(seq_lens) == (end - start)
break
elif state_start is None and count > start:
state_start = i
return SampleBatch(
data,
seq_lens=seq_lens,
_is_training=self.is_training,
_time_major=self.time_major,
_num_grad_updates=self.num_grad_updates,
)
else:
return SampleBatch(
tree.map_structure(lambda value: value[start:end], self),
_is_training=self.is_training,
_time_major=self.time_major,
_num_grad_updates=self.num_grad_updates,
)
def _batch_slice(self, slice_: slice) -> "SampleBatch":
"""Helper method to handle SampleBatch slicing using a slice object.
The returned SampleBatch uses the same underlying data object as
`self`, so changing the slice will also change `self`.
Note that only zero or positive bounds are allowed for both start
and stop values. The slice step must be 1 (or None, which is the
same).
Args:
slice_: The python slice object to slice by.
Returns:
A new SampleBatch, however "linking" into the same data
(sliced) as self.
"""
start = slice_.start or 0
stop = slice_.stop or len(self[SampleBatch.SEQ_LENS])
# If stop goes beyond the length of this batch -> Make it go till the
# end only (including last item).
# Analogous to `l = [0, 1, 2]; l[:100] -> [0, 1, 2];`.
if stop > len(self):
stop = len(self)
assert start >= 0 and stop >= 0 and slice_.step in [1, None]
# Exclude INFOs from regular array slicing as the data under this column might
# be a list (not good for `tree.map_structure` call).
# Furthermore, slicing does not work when the data in the column is
# singular (not a list or array).
infos = self.pop(SampleBatch.INFOS, None)
data = tree.map_structure(lambda value: value[start:stop], self)
if infos is not None:
# Slice infos according to SEQ_LENS.
info_slice_start = int(sum(self[SampleBatch.SEQ_LENS][:start]))
info_slice_stop = int(sum(self[SampleBatch.SEQ_LENS][start:stop]))
data[SampleBatch.INFOS] = infos[info_slice_start:info_slice_stop]
# Put infos back into `self`.
self[Columns.INFOS] = infos
return SampleBatch(
data,
_is_training=self.is_training,
_time_major=self.time_major,
_num_grad_updates=self.num_grad_updates,
)
@PublicAPI
def timeslices(
self,
size: Optional[int] = None,
num_slices: Optional[int] = None,
k: Optional[int] = None,
) -> List["SampleBatch"]:
"""Returns SampleBatches, each one representing a k-slice of this one.
Will start from timestep 0 and produce slices of size=k.
Args:
size: The size (in timesteps) of each returned SampleBatch.
num_slices: The number of slices to produce.
k: Deprecated: Use size or num_slices instead. The size
(in timesteps) of each returned SampleBatch.
Returns:
The list of `num_slices` (new) SampleBatches or n (new)
SampleBatches each one of size `size`.
"""
if size is None and num_slices is None:
deprecation_warning("k", "size or num_slices")
assert k is not None
size = k
if size is None:
assert isinstance(num_slices, int)
slices = []
left = len(self)
start = 0
while left:
len_ = left // (num_slices - len(slices))
stop = start + len_
slices.append(self[start:stop])
left -= len_
start = stop
return slices
else:
assert isinstance(size, int)
slices = []
left = len(self)
start = 0
while left:
stop = start + size
slices.append(self[start:stop])
left -= size
start = stop
return slices
@Deprecated(new="SampleBatch.right_zero_pad", error=True)
def zero_pad(self, max_seq_len, exclude_states=True):
pass
def right_zero_pad(self, max_seq_len: int, exclude_states: bool = True):
"""Right (adding zeros at end) zero-pads this SampleBatch in-place.
This will set the `self.zero_padded` flag to True and
`self.max_seq_len` to the given `max_seq_len` value.
Args:
max_seq_len: The max (total) length to zero pad to.
exclude_states: If False, also right-zero-pad all
`state_in_x` data. If True, leave `state_in_x` keys
as-is.
Returns:
This very (now right-zero-padded) SampleBatch.
Raises:
ValueError: If self[SampleBatch.SEQ_LENS] is None (not defined).
.. testcode::
:skipif: True
from ray.rllib.policy.sample_batch import SampleBatch
batch = SampleBatch(
{"a": [1, 2, 3], "seq_lens": [1, 2]})
print(batch.right_zero_pad(max_seq_len=4))
batch = SampleBatch({"a": [1, 2, 3],
"state_in_0": [1.0, 3.0],
"seq_lens": [1, 2]})
print(batch.right_zero_pad(max_seq_len=5))
.. testoutput::
{"a": [1, 0, 0, 0, 2, 3, 0, 0], "seq_lens": [1, 2]}
{"a": [1, 0, 0, 0, 0, 2, 3, 0, 0, 0],
"state_in_0": [1.0, 3.0], # <- all state-ins remain as-is
"seq_lens": [1, 2]}
"""
seq_lens = self.get(SampleBatch.SEQ_LENS)
if seq_lens is None:
raise ValueError(
"Cannot right-zero-pad SampleBatch if no `seq_lens` field "
f"present! SampleBatch={self}"
)
length = len(seq_lens) * max_seq_len
def _zero_pad_in_place(path, value):
# Skip "state_in_..." columns and "seq_lens".
if (exclude_states is True and path[0].startswith("state_in_")) or path[
0
] == SampleBatch.SEQ_LENS:
return
# Generate zero-filled primer of len=max_seq_len.
if value.dtype == object or value.dtype.type is np.str_:
f_pad = [None] * length
else:
# Make sure type doesn't change.
f_pad = np.zeros((length,) + np.shape(value)[1:], dtype=value.dtype)
# Fill primer with data.
f_pad_base = f_base = 0
for len_ in self[SampleBatch.SEQ_LENS]:
f_pad[f_pad_base : f_pad_base + len_] = value[f_base : f_base + len_]
f_pad_base += max_seq_len
f_base += len_
assert f_base == len(value), value
# Update our data in-place.
curr = self
for i, p in enumerate(path):
if i == len(path) - 1:
curr[p] = f_pad
curr = curr[p]
self_as_dict = dict(self)
tree.map_structure_with_path(_zero_pad_in_place, self_as_dict)
# Set flags to indicate, we are now zero-padded (and to what extend).
self.zero_padded = True
self.max_seq_len = max_seq_len
return self
@ExperimentalAPI
def to_device(
self,
device,
framework: str = "torch",
pin_memory: bool = False,
use_stream: bool = False,
stream: Optional[Union["torch.cuda.Stream", "torch.cuda.Stream"]] = None,
):
"""TODO: transfer batch to given device as framework tensor."""
if framework == "torch":
assert torch is not None
for k, v in self.items():
self[k] = convert_to_torch_tensor(
v,
device,
pin_memory=pin_memory,
use_stream=use_stream,
stream=stream,
)
else:
raise NotImplementedError
return self
@PublicAPI
def size_bytes(self) -> int:
"""Returns sum over number of bytes of all data buffers.
For numpy arrays, we use ``.nbytes``. For all other value types, we use
sys.getsizeof(...).
Returns:
The overall size in bytes of the data buffer (all columns).
"""
return sum(
v.nbytes if isinstance(v, np.ndarray) else sys.getsizeof(v)
for v in tree.flatten(self)
)
def get(self, key, default=None):
"""Returns one column (by key) from the data or a default value."""
try:
return self.__getitem__(key)
except KeyError:
return default
@PublicAPI
def as_multi_agent(self, module_id: Optional[ModuleID] = None) -> "MultiAgentBatch":
"""Returns the respective MultiAgentBatch
Note, if `module_id` is not provided uses `DEFAULT_POLICY`_ID`.
Args;
module_id: An optional module ID. If `None` the `DEFAULT_POLICY_ID`
is used.
Returns:
The MultiAgentBatch (using DEFAULT_POLICY_ID) corresponding
to this SampleBatch.
"""
return MultiAgentBatch({module_id or DEFAULT_POLICY_ID: self}, self.count)
@PublicAPI
def __getitem__(self, key: Union[str, slice]) -> TensorType:
"""Returns one column (by key) from the data or a sliced new batch.
Args:
key: The key (column name) to return or
a slice object for slicing this SampleBatch.
Returns:
The data under the given key or a sliced version of this batch.
"""
if isinstance(key, slice):
return self._slice(key)
# Special key DONES -> Translate to `TERMINATEDS | TRUNCATEDS` to reflect
# the old meaning of DONES.
if key == SampleBatch.DONES:
return self[SampleBatch.TERMINATEDS]
# Backward compatibility for when "input-dicts" were used.
elif key == "is_training":
if log_once("SampleBatch['is_training']"):
deprecation_warning(
old="SampleBatch['is_training']",
new="SampleBatch.is_training",
error=False,
)
return self.is_training
if not hasattr(self, key) and key in self:
self.accessed_keys.add(key)
value = dict.__getitem__(self, key)
if self.get_interceptor is not None:
if key not in self.intercepted_values:
self.intercepted_values[key] = self.get_interceptor(value)
value = self.intercepted_values[key]
return value
@PublicAPI
def __setitem__(self, key, item) -> None:
"""Inserts (overrides) an entire column (by key) in the data buffer.
Args:
key: The column name to set a value for.
item: The data to insert.
"""
# Disallow setting DONES key directly.
if key == SampleBatch.DONES:
raise KeyError(
"Cannot set `DONES` anymore in a SampleBatch! "
"Instead, set the new TERMINATEDS and TRUNCATEDS keys. The values under"
" DONES will then be automatically computed using terminated|truncated."
)
# Defend against creating SampleBatch via pickle (no property
# `added_keys` and first item is already set).
elif not hasattr(self, "added_keys"):
dict.__setitem__(self, key, item)
return
# Backward compatibility for when "input-dicts" were used.
if key == "is_training":
if log_once("SampleBatch['is_training']"):
deprecation_warning(
old="SampleBatch['is_training']",
new="SampleBatch.is_training",
error=False,
)
self._is_training = item
return
if key not in self:
self.added_keys.add(key)
dict.__setitem__(self, key, item)
if key in self.intercepted_values:
self.intercepted_values[key] = item
@property
def is_training(self):
if self.get_interceptor is not None and isinstance(self._is_training, bool):
if "_is_training" not in self.intercepted_values:
self.intercepted_values["_is_training"] = self.get_interceptor(
self._is_training
)
return self.intercepted_values["_is_training"]
return self._is_training
def set_training(self, training: Union[bool, "tf1.placeholder"] = True):
"""Sets the `is_training` flag for this SampleBatch."""
self._is_training = training
self.intercepted_values.pop("_is_training", None)
@PublicAPI
def __delitem__(self, key):
self.deleted_keys.add(key)
dict.__delitem__(self, key)
@DeveloperAPI
def compress(
self, bulk: bool = False, columns: Set[str] = frozenset(["obs", "new_obs"])
) -> "SampleBatch":
"""Compresses the data buffers (by column) in place.
Args:
bulk: Whether to compress across the batch dimension (0)
as well. If False will compress n separate list items, where n
is the batch size.
columns: The columns to compress. Default: Only
compress the obs and new_obs columns.
Returns:
This very (now compressed) SampleBatch.
"""
def _compress_in_place(path, value):
if path[0] not in columns:
return
curr = self
for i, p in enumerate(path):
if i == len(path) - 1:
if bulk:
curr[p] = pack(value)
else:
curr[p] = np.array([pack(o) for o in value])
curr = curr[p]
tree.map_structure_with_path(_compress_in_place, self)
return self
@DeveloperAPI
def decompress_if_needed(
self, columns: Set[str] = frozenset(["obs", "new_obs"])
) -> "SampleBatch":
"""Decompresses data buffers (per column if not compressed) in place.
Args:
columns: The columns to decompress. Default: Only
decompress the obs and new_obs columns.
Returns:
This very (now uncompressed) SampleBatch.
"""
def _decompress_in_place(path, value):
if path[0] not in columns:
return
curr = self
for p in path[:-1]:
curr = curr[p]
# Bulk compressed.
if is_compressed(value):
curr[path[-1]] = unpack(value)
# Non bulk compressed.
elif len(value) > 0 and is_compressed(value[0]):
curr[path[-1]] = np.array([unpack(o) for o in value])
tree.map_structure_with_path(_decompress_in_place, self)
return self
@DeveloperAPI
def set_get_interceptor(self, fn):
"""Sets a function to be called on every getitem."""
# If get-interceptor changes, must erase old intercepted values.
if fn is not self.get_interceptor:
self.intercepted_values = {}
self.get_interceptor = fn
def __repr__(self):
keys = list(self.keys())
if self.get(SampleBatch.SEQ_LENS) is None:
return f"SampleBatch({self.count}: {keys})"
else:
keys.remove(SampleBatch.SEQ_LENS)
return (
f"SampleBatch({self.count} " f"(seqs={len(self['seq_lens'])}): {keys})"
)
def _slice(self, slice_: slice) -> "SampleBatch":
"""Helper method to handle SampleBatch slicing using a slice object.
The returned SampleBatch uses the same underlying data object as
`self`, so changing the slice will also change `self`.
Note that only zero or positive bounds are allowed for both start
and stop values. The slice step must be 1 (or None, which is the
same).
Args:
slice_: The python slice object to slice by.
Returns:
A new SampleBatch, however "linking" into the same data
(sliced) as self.
"""
if self._slice_seq_lens_in_B:
return self._batch_slice(slice_)
start = slice_.start or 0
stop = slice_.stop or len(self)
# If stop goes beyond the length of this batch -> Make it go till the
# end only (including last item).
# Analogous to `l = [0, 1, 2]; l[:100] -> [0, 1, 2];`.
if stop > len(self):
stop = len(self)
if (
self.get(SampleBatch.SEQ_LENS) is not None
and len(self[SampleBatch.SEQ_LENS]) > 0
):
# Build our slice-map, if not done already.
if not self._slice_map:
sum_ = 0
for i, l in enumerate(map(int, self[SampleBatch.SEQ_LENS])):
self._slice_map.extend([(i, sum_)] * l)
sum_ = sum_ + l
# In case `stop` points to the very end (lengths of this
# batch), return the last sequence (the -1 here makes sure we
# never go beyond it; would result in an index error below).
self._slice_map.append((len(self[SampleBatch.SEQ_LENS]), sum_))
start_seq_len, start_unpadded = self._slice_map[start]
stop_seq_len, stop_unpadded = self._slice_map[stop]
start_padded = start_unpadded
stop_padded = stop_unpadded
if self.zero_padded:
start_padded = start_seq_len * self.max_seq_len
stop_padded = stop_seq_len * self.max_seq_len
def map_(path, value):
if path[0] != SampleBatch.SEQ_LENS and not path[0].startswith(
"state_in_"
):
return value[start_padded:stop_padded]
else:
return value[start_seq_len:stop_seq_len]
infos = self.pop(SampleBatch.INFOS, None)
data = tree.map_structure_with_path(map_, self)
if infos is not None and isinstance(infos, (list, np.ndarray)):
self[SampleBatch.INFOS] = infos
data[SampleBatch.INFOS] = infos[start_unpadded:stop_unpadded]
return SampleBatch(
data,
_is_training=self.is_training,
_time_major=self.time_major,
_zero_padded=self.zero_padded,
_max_seq_len=self.max_seq_len if self.zero_padded else None,
_num_grad_updates=self.num_grad_updates,
)
else:
infos = self.pop(SampleBatch.INFOS, None)
data = tree.map_structure(lambda s: s[start:stop], self)
if infos is not None and isinstance(infos, (list, np.ndarray)):
self[SampleBatch.INFOS] = infos
data[SampleBatch.INFOS] = infos[start:stop]
return SampleBatch(
data,
_is_training=self.is_training,
_time_major=self.time_major,
_num_grad_updates=self.num_grad_updates,
)
@Deprecated(error=False)
def _get_slice_indices(self, slice_size):
data_slices = []
data_slices_states = []
if (
self.get(SampleBatch.SEQ_LENS) is not None
and len(self[SampleBatch.SEQ_LENS]) > 0
):
assert np.all(self[SampleBatch.SEQ_LENS] < slice_size), (
"ERROR: `slice_size` must be larger than the max. seq-len "
"in the batch!"
)
start_pos = 0
current_slize_size = 0
actual_slice_idx = 0
start_idx = 0
idx = 0
while idx < len(self[SampleBatch.SEQ_LENS]):
seq_len = self[SampleBatch.SEQ_LENS][idx]
current_slize_size += seq_len
actual_slice_idx += (
seq_len if not self.zero_padded else self.max_seq_len
)
# Complete minibatch -> Append to data_slices.
if current_slize_size >= slice_size:
end_idx = idx + 1
# We are not zero-padded yet; all sequences are
# back-to-back.
if not self.zero_padded:
data_slices.append((start_pos, start_pos + slice_size))
start_pos += slice_size
if current_slize_size > slice_size:
overhead = current_slize_size - slice_size
start_pos -= seq_len - overhead
idx -= 1
# We are already zero-padded: Cut in chunks of max_seq_len.
else:
data_slices.append((start_pos, actual_slice_idx))
start_pos = actual_slice_idx
data_slices_states.append((start_idx, end_idx))
current_slize_size = 0
start_idx = idx + 1
idx += 1
else:
i = 0
while i < self.count:
data_slices.append((i, i + slice_size))
i += slice_size
return data_slices, data_slices_states
@ExperimentalAPI
def get_single_step_input_dict(
self,
view_requirements: ViewRequirementsDict,
index: Union[str, int] = "last",
) -> "SampleBatch":
"""Creates single ts SampleBatch at given index from `self`.
For usage as input-dict for model (action or value function) calls.
Args:
view_requirements: A view requirements dict from the model for
which to produce the input_dict.
index: An integer index value indicating the
position in the trajectory for which to generate the
compute_actions input dict. Set to "last" to generate the dict
at the very end of the trajectory (e.g. for value estimation).
Note that "last" is different from -1, as "last" will use the
final NEXT_OBS as observation input.
Returns:
The (single-timestep) input dict for ModelV2 calls.
"""
last_mappings = {
SampleBatch.OBS: SampleBatch.NEXT_OBS,
SampleBatch.PREV_ACTIONS: SampleBatch.ACTIONS,
SampleBatch.PREV_REWARDS: SampleBatch.REWARDS,
}
input_dict = {}
for view_col, view_req in view_requirements.items():
if view_req.used_for_compute_actions is False:
continue
# Create batches of size 1 (single-agent input-dict).
data_col = view_req.data_col or view_col
if index == "last":
data_col = last_mappings.get(data_col, data_col)
# Range needed.
if view_req.shift_from is not None:
# Batch repeat value > 1: We have single frames in the
# batch at each timestep (for the `data_col`).
data = self[view_col][-1]
traj_len = len(self[data_col])
missing_at_end = traj_len % view_req.batch_repeat_value
# Index into the observations column must be shifted by
# -1 b/c index=0 for observations means the current (last
# seen) observation (after having taken an action).
obs_shift = (
-1 if data_col in [SampleBatch.OBS, SampleBatch.NEXT_OBS] else 0
)
from_ = view_req.shift_from + obs_shift
to_ = view_req.shift_to + obs_shift + 1
if to_ == 0:
to_ = None
input_dict[view_col] = np.array(
[
np.concatenate([data, self[data_col][-missing_at_end:]])[
from_:to_
]
]
)
# Single index.
else:
input_dict[view_col] = tree.map_structure(
lambda v: v[-1:], # keep as array (w/ 1 element)
self[data_col],
)
# Single index somewhere inside the trajectory (non-last).
else:
input_dict[view_col] = self[data_col][
index : index + 1 if index != -1 else None
]
return SampleBatch(input_dict, seq_lens=np.array([1], dtype=np.int32))
@PublicAPI
class MultiAgentBatch:
"""A batch of experiences from multiple agents in the environment.
Attributes:
policy_batches (Dict[PolicyID, SampleBatch]): Dict mapping policy IDs to
SampleBatches of experiences.
count: The number of env steps in this batch.
"""
@PublicAPI
def __init__(self, policy_batches: Dict[PolicyID, SampleBatch], env_steps: int):
"""Initialize a MultiAgentBatch instance.
Args:
policy_batches: Dict mapping policy IDs to SampleBatches of experiences.
env_steps: The number of environment steps in the environment
this batch contains. This will be less than the number of
transitions this batch contains across all policies in total.
"""
for v in policy_batches.values():
assert isinstance(v, SampleBatch)
self.policy_batches = policy_batches
# Called "count" for uniformity with SampleBatch.
# Prefer to access this via the `env_steps()` method when possible
# for clarity.
self.count = env_steps
@PublicAPI
def env_steps(self) -> int:
"""The number of env steps (there are >= 1 agent steps per env step).
Returns:
The number of environment steps contained in this batch.
"""
return self.count
@PublicAPI
def __len__(self) -> int:
"""Same as `self.env_steps()`."""
return self.count
@PublicAPI
def agent_steps(self) -> int:
"""The number of agent steps (there are >= 1 agent steps per env step).
Returns:
The number of agent steps total in this batch.
"""
ct = 0
for batch in self.policy_batches.values():
ct += batch.count
return ct
@PublicAPI
def timeslices(self, k: int) -> List["MultiAgentBatch"]:
"""Returns k-step batches holding data for each agent at those steps.
For examples, suppose we have agent1 observations [a1t1, a1t2, a1t3],
for agent2, [a2t1, a2t3], and for agent3, [a3t3] only.
Calling timeslices(1) would return three MultiAgentBatches containing
[a1t1, a2t1], [a1t2], and [a1t3, a2t3, a3t3].
Calling timeslices(2) would return two MultiAgentBatches containing
[a1t1, a1t2, a2t1], and [a1t3, a2t3, a3t3].
This method is used to implement "lockstep" replay mode. Note that this
method does not guarantee each batch contains only data from a single
unroll. Batches might contain data from multiple different envs.
"""
from ray.rllib.evaluation.sample_batch_builder import SampleBatchBuilder
# Build a sorted set of (eps_id, t, policy_id, data...)
steps = []
for policy_id, batch in self.policy_batches.items():
for row in batch.rows():
steps.append(
(
row[SampleBatch.EPS_ID],
row[SampleBatch.T],
row[SampleBatch.AGENT_INDEX],
policy_id,
row,
)
)
steps.sort()
finished_slices = []
cur_slice = collections.defaultdict(SampleBatchBuilder)
cur_slice_size = 0
def finish_slice():
nonlocal cur_slice_size
assert cur_slice_size > 0
batch = MultiAgentBatch(
{k: v.build_and_reset() for k, v in cur_slice.items()}, cur_slice_size
)
cur_slice_size = 0
cur_slice.clear()
finished_slices.append(batch)
# For each unique env timestep.
for _, group in itertools.groupby(steps, lambda x: x[:2]):
# Accumulate into the current slice.
for _, _, _, policy_id, row in group:
cur_slice[policy_id].add_values(**row)
cur_slice_size += 1
# Slice has reached target number of env steps.
if cur_slice_size >= k:
finish_slice()
assert cur_slice_size == 0
if cur_slice_size > 0:
finish_slice()
assert len(finished_slices) > 0, finished_slices
return finished_slices
@staticmethod
@PublicAPI
def wrap_as_needed(
policy_batches: Dict[PolicyID, SampleBatch], env_steps: int
) -> Union[SampleBatch, "MultiAgentBatch"]:
"""Returns SampleBatch or MultiAgentBatch, depending on given policies.
If policy_batches is empty (i.e. {}) it returns an empty MultiAgentBatch.
Args:
policy_batches: Mapping from policy ids to SampleBatch.
env_steps: Number of env steps in the batch.
Returns:
The single default policy's SampleBatch or a MultiAgentBatch
(more than one policy).
"""
if len(policy_batches) == 1 and DEFAULT_POLICY_ID in policy_batches:
return policy_batches[DEFAULT_POLICY_ID]
return MultiAgentBatch(policy_batches=policy_batches, env_steps=env_steps)
@staticmethod
@PublicAPI
@Deprecated(new="concat_samples() from rllib.policy.sample_batch", error=True)
def concat_samples(samples: List["MultiAgentBatch"]) -> "MultiAgentBatch":
return concat_samples_into_ma_batch(samples)
@PublicAPI
def copy(self) -> "MultiAgentBatch":
"""Deep-copies self into a new MultiAgentBatch.
Returns:
The copy of self with deep-copied data.
"""
return MultiAgentBatch(
{k: v.copy() for (k, v) in self.policy_batches.items()}, self.count
)
@ExperimentalAPI
def to_device(
self,
device,
framework="torch",
pin_memory: bool = False,
use_stream: bool = False,
stream: Optional[Union["torch.cuda.Stream", "torch.cuda.Stream"]] = None,
):
"""TODO: transfer batch to given device as framework tensor."""
if framework == "torch":
assert torch is not None
for pid, policy_batch in self.policy_batches.items():
self.policy_batches[pid] = policy_batch.to_device(
device,
framework=framework,
pin_memory=pin_memory,
use_stream=use_stream,
stream=stream,
)
else:
raise NotImplementedError
return self
@PublicAPI
def size_bytes(self) -> int:
"""
Returns:
The overall size in bytes of all policy batches (all columns).
"""
return sum(b.size_bytes() for b in self.policy_batches.values())
@DeveloperAPI
def compress(
self, bulk: bool = False, columns: Set[str] = frozenset(["obs", "new_obs"])
) -> None:
"""Compresses each policy batch (per column) in place.
Args:
bulk: Whether to compress across the batch dimension (0)
as well. If False will compress n separate list items, where n
is the batch size.
columns: Set of column names to compress.
"""
for batch in self.policy_batches.values():
batch.compress(bulk=bulk, columns=columns)
@DeveloperAPI
def decompress_if_needed(
self, columns: Set[str] = frozenset(["obs", "new_obs"])
) -> "MultiAgentBatch":
"""Decompresses each policy batch (per column), if already compressed.
Args:
columns: Set of column names to decompress.
Returns:
Self.
"""
for batch in self.policy_batches.values():
batch.decompress_if_needed(columns)
return self
@DeveloperAPI
def as_multi_agent(self) -> "MultiAgentBatch":
"""Simply returns `self` (already a MultiAgentBatch).
Returns:
This very instance of MultiAgentBatch.
"""
return self
def __getitem__(self, key: str) -> SampleBatch:
"""Returns the SampleBatch for the given policy id."""
return self.policy_batches[key]
def __str__(self):
return "MultiAgentBatch({}, env_steps={})".format(
str(self.policy_batches), self.count
)
def __repr__(self):
return "MultiAgentBatch({}, env_steps={})".format(
str(self.policy_batches), self.count
)
@PublicAPI
def concat_samples(samples: List[SampleBatchType]) -> SampleBatchType:
"""Concatenates a list of SampleBatches or MultiAgentBatches.
If all items in the list are or SampleBatch typ4, the output will be
a SampleBatch type. Otherwise, the output will be a MultiAgentBatch type.
If input is a mixture of SampleBatch and MultiAgentBatch types, it will treat
SampleBatch objects as MultiAgentBatch types with 'default_policy' key and
concatenate it with th rest of MultiAgentBatch objects.
Empty samples are simply ignored.
Args:
samples: List of SampleBatches or MultiAgentBatches to be
concatenated.
Returns:
A new (concatenated) SampleBatch or MultiAgentBatch.
.. testcode::
:skipif: True
import numpy as np
from ray.rllib.policy.sample_batch import SampleBatch
b1 = SampleBatch({"a": np.array([1, 2]),
"b": np.array([10, 11])})
b2 = SampleBatch({"a": np.array([3]),
"b": np.array([12])})
print(concat_samples([b1, b2]))
c1 = MultiAgentBatch({'default_policy': {
"a": np.array([1, 2]),
"b": np.array([10, 11])
}}, env_steps=2)
c2 = SampleBatch({"a": np.array([3]),
"b": np.array([12])})
print(concat_samples([b1, b2]))
.. testoutput::
{"a": np.array([1, 2, 3]), "b": np.array([10, 11, 12])}
MultiAgentBatch = {'default_policy': {"a": np.array([1, 2, 3]),
"b": np.array([10, 11, 12])}}
"""
if any(isinstance(s, MultiAgentBatch) for s in samples):
return concat_samples_into_ma_batch(samples)
# the output is a SampleBatch type
concatd_seq_lens = []
concatd_num_grad_updates = [0, 0.0] # [0]=count; [1]=weighted sum values
concated_samples = []
# Make sure these settings are consistent amongst all batches.
zero_padded = max_seq_len = time_major = None
for s in samples:
if s.count <= 0:
continue
if max_seq_len is None:
zero_padded = s.zero_padded
max_seq_len = s.max_seq_len
time_major = s.time_major
# Make sure these settings are consistent amongst all batches.
if s.zero_padded != zero_padded or s.time_major != time_major:
raise ValueError(
"All SampleBatches' `zero_padded` and `time_major` settings "
"must be consistent!"
)
if (
s.max_seq_len is None or max_seq_len is None
) and s.max_seq_len != max_seq_len:
raise ValueError(
"Samples must consistently either provide or omit `max_seq_len`!"
)
elif zero_padded and s.max_seq_len != max_seq_len:
raise ValueError(
"For `zero_padded` SampleBatches, the values of `max_seq_len` "
"must be consistent!"
)
if max_seq_len is not None:
max_seq_len = max(max_seq_len, s.max_seq_len)
if s.get(SampleBatch.SEQ_LENS) is not None:
concatd_seq_lens.extend(s[SampleBatch.SEQ_LENS])
if s.num_grad_updates is not None:
concatd_num_grad_updates[0] += s.count
concatd_num_grad_updates[1] += s.num_grad_updates * s.count
concated_samples.append(s)
# If we don't have any samples (0 or only empty SampleBatches),
# return an empty SampleBatch here.
if len(concated_samples) == 0:
return SampleBatch()
# Collect the concat'd data.
concatd_data = {}
for k in concated_samples[0].keys():
if k == SampleBatch.INFOS:
concatd_data[k] = _concat_values(
*[s[k] for s in concated_samples],
time_major=time_major,
)
else:
values_to_concat = [c[k] for c in concated_samples]
_concat_values_w_time = partial(_concat_values, time_major=time_major)
concatd_data[k] = tree.map_structure(
_concat_values_w_time, *values_to_concat
)
if concatd_seq_lens != [] and torch and torch.is_tensor(concatd_seq_lens[0]):
concatd_seq_lens = torch.Tensor(concatd_seq_lens)
elif concatd_seq_lens != [] and tf and tf.is_tensor(concatd_seq_lens[0]):
concatd_seq_lens = tf.convert_to_tensor(concatd_seq_lens)
# Return a new (concat'd) SampleBatch.
return SampleBatch(
concatd_data,
seq_lens=concatd_seq_lens,
_time_major=time_major,
_zero_padded=zero_padded,
_max_seq_len=max_seq_len,
# Compute weighted average of the num_grad_updates for the batches
# (assuming they all come from the same policy).
_num_grad_updates=(
concatd_num_grad_updates[1] / (concatd_num_grad_updates[0] or 1.0)
),
)
@PublicAPI
def concat_samples_into_ma_batch(samples: List[SampleBatchType]) -> "MultiAgentBatch":
"""Concatenates a list of SampleBatchTypes to a single MultiAgentBatch type.
This function, as opposed to concat_samples() forces the output to always be
MultiAgentBatch which is more generic than SampleBatch.
Args:
samples: List of SampleBatches or MultiAgentBatches to be
concatenated.
Returns:
A new (concatenated) MultiAgentBatch.
.. testcode::
:skipif: True
import numpy as np
from ray.rllib.policy.sample_batch import SampleBatch
b1 = MultiAgentBatch({'default_policy': {
"a": np.array([1, 2]),
"b": np.array([10, 11])
}}, env_steps=2)
b2 = SampleBatch({"a": np.array([3]),
"b": np.array([12])})
print(concat_samples([b1, b2]))
.. testoutput::
{'default_policy': {"a": np.array([1, 2, 3]),
"b": np.array([10, 11, 12])}}
"""
policy_batches = collections.defaultdict(list)
env_steps = 0
for s in samples:
# Some batches in `samples` may be SampleBatch.
if isinstance(s, SampleBatch):
# If empty SampleBatch: ok (just ignore).
if len(s) <= 0:
continue
else:
# if non-empty: just convert to MA-batch and move forward
s = s.as_multi_agent()
elif not isinstance(s, MultiAgentBatch):
# Otherwise: Error.
raise ValueError(
"`concat_samples_into_ma_batch` can only concat "
"SampleBatch|MultiAgentBatch objects, not {}!".format(type(s).__name__)
)
for key, batch in s.policy_batches.items():
policy_batches[key].append(batch)
env_steps += s.env_steps()
out = {}
for key, batches in policy_batches.items():
out[key] = concat_samples(batches)
return MultiAgentBatch(out, env_steps)
def _concat_values(*values, time_major=None) -> TensorType:
"""Concatenates a list of values.
Args:
values: The values to concatenate.
time_major: Whether to concatenate along the first axis
(time_major=False) or the second axis (time_major=True).
"""
if torch and torch.is_tensor(values[0]):
return torch.cat(values, dim=1 if time_major else 0)
elif isinstance(values[0], np.ndarray):
return np.concatenate(values, axis=1 if time_major else 0)
elif tf and tf.is_tensor(values[0]):
return tf.concat(values, axis=1 if time_major else 0)
elif isinstance(values[0], list):
concatenated_list = []
for sublist in values:
concatenated_list.extend(sublist)
return concatenated_list
else:
raise ValueError(
f"Unsupported type for concatenation: {type(values[0])} "
f"first element: {values[0]}"
)
@DeveloperAPI
def convert_ma_batch_to_sample_batch(batch: SampleBatchType) -> SampleBatch:
"""Converts a MultiAgentBatch to a SampleBatch if necessary.
Args:
batch: The SampleBatchType to convert.
Returns:
batch: the converted SampleBatch
Raises:
ValueError if the MultiAgentBatch has more than one policy_id
or if the policy_id is not `DEFAULT_POLICY_ID`
"""
if isinstance(batch, MultiAgentBatch):
policy_keys = batch.policy_batches.keys()
if len(policy_keys) == 1 and DEFAULT_POLICY_ID in policy_keys:
batch = batch.policy_batches[DEFAULT_POLICY_ID]
else:
raise ValueError(
"RLlib tried to convert a multi agent-batch with data from more "
"than one policy to a single-agent batch. This is not supported and "
"may be due to a number of issues. Here are two possible ones:"
"1) Off-Policy Estimation is not implemented for "
"multi-agent batches. You can set `off_policy_estimation_methods: {}` "
"to resolve this."
"2) Loading multi-agent data for offline training is not implemented."
"Load single-agent data instead to resolve this."
)
return batch
Reference in New Issue View Git Blame Copy Permalink