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/env/multi_agent_episode.py
T
wehub-resource-sync f1825c8ceb chore: import upstream snapshot with attribution
2026-07-13 13:17:40 +08:00

2874 lines
136 KiB
Python
Raw Permalink Blame History

import copy
import time
import uuid
from collections import defaultdict
from typing import (
Any,
Callable,
Collection,
DefaultDict,
Dict,
List,
Optional,
Set,
Union,
)
import gymnasium as gym
from ray._common.deprecation import Deprecated
from ray.rllib.env.single_agent_episode import SingleAgentEpisode
from ray.rllib.env.utils.infinite_lookback_buffer import InfiniteLookbackBuffer
from ray.rllib.policy.sample_batch import MultiAgentBatch
from ray.rllib.utils import force_list
from ray.rllib.utils.error import MultiAgentEnvError
from ray.rllib.utils.spaces.space_utils import batch
from ray.rllib.utils.typing import AgentID, ModuleID, MultiAgentDict
from ray.util.annotations import PublicAPI
# TODO (simon): Include cases in which the number of agents in an
# episode are shrinking or growing during the episode itself.
@PublicAPI(stability="alpha")
class MultiAgentEpisode:
"""Stores multi-agent episode data.
The central attribute of the class is the timestep mapping
`self.env_t_to_agent_t` that maps AgentIDs to their specific environment steps to
the agent's own scale/timesteps.
Each AgentID in the `MultiAgentEpisode` has its own `SingleAgentEpisode` object
in which this agent's data is stored. Together with the env_t_to_agent_t mapping,
we can extract information either on any individual agent's timescale or from
the (global) multi-agent environment timescale.
Extraction of data from a MultiAgentEpisode happens via the getter APIs, e.g.
`get_observations()`, which work analogous to the ones implemented in the
`SingleAgentEpisode` class.
Note that recorded `terminateds`/`truncateds` come as simple
`MultiAgentDict`s mapping AgentID to bools and thus have no assignment to a
certain timestep (analogous to a SingleAgentEpisode's single `terminated/truncated`
boolean flag). Instead we assign it to the last observation recorded.
Theoretically, there could occur edge cases in some environments
where an agent receives partial rewards and then terminates without
a last observation. In these cases, we duplicate the last observation.
Also, if no initial observation has been received yet for an agent, but
some rewards for this same agent already occurred, we delete the agent's data
up to here, b/c there is nothing to learn from these "premature" rewards.
"""
__slots__ = (
"id_",
"agent_to_module_mapping_fn",
"_agent_to_module_mapping",
"observation_space",
"action_space",
"env_t_started",
"env_t",
"agent_t_started",
"env_t_to_agent_t",
"_hanging_actions_end",
"_hanging_extra_model_outputs_end",
"_hanging_rewards_end",
"_hanging_rewards_begin",
"is_terminated",
"is_truncated",
"agent_episodes",
"_last_step_time",
"_len_lookback_buffers",
"_start_time",
"_custom_data",
)
SKIP_ENV_TS_TAG = "S"
def __init__(
self,
id_: Optional[str] = None,
*,
observations: Optional[List[MultiAgentDict]] = None,
observation_space: Optional[gym.Space] = None,
infos: Optional[List[MultiAgentDict]] = None,
actions: Optional[List[MultiAgentDict]] = None,
action_space: Optional[gym.Space] = None,
rewards: Optional[List[MultiAgentDict]] = None,
terminateds: Union[MultiAgentDict, bool] = False,
truncateds: Union[MultiAgentDict, bool] = False,
extra_model_outputs: Optional[List[MultiAgentDict]] = None,
env_t_started: Optional[int] = None,
agent_t_started: Optional[Dict[AgentID, int]] = None,
len_lookback_buffer: Union[int, str] = "auto",
agent_episode_ids: Optional[Dict[AgentID, str]] = None,
agent_module_ids: Optional[Dict[AgentID, ModuleID]] = None,
agent_to_module_mapping_fn: Optional[
Callable[[AgentID, "MultiAgentEpisode"], ModuleID]
] = None,
):
"""Initializes a `MultiAgentEpisode`.
Args:
id_: Optional. Either a string to identify an episode or None.
If None, a hexadecimal id is created. In case of providing
a string, make sure that it is unique, as episodes get
concatenated via this string.
observations: A list of dictionaries mapping agent IDs to observations.
Can be None. If provided, should match all other episode data
(actions, rewards, etc.) in terms of list lengths and agent IDs.
observation_space: An optional gym.spaces.Dict mapping agent IDs to
individual agents' spaces, which all (individual agents') observations
should abide to. If not None and this MultiAgentEpisode is numpy'ized
(via the `self.to_numpy()` method), and data is appended or set, the new
data will be checked for correctness.
infos: A list of dictionaries mapping agent IDs to info dicts.
Can be None. If provided, should match all other episode data
(observations, rewards, etc.) in terms of list lengths and agent IDs.
actions: A list of dictionaries mapping agent IDs to actions.
Can be None. If provided, should match all other episode data
(observations, rewards, etc.) in terms of list lengths and agent IDs.
action_space: An optional gym.spaces.Dict mapping agent IDs to
individual agents' spaces, which all (individual agents') actions
should abide to. If not None and this MultiAgentEpisode is numpy'ized
(via the `self.to_numpy()` method), and data is appended or set, the new
data will be checked for correctness.
rewards: A list of dictionaries mapping agent IDs to rewards.
Can be None. If provided, should match all other episode data
(actions, rewards, etc.) in terms of list lengths and agent IDs.
terminateds: A boolean defining if an environment has
terminated OR a MultiAgentDict mapping individual agent ids
to boolean flags indicating whether individual agents have terminated.
A special __all__ key in these dicts indicates, whether the episode
is terminated for all agents.
The default is `False`, i.e. the episode has not been terminated.
truncateds: A boolean defining if the environment has been
truncated OR a MultiAgentDict mapping individual agent ids
to boolean flags indicating whether individual agents have been
truncated. A special __all__ key in these dicts indicates, whether the
episode is truncated for all agents.
The default is `False`, i.e. the episode has not been truncated.
extra_model_outputs: A list of dictionaries mapping agent IDs to their
corresponding extra model outputs. Each of these "outputs" is a dict
mapping keys (str) to model output values, for example for
`key=STATE_OUT`, the values would be the internal state outputs for
that agent.
env_t_started: The env timestep (int) that defines the starting point
of the episode. This is only larger zero, if an already ongoing episode
chunk is being created, for example by slicing an ongoing episode or
by calling the `cut()` method on an ongoing episode.
agent_t_started: A dict mapping AgentIDs to the agent's timestep
(not global env timestep) at which its SingleAgentEpisode chunk started.
len_lookback_buffer: The size of the lookback buffers to keep in
front of this Episode for each type of data (observations, actions,
etc..). If larger 0, will interpret the first `len_lookback_buffer`
items in each type of data as NOT part of this actual
episode chunk, but instead serve as "historical" record that may be
viewed and used to derive new data from. For example, it might be
necessary to have a lookback buffer of four if you would like to do
observation frame stacking and your episode has been cut and you are now
operating on a new chunk (continuing from the cut one). Then, for the
first 3 items, you would have to be able to look back into the old
chunk's data.
If `len_lookback_buffer` is "auto" (default), will interpret all
provided data in the constructor as part of the lookback buffers.
agent_episode_ids: An optional dict mapping AgentIDs
to their corresponding `SingleAgentEpisode`. If None, each
`SingleAgentEpisode` in `MultiAgentEpisode.agent_episodes`
will generate a hexadecimal code. If a dictionary is provided,
make sure that IDs are unique, because the agents' `SingleAgentEpisode`
instances are concatenated or recreated by it.
agent_module_ids: An optional dict mapping AgentIDs to their respective
ModuleIDs (these mapping are always valid for an entire episode and
thus won't change during the course of this episode). If a mapping from
agent to module has already been provided via this dict, the (optional)
`agent_to_module_mapping_fn` will NOT be used again to map the same
agent (agents do not change their assigned module in the course of
one episode).
agent_to_module_mapping_fn: A callable taking an AgentID and a
MultiAgentEpisode as args and returning a ModuleID. Used to map agents
that have not been mapped yet (because they just entered this episode)
to a ModuleID. The resulting ModuleID is only stored inside the agent's
SingleAgentEpisode object.
"""
self.id_: str = id_ or uuid.uuid4().hex
if agent_to_module_mapping_fn is None:
from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
agent_to_module_mapping_fn = (
AlgorithmConfig.DEFAULT_AGENT_TO_MODULE_MAPPING_FN
)
self.agent_to_module_mapping_fn = agent_to_module_mapping_fn
# In case a user - e.g. via callbacks - already forces a mapping to happen
# via the `module_for()` API even before the agent has entered the episode
# (and has its SingleAgentEpisode created), we store all aldeary done mappings
# in this dict here.
self._agent_to_module_mapping: Dict[AgentID, ModuleID] = agent_module_ids or {}
# Lookback buffer length is not provided. Interpret all provided data as
# lookback buffer.
if len_lookback_buffer == "auto":
len_lookback_buffer = len(rewards or [])
self._len_lookback_buffers = len_lookback_buffer
self.observation_space = observation_space or {}
self.action_space = action_space or {}
terminateds = terminateds or {}
truncateds = truncateds or {}
# The global last timestep of the episode and the timesteps when this chunk
# started (excluding a possible lookback buffer).
self.env_t_started = env_t_started or 0
self.env_t = (
(len(rewards) if rewards is not None else 0)
- self._len_lookback_buffers
+ self.env_t_started
)
self.agent_t_started = defaultdict(int, agent_t_started or {})
# Keeps track of the correspondence between agent steps and environment steps.
# Under each AgentID as key is a InfiniteLookbackBuffer with the following
# data in it:
# The indices of the items in the data represent environment timesteps,
# starting from index=0 for the `env.reset()` and with each `env.step()` call
# increase by 1.
# The values behind these (env timestep) indices represent the agent timesteps
# happening at these env timesteps and the special value of
# `self.SKIP_ENV_TS_TAG` means that the agent did NOT step at the given env
# timestep.
# Thus, agents that are part of the reset obs, will start their mapping data
# with a [0 ...], all other agents will start their mapping data with:
# [self.SKIP_ENV_TS_TAG, ...].
self.env_t_to_agent_t: DefaultDict[
AgentID, InfiniteLookbackBuffer
] = defaultdict(InfiniteLookbackBuffer)
# Create caches for hanging actions/rewards/extra_model_outputs.
# When an agent gets an observation (and then sends an action), but does not
# receive immediately a next observation, we store the "hanging" action (and
# related rewards and extra model outputs) in the caches postfixed w/ `_end`
# until the next observation is received.
self._hanging_actions_end = {}
self._hanging_extra_model_outputs_end = defaultdict(dict)
self._hanging_rewards_end = defaultdict(float)
# In case of a `cut()` or `slice()`, we also need to store the hanging actions,
# rewards, and extra model outputs that were already "hanging" in preceeding
# episode slice.
self._hanging_rewards_begin = defaultdict(float)
# If this is an ongoing episode than the last `__all__` should be `False`
self.is_terminated: bool = (
terminateds
if isinstance(terminateds, bool)
else terminateds.get("__all__", False)
)
# If this is an ongoing episode than the last `__all__` should be `False`
self.is_truncated: bool = (
truncateds
if isinstance(truncateds, bool)
else truncateds.get("__all__", False)
)
# The individual agent SingleAgentEpisode objects.
self.agent_episodes: Dict[AgentID, SingleAgentEpisode] = {}
self._init_single_agent_episodes(
agent_module_ids=agent_module_ids,
agent_episode_ids=agent_episode_ids,
observations=observations,
infos=infos,
actions=actions,
rewards=rewards,
terminateds=terminateds,
truncateds=truncateds,
extra_model_outputs=extra_model_outputs,
)
# Cache for custom data. May be used to store custom metrics from within a
# callback for the ongoing episode (e.g. render images).
self._custom_data = {}
# Keep timer stats on deltas between steps.
self._start_time = None
self._last_step_time = None
# Validate ourselves.
self.validate()
def add_env_reset(
self,
*,
observations: MultiAgentDict,
infos: Optional[MultiAgentDict] = None,
) -> None:
"""Stores initial observation.
Args:
observations: A dictionary mapping agent IDs to initial observations.
Note that some agents may not have an initial observation.
infos: A dictionary mapping agent IDs to initial info dicts.
Note that some agents may not have an initial info dict. If not None,
the agent IDs in `infos` must be a subset of those in `observations`
meaning it would not be allowed to have an agent with an info dict,
but not with an observation.
"""
assert not self.is_done
# Assume that this episode is completely empty and has not stepped yet.
# Leave self.env_t (and self.env_t_started) at 0.
assert self.env_t == self.env_t_started == 0
infos = infos or {}
# Note, all agents will have an initial observation, some may have an initial
# info dict as well.
for agent_id, agent_obs in observations.items():
# Update env_t_to_agent_t mapping (all agents that are part of the reset
# obs have their first mapping 0 (env_t) -> 0 (agent_t)).
self.env_t_to_agent_t[agent_id].append(0)
# Create SingleAgentEpisode, if necessary.
if agent_id not in self.agent_episodes:
self.agent_episodes[agent_id] = SingleAgentEpisode(
agent_id=agent_id,
module_id=self.module_for(agent_id),
multi_agent_episode_id=self.id_,
observation_space=self.observation_space.get(agent_id),
action_space=self.action_space.get(agent_id),
)
# Add initial observations (and infos) to the agent's episode.
self.agent_episodes[agent_id].add_env_reset(
observation=agent_obs,
infos=infos.get(agent_id),
)
# Validate our data.
self.validate()
# Start the timer for this episode.
self._start_time = time.perf_counter()
def add_env_step(
self,
observations: MultiAgentDict,
actions: MultiAgentDict,
rewards: MultiAgentDict,
infos: Optional[MultiAgentDict] = None,
*,
terminateds: Optional[MultiAgentDict] = None,
truncateds: Optional[MultiAgentDict] = None,
extra_model_outputs: Optional[MultiAgentDict] = None,
) -> None:
"""Adds a timestep to the episode.
Args:
observations: A dictionary mapping agent IDs to their corresponding
next observations. Note that some agents may not have stepped at this
timestep.
actions: Mandatory. A dictionary mapping agent IDs to their
corresponding actions. Note that some agents may not have stepped at
this timestep.
rewards: Mandatory. A dictionary mapping agent IDs to their
corresponding observations. Note that some agents may not have stepped
at this timestep.
infos: A dictionary mapping agent IDs to their
corresponding info. Note that some agents may not have stepped at this
timestep.
terminateds: A dictionary mapping agent IDs to their `terminated` flags,
indicating, whether the environment has been terminated for them.
A special `__all__` key indicates that the episode is terminated for
all agent IDs.
terminateds: A dictionary mapping agent IDs to their `truncated` flags,
indicating, whether the environment has been truncated for them.
A special `__all__` key indicates that the episode is `truncated` for
all agent IDs.
extra_model_outputs: A dictionary mapping agent IDs to their
corresponding specific model outputs (also in a dictionary; e.g.
`vf_preds` for PPO).
"""
# Cannot add data to an already done episode.
if self.is_done:
raise MultiAgentEnvError(
"Cannot call `add_env_step` on a MultiAgentEpisode that is already "
"done!"
)
infos = infos or {}
terminateds = terminateds or {}
truncateds = truncateds or {}
extra_model_outputs = extra_model_outputs or {}
# Increase (global) env step by one.
self.env_t += 1
# Find out, whether this episode is terminated/truncated (for all agents).
# Case 1: all agents are terminated or all are truncated.
self.is_terminated = terminateds.get("__all__", False)
self.is_truncated = truncateds.get("__all__", False)
# Find all agents that were done at prior timesteps and add the agents that are
# done at the present timestep.
agents_done = set(
[aid for aid, sa_eps in self.agent_episodes.items() if sa_eps.is_done]
+ [aid for aid in terminateds if terminateds[aid]]
+ [aid for aid in truncateds if truncateds[aid]]
)
# Case 2: Some agents are truncated and the others are terminated -> Declare
# this episode as terminated.
if all(aid in set(agents_done) for aid in self.agent_ids):
self.is_terminated = True
# For all agents that are not stepping in this env step, but that are not done
# yet -> Add a skip tag to their env- to agent-step mappings.
stepped_agent_ids = set(observations.keys())
for agent_id, env_t_to_agent_t in self.env_t_to_agent_t.items():
if agent_id not in stepped_agent_ids:
env_t_to_agent_t.append(self.SKIP_ENV_TS_TAG)
# Loop through all agent IDs that we received data for in this step:
# Those found in observations, actions, and rewards.
agent_ids_with_data = (
set(observations.keys())
| set(actions.keys())
| set(rewards.keys())
| set(terminateds.keys())
| set(truncateds.keys())
| set(
self.agent_episodes.keys()
if terminateds.get("__all__") or truncateds.get("__all__")
else set()
)
) - {"__all__"}
for agent_id in agent_ids_with_data:
if agent_id not in self.agent_episodes:
sa_episode = SingleAgentEpisode(
agent_id=agent_id,
module_id=self.module_for(agent_id),
multi_agent_episode_id=self.id_,
observation_space=self.observation_space.get(agent_id),
action_space=self.action_space.get(agent_id),
)
else:
sa_episode = self.agent_episodes[agent_id]
# Collect value to be passed (at end of for-loop) into `add_env_step()`
# call.
_observation = observations.get(agent_id)
_action = actions.get(agent_id)
_reward = rewards.get(agent_id)
_infos = infos.get(agent_id)
_terminated = terminateds.get(agent_id, False) or self.is_terminated
_truncated = truncateds.get(agent_id, False) or self.is_truncated
_extra_model_outputs = extra_model_outputs.get(agent_id)
# The value to place into the env- to agent-step map for this agent ID.
# _agent_step = self.SKIP_ENV_TS_TAG
# Agents, whose SingleAgentEpisode had already been done before this
# step should NOT have received any data in this step.
if sa_episode.is_done and any(
v is not None
for v in [_observation, _action, _reward, _infos, _extra_model_outputs]
):
raise MultiAgentEnvError(
f"Agent {agent_id} already had its `SingleAgentEpisode.is_done` "
f"set to True, but still received data in a following step! "
f"obs={_observation} act={_action} rew={_reward} info={_infos} "
f"extra_model_outputs={_extra_model_outputs}."
)
_reward = _reward or 0.0
# CASE 1: A complete agent step is available (in one env step).
# -------------------------------------------------------------
# We have an observation and an action for this agent ->
# Add the agent step to the single agent episode.
# ... action -> next obs + reward ...
if _observation is not None and _action is not None:
if agent_id not in rewards:
raise MultiAgentEnvError(
f"Agent {agent_id} acted (and received next obs), but did NOT "
f"receive any reward from the env!"
)
# CASE 2: Step gets completed with a hanging action OR first observation.
# ------------------------------------------------------------------------
# We have an observation, but no action ->
# a) Action (and extra model outputs) must be hanging already. Also use
# collected hanging rewards and extra_model_outputs.
# b) The observation is the first observation for this agent ID.
elif _observation is not None and _action is None:
_action = self._hanging_actions_end.pop(agent_id, None)
# We have a hanging action (the agent had acted after the previous
# observation, but the env had not responded - until now - with another
# observation).
# ...[hanging action] ... ... -> next obs + (reward)? ...
if _action is not None:
# Get the extra model output if available.
_extra_model_outputs = self._hanging_extra_model_outputs_end.pop(
agent_id, None
)
_reward = self._hanging_rewards_end.pop(agent_id, 0.0) + _reward
# First observation for this agent, we have no hanging action.
# ... [done]? ... -> [1st obs for agent ID]
else:
# The agent is already done -> The agent thus has never stepped once
# and we do not have to create a SingleAgentEpisode for it.
if _terminated or _truncated:
self._del_hanging(agent_id)
continue
# This must be the agent's initial observation.
else:
# Prepend n skip tags to this agent's mapping + the initial [0].
assert agent_id not in self.env_t_to_agent_t
self.env_t_to_agent_t[agent_id].extend(
[self.SKIP_ENV_TS_TAG] * self.env_t + [0]
)
self.env_t_to_agent_t[
agent_id
].lookback = self._len_lookback_buffers
# Make `add_env_reset` call and continue with next agent.
sa_episode.add_env_reset(observation=_observation, infos=_infos)
# Add possible reward to begin cache.
self._hanging_rewards_begin[agent_id] += _reward
# Now that the SAEps is valid, add it to our dict.
self.agent_episodes[agent_id] = sa_episode
continue
# CASE 3: Step is started (by an action), but not completed (no next obs).
# ------------------------------------------------------------------------
# We have no observation, but we have a hanging action (used when we receive
# the next obs for this agent in the future).
elif agent_id not in observations and agent_id in actions:
# Agent got truncated -> Error b/c we would need a last (truncation)
# observation for this (otherwise, e.g. bootstrapping would not work).
# [previous obs] [action] (hanging) ... ... [truncated]
if _truncated:
raise MultiAgentEnvError(
f"Agent {agent_id} acted and then got truncated, but did NOT "
"receive a last (truncation) observation, required for e.g. "
"value function bootstrapping!"
)
# Agent got terminated.
# [previous obs] [action] (hanging) ... ... [terminated]
elif _terminated:
# If the agent was terminated and no observation is provided,
# duplicate the previous one (this is a technical "fix" to properly
# complete the single agent episode; this last observation is never
# used for learning anyway).
_observation = sa_episode._last_added_observation
_infos = sa_episode._last_added_infos
# Agent is still alive.
# [previous obs] [action] (hanging) ...
else:
# Hanging action, reward, and extra_model_outputs.
assert agent_id not in self._hanging_actions_end
self._hanging_actions_end[agent_id] = _action
self._hanging_rewards_end[agent_id] = _reward
self._hanging_extra_model_outputs_end[
agent_id
] = _extra_model_outputs
# CASE 4: Step has started in the past and is still ongoing (no observation,
# no action).
# --------------------------------------------------------------------------
# Record reward and terminated/truncated flags.
else:
_action = self._hanging_actions_end.get(agent_id)
# Agent is done.
if _terminated or _truncated:
# If the agent has NOT stepped, we treat it as not being
# part of this episode.
# ... ... [other agents doing stuff] ... ... [agent done]
if _action is None:
self._del_hanging(agent_id)
continue
# Agent got truncated -> Error b/c we would need a last (truncation)
# observation for this (otherwise, e.g. bootstrapping would not
# work).
if _truncated:
raise MultiAgentEnvError(
f"Agent {agent_id} acted and then got truncated, but did "
"NOT receive a last (truncation) observation, required "
"for e.g. value function bootstrapping!"
)
# [obs] ... ... [hanging action] ... ... [done]
# If the agent was terminated and no observation is provided,
# duplicate the previous one (this is a technical "fix" to properly
# complete the single agent episode; this last observation is never
# used for learning anyway).
_observation = sa_episode._last_added_observation
_infos = sa_episode._last_added_infos
# `_action` is already `get` above. We don't need to pop out from
# the cache as it gets wiped out anyway below b/c the agent is
# done.
_extra_model_outputs = self._hanging_extra_model_outputs_end.pop(
agent_id, None
)
_reward = self._hanging_rewards_end.pop(agent_id, 0.0) + _reward
# The agent is still alive, just add current reward to cache.
else:
# But has never stepped in this episode -> add to begin cache.
if agent_id not in self.agent_episodes:
self._hanging_rewards_begin[agent_id] += _reward
# Otherwise, add to end cache.
else:
self._hanging_rewards_end[agent_id] += _reward
# If agent is stepping, add timestep to `SingleAgentEpisode`.
if _observation is not None:
sa_episode.add_env_step(
observation=_observation,
action=_action,
reward=_reward,
infos=_infos,
terminated=_terminated,
truncated=_truncated,
extra_model_outputs=_extra_model_outputs,
)
# Update the env- to agent-step mapping.
self.env_t_to_agent_t[agent_id].append(
len(sa_episode) + self.agent_t_started[agent_id]
)
# Agent is also done. -> Erase all hanging values for this agent
# (they should be empty at this point anyways).
if _terminated or _truncated:
self._del_hanging(agent_id)
# Validate our data.
self.validate()
# Step time stats.
self._last_step_time = time.perf_counter()
if self._start_time is None:
self._start_time = self._last_step_time
def validate(self) -> None:
"""Validates the episode's data.
This function ensures that the data stored to a `MultiAgentEpisode` is
in order (e.g. that the correct number of observations, actions, rewards
are there).
"""
for eps in self.agent_episodes.values():
eps.validate()
# TODO (sven): Validate MultiAgentEpisode specifics, like the timestep mappings,
# action/reward caches, etc..
@property
def custom_data(self):
return self._custom_data
@property
def is_reset(self) -> bool:
"""Returns True if `self.add_env_reset()` has already been called."""
return any(
len(sa_episode.observations) > 0
for sa_episode in self.agent_episodes.values()
)
@property
def is_numpy(self) -> bool:
"""True, if the data in this episode is already stored as numpy arrays."""
is_numpy = next(iter(self.agent_episodes.values())).is_numpy
# Make sure that all single agent's episodes' `is_numpy` flags are the same.
if not all(eps.is_numpy is is_numpy for eps in self.agent_episodes.values()):
raise RuntimeError(
f"Only some SingleAgentEpisode objects in {self} are converted to "
f"numpy, others are not!"
)
return is_numpy
@property
def is_done(self):
"""Whether the episode is actually done (terminated or truncated).
A done episode cannot be continued via `self.add_env_step()` or being
concatenated on its right-side with another episode chunk or being
succeeded via `self.cut()`.
Note that in a multi-agent environment this does not necessarily
correspond to single agents having terminated or being truncated.
`self.is_terminated` should be `True`, if all agents are terminated and
`self.is_truncated` should be `True`, if all agents are truncated. If
only one or more (but not all!) agents are `terminated/truncated the
`MultiAgentEpisode.is_terminated/is_truncated` should be `False`. This
information about single agent's terminated/truncated states can always
be retrieved from the `SingleAgentEpisode`s inside the 'MultiAgentEpisode`
one.
If all agents are either terminated or truncated, but in a mixed fashion,
i.e. some are terminated and others are truncated: This is currently
undefined and could potentially be a problem (if a user really implemented
such a multi-agent env that behaves this way).
Returns:
Boolean defining if an episode has either terminated or truncated.
"""
return self.is_terminated or self.is_truncated
def to_numpy(self) -> "MultiAgentEpisode":
"""Converts this Episode's list attributes to numpy arrays.
This means in particular that this episodes' lists (per single agent) of
(possibly complex) data (e.g. an agent having a dict obs space) will be
converted to (possibly complex) structs, whose leafs are now numpy arrays.
Each of these leaf numpy arrays will have the same length (batch dimension)
as the length of the original lists.
Note that Columns.INFOS are NEVER numpy'ized and will remain a list
(normally, a list of the original, env-returned dicts). This is due to the
heterogeneous nature of INFOS returned by envs, which would make it unwieldy to
convert this information to numpy arrays.
After calling this method, no further data may be added to this episode via
the `self.add_env_step()` method.
Examples:
.. testcode::
import numpy as np
from ray.rllib.env.multi_agent_episode import MultiAgentEpisode
from ray.rllib.env.tests.test_multi_agent_episode import (
TestMultiAgentEpisode
)
# Create some multi-agent episode data.
(
observations,
actions,
rewards,
terminateds,
truncateds,
infos,
) = TestMultiAgentEpisode._mock_multi_agent_records()
# Define the agent ids.
agent_ids = ["agent_1", "agent_2", "agent_3", "agent_4", "agent_5"]
episode = MultiAgentEpisode(
observations=observations,
infos=infos,
actions=actions,
rewards=rewards,
# Note: terminated/truncated have nothing to do with an episode
# being converted `to_numpy` or not (via the `self.to_numpy()` method)!
terminateds=terminateds,
truncateds=truncateds,
len_lookback_buffer=0, # no lookback; all data is actually "in" episode
)
# Episode has not been numpy'ized yet.
assert not episode.is_numpy
# We are still operating on lists.
assert (
episode.get_observations(
indices=[1],
agent_ids="agent_1",
) == {"agent_1": [1]}
)
# Numpy'ized the episode.
episode.to_numpy()
assert episode.is_numpy
# Everything is now numpy arrays (with 0-axis of size
# B=[len of requested slice]).
assert (
isinstance(episode.get_observations(
indices=[1],
agent_ids="agent_1",
)["agent_1"], np.ndarray)
)
Returns:
This `MultiAgentEpisode` object with the converted numpy data.
"""
for agent_id, agent_eps in self.agent_episodes.copy().items():
agent_eps.to_numpy()
return self
def concat_episode(self, other: "MultiAgentEpisode") -> None:
"""Adds the given `other` MultiAgentEpisode to the right side of `self`.
In order for this to work, both chunks (`self` and `other`) must fit
together that are split through `cut`. For sequential multi-agent environments
using slice might cause problems from hanging observation/actions.
This is checked by the IDs (must be identical), the time step counters
(`self.env_t` must be the same as `other.env_t_started`), as well as the
observations/infos of the individual agents at the concatenation boundaries.
Also, `self.is_done` must not be True, meaning `self.is_terminated` and
`self.is_truncated` are both False.
Args:
other: The other `MultiAgentEpisode` to be concatenated to this one.
Returns:
A `MultiAgentEpisode` instance containing the concatenated data
from both episodes (`self` and `other`).
"""
# Make sure the IDs match.
assert other.id_ == self.id_
# NOTE (sven): This is what we agreed on. As the replay buffers must be
# able to concatenate.
assert not self.is_done
# Make sure the timesteps match.
assert self.env_t == other.env_t_started
# Validate `other`.
other.validate()
# Concatenate the individual SingleAgentEpisodes from both chunks.
all_agent_ids = set(self.agent_ids) | set(other.agent_ids)
for agent_id in all_agent_ids:
sa_episode = self.agent_episodes.get(agent_id)
# If agent is only in the new episode chunk -> Store all the data of `other`
# wrt agent in `self`.
if sa_episode is None:
self.agent_episodes[agent_id] = other.agent_episodes[agent_id]
self.agent_t_started[agent_id] = other.agent_t_started[agent_id]
# If agent only has the first reset observation then no episode exists but `env_t_to_agent_t` does
if agent_id not in self.env_t_to_agent_t:
self.env_t_to_agent_t[agent_id] = other.env_t_to_agent_t[agent_id]
else:
# For a cut episode, the first timestep is a copy of the last timestep from the previous episode
for val in other.env_t_to_agent_t[agent_id][1:]:
self.env_t_to_agent_t[agent_id].append(val)
self._copy_hanging(agent_id, other)
# If the agent was done in `self`, ignore and continue. There should not be
# any data of that agent in `other`.
elif sa_episode.is_done:
continue
# If the agent has data in both chunks, concatenate on the single-agent
# level, thereby making sure the hanging values (begin and end) match.
elif agent_id in other.agent_episodes:
sa_episode.concat_episode(other.agent_episodes[agent_id])
# Override `self`'s hanging (end) values with `other`'s hanging (end).
if agent_id in other._hanging_actions_end:
self._hanging_actions_end[agent_id] = copy.deepcopy(
other._hanging_actions_end[agent_id]
)
self._hanging_rewards_end[agent_id] = other._hanging_rewards_end[
agent_id
]
self._hanging_extra_model_outputs_end[agent_id] = copy.deepcopy(
other._hanging_extra_model_outputs_end[agent_id]
)
# Concatenate the env- to agent-timestep mappings.
# Skip the first element (overlapping boundary) and append the rest.
# Values are agent timesteps, so append them directly.
for val in other.env_t_to_agent_t[agent_id][1:]:
self.env_t_to_agent_t[agent_id].append(val)
# Otherwise, the agent is only in `self` and not done. All data is stored
# already -> skip
# else: pass
# Update all timestep counters.
self.env_t = other.env_t
# Check, if the episode is terminated or truncated.
if other.is_terminated:
self.is_terminated = True
elif other.is_truncated:
self.is_truncated = True
# Merge with `other`'s custom_data, but give `other` priority b/c we assume
# that as a follow-up chunk of `self` other has a more complete version of
# `custom_data`.
self.custom_data.update(other.custom_data)
# Validate.
self.validate()
def cut(self, len_lookback_buffer: int = 0) -> "MultiAgentEpisode":
"""Returns a successor episode chunk (of len=0) continuing from this Episode.
The successor will have the same ID as `self`.
If no lookback buffer is requested (len_lookback_buffer=0), the successor's
observations will be the last observation(s) of `self` and its length will
therefore be 0 (no further steps taken yet). If `len_lookback_buffer` > 0,
the returned successor will have `len_lookback_buffer` observations (and
actions, rewards, etc..) taken from the right side (end) of `self`. For example
if `len_lookback_buffer=2`, the returned successor's lookback buffer actions
will be identical to the results of `self.get_actions([-2, -1])`.
This method is useful if you would like to discontinue building an episode
chunk (b/c you have to return it from somewhere), but would like to have a new
episode instance to continue building the actual gym.Env episode at a later
time. Vie the `len_lookback_buffer` argument, the continuing chunk (successor)
will still be able to "look back" into this predecessor episode's data (at
least to some extend, depending on the value of `len_lookback_buffer`).
Args:
len_lookback_buffer: The number of environment timesteps to take along into
the new chunk as "lookback buffer". A lookback buffer is additional data
on the left side of the actual episode data for visibility purposes
(but without actually being part of the new chunk). For example, if
`self` ends in actions: agent_1=5,6,7 and agent_2=6,7, and we call
`self.cut(len_lookback_buffer=2)`, the returned chunk will have
actions 6 and 7 for both agents already in it, but still
`t_started`==t==8 (not 7!) and a length of 0. If there is not enough
data in `self` yet to fulfil the `len_lookback_buffer` request, the
value of `len_lookback_buffer` is automatically adjusted (lowered).
Returns:
The successor Episode chunk of this one with the same ID and state and the
only observation being the last observation in self.
"""
assert len_lookback_buffer >= 0
if self.is_done:
raise RuntimeError(
"Can't call `MultiAgentEpisode.cut()` when the episode is already done!"
)
# If there is hanging data (e.g. actions) in the agents' caches, we might have
# to re-adjust the lookback len further into the past to make sure that these
# agents have at least one observation to look back to. Otherwise, the timestep
# that got cut into will be "lost" for learning from it.
orig_len_lb = len_lookback_buffer
for agent_id, agent_actions in self._hanging_actions_end.items():
assert self.env_t_to_agent_t[agent_id].get(-1) == self.SKIP_ENV_TS_TAG
for i in range(orig_len_lb, len(self.env_t_to_agent_t[agent_id].data) + 1):
if self.env_t_to_agent_t[agent_id].get(-i) != self.SKIP_ENV_TS_TAG:
len_lookback_buffer = max(len_lookback_buffer, i - 1)
break
# Initialize this episode chunk with the most recent observations
# and infos (even if lookback is zero). Similar to an initial `env.reset()`
indices_obs_and_infos = slice(-len_lookback_buffer - 1, None)
indices_rest = (
slice(-len_lookback_buffer, None)
if len_lookback_buffer > 0
else slice(None, 0) # -> empty slice
)
observations = self.get_observations(
indices=indices_obs_and_infos, return_list=True
)
infos = self.get_infos(indices=indices_obs_and_infos, return_list=True)
actions = self.get_actions(indices=indices_rest, return_list=True)
rewards = self.get_rewards(indices=indices_rest, return_list=True)
extra_model_outputs = self.get_extra_model_outputs(
key=None, # all keys
indices=indices_rest,
return_list=True,
)
successor = MultiAgentEpisode(
# Same ID.
id_=self.id_,
observations=observations,
observation_space=self.observation_space,
infos=infos,
actions=actions,
action_space=self.action_space,
rewards=rewards,
# List of MADicts, mapping agent IDs to their respective extra model output
# dicts.
extra_model_outputs=extra_model_outputs,
terminateds=self.get_terminateds(),
truncateds=self.get_truncateds(),
# Continue with `self`'s current timesteps.
env_t_started=self.env_t,
agent_t_started={
aid: self.agent_episodes[aid].t
for aid in self.agent_ids
if not self.agent_episodes[aid].is_done
},
# Same AgentIDs and SingleAgentEpisode IDs.
agent_episode_ids=self.agent_episode_ids,
agent_module_ids={
aid: self.agent_episodes[aid].module_id for aid in self.agent_ids
},
agent_to_module_mapping_fn=self.agent_to_module_mapping_fn,
# All data we provided to the c'tor goes into the lookback buffer.
len_lookback_buffer="auto",
)
# Copy over the hanging (end) values into the hanging (begin) caches of the
# successor.
successor._hanging_rewards_begin = self._hanging_rewards_end.copy()
# Deepcopy all custom data in `self` to be continued in the cut episode.
successor._custom_data = copy.deepcopy(self.custom_data)
return successor
@property
def agent_ids(self) -> Set[AgentID]:
"""Returns the agent ids."""
return set(self.agent_episodes.keys())
@property
def agent_episode_ids(self) -> MultiAgentDict:
"""Returns ids from each agent's `SingleAgentEpisode`."""
return {
agent_id: agent_eps.id_
for agent_id, agent_eps in self.agent_episodes.items()
}
def module_for(self, agent_id: AgentID) -> Optional[ModuleID]:
"""Returns the ModuleID for a given AgentID.
Forces the agent-to-module mapping to be performed (via
`self.agent_to_module_mapping_fn`), if this has not been done yet.
Note that all such mappings are stored in the `self._agent_to_module_mapping`
property.
Args:
agent_id: The AgentID to get a mapped ModuleID for.
Returns:
The ModuleID mapped to from the given `agent_id`.
"""
if agent_id not in self._agent_to_module_mapping:
module_id = self._agent_to_module_mapping[
agent_id
] = self.agent_to_module_mapping_fn(agent_id, self)
return module_id
else:
return self._agent_to_module_mapping[agent_id]
def get_observations(
self,
indices: Optional[Union[int, List[int], slice]] = None,
agent_ids: Optional[Union[Collection[AgentID], AgentID]] = None,
*,
env_steps: bool = True,
# global_indices: bool = False,
neg_index_as_lookback: bool = False,
fill: Optional[Any] = None,
one_hot_discrete: bool = False,
return_list: bool = False,
) -> Union[MultiAgentDict, List[MultiAgentDict]]:
"""Returns agents' observations or batched ranges thereof from this episode.
Args:
indices: A single int is interpreted as an index, from which to return the
individual observation stored at this index.
A list of ints is interpreted as a list of indices from which to gather
individual observations in a batch of size len(indices).
A slice object is interpreted as a range of observations to be returned.
Thereby, negative indices by default are interpreted as "before the end"
unless the `neg_index_as_lookback=True` option is used, in which case
negative indices are interpreted as "before ts=0", meaning going back
into the lookback buffer.
If None, will return all observations (from ts=0 to the end).
agent_ids: An optional collection of AgentIDs or a single AgentID to get
observations for. If None, will return observations for all agents in
this episode.
env_steps: Whether `indices` should be interpreted as environment time steps
(True) or per-agent timesteps (False).
neg_index_as_lookback: If True, negative values in `indices` are
interpreted as "before ts=0", meaning going back into the lookback
buffer. For example, an episode with agent A's observations
[4, 5, 6, 7, 8, 9], where [4, 5, 6] is the lookback buffer range
(ts=0 item is 7), will respond to `get_observations(-1, agent_ids=[A],
neg_index_as_lookback=True)` with {A: `6`} and to
`get_observations(slice(-2, 1), agent_ids=[A],
neg_index_as_lookback=True)` with {A: `[5, 6, 7]`}.
fill: An optional value to use for filling up the returned results at
the boundaries. This filling only happens if the requested index range's
start/stop boundaries exceed the episode's boundaries (including the
lookback buffer on the left side). This comes in very handy, if users
don't want to worry about reaching such boundaries and want to zero-pad.
For example, an episode with agent A' observations [10, 11, 12, 13, 14]
and lookback buffer size of 2 (meaning observations `10` and `11` are
part of the lookback buffer) will respond to
`get_observations(slice(-7, -2), agent_ids=[A], fill=0.0)` with
`{A: [0.0, 0.0, 10, 11, 12]}`.
one_hot_discrete: If True, will return one-hot vectors (instead of
int-values) for those sub-components of a (possibly complex) observation
space that are Discrete or MultiDiscrete. Note that if `fill=0` and the
requested `indices` are out of the range of our data, the returned
one-hot vectors will actually be zero-hot (all slots zero).
return_list: Whether to return a list of multi-agent dicts (instead of
a single multi-agent dict of lists/structs). False by default. This
option can only be used when `env_steps` is True due to the fact the
such a list can only be interpreted as one env step per list item
(would not work with agent steps).
Returns:
A dictionary mapping agent IDs to observations (at the given
`indices`). If `env_steps` is True, only agents that have stepped
(were ready) at the given env step `indices` are returned (i.e. not all
agent IDs are necessarily in the keys).
If `return_list` is True, returns a list of MultiAgentDicts (mapping agent
IDs to observations) instead.
"""
return self._get(
what="observations",
indices=indices,
agent_ids=agent_ids,
env_steps=env_steps,
neg_index_as_lookback=neg_index_as_lookback,
fill=fill,
one_hot_discrete=one_hot_discrete,
return_list=return_list,
)
def get_infos(
self,
indices: Optional[Union[int, List[int], slice]] = None,
agent_ids: Optional[Union[Collection[AgentID], AgentID]] = None,
*,
env_steps: bool = True,
neg_index_as_lookback: bool = False,
fill: Optional[Any] = None,
return_list: bool = False,
) -> Union[MultiAgentDict, List[MultiAgentDict]]:
"""Returns agents' info dicts or list (ranges) thereof from this episode.
Args:
indices: A single int is interpreted as an index, from which to return the
individual info dict stored at this index.
A list of ints is interpreted as a list of indices from which to gather
individual info dicts in a list of size len(indices).
A slice object is interpreted as a range of info dicts to be returned.
Thereby, negative indices by default are interpreted as "before the end"
unless the `neg_index_as_lookback=True` option is used, in which case
negative indices are interpreted as "before ts=0", meaning going back
into the lookback buffer.
If None, will return all infos (from ts=0 to the end).
agent_ids: An optional collection of AgentIDs or a single AgentID to get
info dicts for. If None, will return info dicts for all agents in
this episode.
env_steps: Whether `indices` should be interpreted as environment time steps
(True) or per-agent timesteps (False).
neg_index_as_lookback: If True, negative values in `indices` are
interpreted as "before ts=0", meaning going back into the lookback
buffer. For example, an episode with agent A's info dicts
[{"l":4}, {"l":5}, {"l":6}, {"a":7}, {"b":8}, {"c":9}], where the
first 3 items are the lookback buffer (ts=0 item is {"a": 7}), will
respond to `get_infos(-1, agent_ids=A, neg_index_as_lookback=True)`
with `{A: {"l":6}}` and to
`get_infos(slice(-2, 1), agent_ids=A, neg_index_as_lookback=True)`
with `{A: [{"l":5}, {"l":6}, {"a":7}]}`.
fill: An optional value to use for filling up the returned results at
the boundaries. This filling only happens if the requested index range's
start/stop boundaries exceed the episode's boundaries (including the
lookback buffer on the left side). This comes in very handy, if users
don't want to worry about reaching such boundaries and want to
auto-fill. For example, an episode with agent A's infos being
[{"l":10}, {"l":11}, {"a":12}, {"b":13}, {"c":14}] and lookback buffer
size of 2 (meaning infos {"l":10}, {"l":11} are part of the lookback
buffer) will respond to `get_infos(slice(-7, -2), agent_ids=A,
fill={"o": 0.0})` with
`{A: [{"o":0.0}, {"o":0.0}, {"l":10}, {"l":11}, {"a":12}]}`.
return_list: Whether to return a list of multi-agent dicts (instead of
a single multi-agent dict of lists/structs). False by default. This
option can only be used when `env_steps` is True due to the fact the
such a list can only be interpreted as one env step per list item
(would not work with agent steps).
Returns:
A dictionary mapping agent IDs to observations (at the given
`indices`). If `env_steps` is True, only agents that have stepped
(were ready) at the given env step `indices` are returned (i.e. not all
agent IDs are necessarily in the keys).
If `return_list` is True, returns a list of MultiAgentDicts (mapping agent
IDs to infos) instead.
"""
return self._get(
what="infos",
indices=indices,
agent_ids=agent_ids,
env_steps=env_steps,
neg_index_as_lookback=neg_index_as_lookback,
fill=fill,
return_list=return_list,
)
def get_actions(
self,
indices: Optional[Union[int, List[int], slice]] = None,
agent_ids: Optional[Union[Collection[AgentID], AgentID]] = None,
*,
env_steps: bool = True,
neg_index_as_lookback: bool = False,
fill: Optional[Any] = None,
one_hot_discrete: bool = False,
return_list: bool = False,
) -> Union[MultiAgentDict, List[MultiAgentDict]]:
"""Returns agents' actions or batched ranges thereof from this episode.
Args:
indices: A single int is interpreted as an index, from which to return the
individual actions stored at this index.
A list of ints is interpreted as a list of indices from which to gather
individual actions in a batch of size len(indices).
A slice object is interpreted as a range of actions to be returned.
Thereby, negative indices by default are interpreted as "before the end"
unless the `neg_index_as_lookback=True` option is used, in which case
negative indices are interpreted as "before ts=0", meaning going back
into the lookback buffer.
If None, will return all actions (from ts=0 to the end).
agent_ids: An optional collection of AgentIDs or a single AgentID to get
actions for. If None, will return actions for all agents in
this episode.
env_steps: Whether `indices` should be interpreted as environment time steps
(True) or per-agent timesteps (False).
neg_index_as_lookback: If True, negative values in `indices` are
interpreted as "before ts=0", meaning going back into the lookback
buffer. For example, an episode with agent A's actions
[4, 5, 6, 7, 8, 9], where [4, 5, 6] is the lookback buffer range
(ts=0 item is 7), will respond to `get_actions(-1, agent_ids=[A],
neg_index_as_lookback=True)` with {A: `6`} and to
`get_actions(slice(-2, 1), agent_ids=[A],
neg_index_as_lookback=True)` with {A: `[5, 6, 7]`}.
fill: An optional value to use for filling up the returned results at
the boundaries. This filling only happens if the requested index range's
start/stop boundaries exceed the episode's boundaries (including the
lookback buffer on the left side). This comes in very handy, if users
don't want to worry about reaching such boundaries and want to zero-pad.
For example, an episode with agent A' actions [10, 11, 12, 13, 14]
and lookback buffer size of 2 (meaning actions `10` and `11` are
part of the lookback buffer) will respond to
`get_actions(slice(-7, -2), agent_ids=[A], fill=0.0)` with
`{A: [0.0, 0.0, 10, 11, 12]}`.
one_hot_discrete: If True, will return one-hot vectors (instead of
int-values) for those sub-components of a (possibly complex) observation
space that are Discrete or MultiDiscrete. Note that if `fill=0` and the
requested `indices` are out of the range of our data, the returned
one-hot vectors will actually be zero-hot (all slots zero).
return_list: Whether to return a list of multi-agent dicts (instead of
a single multi-agent dict of lists/structs). False by default. This
option can only be used when `env_steps` is True due to the fact the
such a list can only be interpreted as one env step per list item
(would not work with agent steps).
Returns:
A dictionary mapping agent IDs to actions (at the given
`indices`). If `env_steps` is True, only agents that have stepped
(were ready) at the given env step `indices` are returned (i.e. not all
agent IDs are necessarily in the keys).
If `return_list` is True, returns a list of MultiAgentDicts (mapping agent
IDs to actions) instead.
"""
return self._get(
what="actions",
indices=indices,
agent_ids=agent_ids,
env_steps=env_steps,
neg_index_as_lookback=neg_index_as_lookback,
fill=fill,
one_hot_discrete=one_hot_discrete,
return_list=return_list,
)
def get_rewards(
self,
indices: Optional[Union[int, List[int], slice]] = None,
agent_ids: Optional[Union[Collection[AgentID], AgentID]] = None,
*,
env_steps: bool = True,
neg_index_as_lookback: bool = False,
fill: Optional[float] = None,
return_list: bool = False,
) -> Union[MultiAgentDict, List[MultiAgentDict]]:
"""Returns agents' rewards or batched ranges thereof from this episode.
Args:
indices: A single int is interpreted as an index, from which to return the
individual rewards stored at this index.
A list of ints is interpreted as a list of indices from which to gather
individual rewards in a batch of size len(indices).
A slice object is interpreted as a range of rewards to be returned.
Thereby, negative indices by default are interpreted as "before the end"
unless the `neg_index_as_lookback=True` option is used, in which case
negative indices are interpreted as "before ts=0", meaning going back
into the lookback buffer.
If None, will return all rewards (from ts=0 to the end).
agent_ids: An optional collection of AgentIDs or a single AgentID to get
rewards for. If None, will return rewards for all agents in
this episode.
env_steps: Whether `indices` should be interpreted as environment time steps
(True) or per-agent timesteps (False).
neg_index_as_lookback: If True, negative values in `indices` are
interpreted as "before ts=0", meaning going back into the lookback
buffer. For example, an episode with agent A's rewards
[4, 5, 6, 7, 8, 9], where [4, 5, 6] is the lookback buffer range
(ts=0 item is 7), will respond to `get_rewards(-1, agent_ids=[A],
neg_index_as_lookback=True)` with {A: `6`} and to
`get_rewards(slice(-2, 1), agent_ids=[A],
neg_index_as_lookback=True)` with {A: `[5, 6, 7]`}.
fill: An optional float value to use for filling up the returned results at
the boundaries. This filling only happens if the requested index range's
start/stop boundaries exceed the episode's boundaries (including the
lookback buffer on the left side). This comes in very handy, if users
don't want to worry about reaching such boundaries and want to zero-pad.
For example, an episode with agent A' rewards [10, 11, 12, 13, 14]
and lookback buffer size of 2 (meaning rewards `10` and `11` are
part of the lookback buffer) will respond to
`get_rewards(slice(-7, -2), agent_ids=[A], fill=0.0)` with
`{A: [0.0, 0.0, 10, 11, 12]}`.
return_list: Whether to return a list of multi-agent dicts (instead of
a single multi-agent dict of lists/structs). False by default. This
option can only be used when `env_steps` is True due to the fact the
such a list can only be interpreted as one env step per list item
(would not work with agent steps).
Returns:
A dictionary mapping agent IDs to rewards (at the given
`indices`). If `env_steps` is True, only agents that have stepped
(were ready) at the given env step `indices` are returned (i.e. not all
agent IDs are necessarily in the keys).
If `return_list` is True, returns a list of MultiAgentDicts (mapping agent
IDs to rewards) instead.
"""
return self._get(
what="rewards",
indices=indices,
agent_ids=agent_ids,
env_steps=env_steps,
neg_index_as_lookback=neg_index_as_lookback,
fill=fill,
return_list=return_list,
)
def get_extra_model_outputs(
self,
key: Optional[str] = None,
indices: Optional[Union[int, List[int], slice]] = None,
agent_ids: Optional[Union[Collection[AgentID], AgentID]] = None,
*,
env_steps: bool = True,
neg_index_as_lookback: bool = False,
fill: Optional[Any] = None,
return_list: bool = False,
) -> Union[MultiAgentDict, List[MultiAgentDict]]:
"""Returns agents' actions or batched ranges thereof from this episode.
Args:
key: The `key` within each agents' extra_model_outputs dict to extract
data for. If None, return data of all extra model output keys.
indices: A single int is interpreted as an index, from which to return the
individual extra model outputs stored at this index.
A list of ints is interpreted as a list of indices from which to gather
individual extra model outputs in a batch of size len(indices).
A slice object is interpreted as a range of extra model outputs to be
returned.
Thereby, negative indices by default are interpreted as "before the end"
unless the `neg_index_as_lookback=True` option is used, in which case
negative indices are interpreted as "before ts=0", meaning going back
into the lookback buffer.
If None, will return all extra model outputs (from ts=0 to the end).
agent_ids: An optional collection of AgentIDs or a single AgentID to get
extra model outputs for. If None, will return extra model outputs for
all agents in this episode.
env_steps: Whether `indices` should be interpreted as environment time steps
(True) or per-agent timesteps (False).
neg_index_as_lookback: If True, negative values in `indices` are
interpreted as "before ts=0", meaning going back into the lookback
buffer. For example, an episode with agent A's actions
[4, 5, 6, 7, 8, 9], where [4, 5, 6] is the lookback buffer range
(ts=0 item is 7), will respond to `get_actions(-1, agent_ids=[A],
neg_index_as_lookback=True)` with {A: `6`} and to
`get_actions(slice(-2, 1), agent_ids=[A],
neg_index_as_lookback=True)` with {A: `[5, 6, 7]`}.
fill: An optional value to use for filling up the returned results at
the boundaries. This filling only happens if the requested index range's
start/stop boundaries exceed the episode's boundaries (including the
lookback buffer on the left side). This comes in very handy, if users
don't want to worry about reaching such boundaries and want to zero-pad.
For example, an episode with agent A' actions [10, 11, 12, 13, 14]
and lookback buffer size of 2 (meaning actions `10` and `11` are
part of the lookback buffer) will respond to
`get_actions(slice(-7, -2), agent_ids=[A], fill=0.0)` with
`{A: [0.0, 0.0, 10, 11, 12]}`.
one_hot_discrete: If True, will return one-hot vectors (instead of
int-values) for those sub-components of a (possibly complex) observation
space that are Discrete or MultiDiscrete. Note that if `fill=0` and the
requested `indices` are out of the range of our data, the returned
one-hot vectors will actually be zero-hot (all slots zero).
return_list: Whether to return a list of multi-agent dicts (instead of
a single multi-agent dict of lists/structs). False by default. This
option can only be used when `env_steps` is True due to the fact the
such a list can only be interpreted as one env step per list item
(would not work with agent steps).
Returns:
A dictionary mapping agent IDs to actions (at the given
`indices`). If `env_steps` is True, only agents that have stepped
(were ready) at the given env step `indices` are returned (i.e. not all
agent IDs are necessarily in the keys).
If `return_list` is True, returns a list of MultiAgentDicts (mapping agent
IDs to extra_model_outputs) instead.
"""
return self._get(
what="extra_model_outputs",
extra_model_outputs_key=key,
indices=indices,
agent_ids=agent_ids,
env_steps=env_steps,
neg_index_as_lookback=neg_index_as_lookback,
fill=fill,
return_list=return_list,
)
def get_terminateds(self) -> MultiAgentDict:
"""Gets the terminateds at given indices."""
terminateds = {
agent_id: self.agent_episodes[agent_id].is_terminated
for agent_id in self.agent_ids
}
terminateds.update({"__all__": self.is_terminated})
return terminateds
def get_truncateds(self) -> MultiAgentDict:
truncateds = {
agent_id: self.agent_episodes[agent_id].is_truncated
for agent_id in self.agent_ids
}
truncateds.update({"__all__": self.is_terminated})
return truncateds
def slice(
self,
slice_: slice,
*,
len_lookback_buffer: Optional[int] = None,
) -> "MultiAgentEpisode":
"""Returns a slice of this episode with the given slice object.
Works analogous to
:py:meth:`~ray.rllib.env.single_agent_episode.SingleAgentEpisode.slice`
However, the important differences are:
- `slice_` is provided in (global) env steps, not agent steps.
- In case `slice_` ends - for a certain agent - in an env step, where that
particular agent does not have an observation, the previous observation will
be included, but the next action and sum of rewards until this point will
be stored in the agent's hanging values caches for the returned
MultiAgentEpisode slice.
.. testcode::
from ray.rllib.env.multi_agent_episode import MultiAgentEpisode
from ray.rllib.utils.test_utils import check
# Generate a simple multi-agent episode.
observations = [
{"a0": 0, "a1": 0}, # 0
{ "a1": 1}, # 1
{ "a1": 2}, # 2
{"a0": 3, "a1": 3}, # 3
{"a0": 4}, # 4
]
# Actions are the same as observations (except for last obs, which doesn't
# have an action).
actions = observations[:-1]
# Make up a reward for each action.
rewards = [
{aid: r / 10 + 0.1 for aid, r in o.items()}
for o in observations
]
episode = MultiAgentEpisode(
observations=observations,
actions=actions,
rewards=rewards,
len_lookback_buffer=0,
)
# Slice the episode and check results.
slice = episode[1:3]
a0 = slice.agent_episodes["a0"]
a1 = slice.agent_episodes["a1"]
check((a0.observations, a1.observations), ([3], [1, 2, 3]))
check((a0.actions, a1.actions), ([], [1, 2]))
check((a0.rewards, a1.rewards), ([], [0.2, 0.3]))
check((a0.is_done, a1.is_done), (False, False))
# If a slice ends in a "gap" for an agent, expect actions and rewards to be
# cached for this agent.
slice = episode[:2]
a0 = slice.agent_episodes["a0"]
check(a0.observations, [0])
check(a0.actions, [])
check(a0.rewards, [])
check(slice._hanging_actions_end["a0"], 0)
check(slice._hanging_rewards_end["a0"], 0.1)
Args:
slice_: The slice object to use for slicing. This should exclude the
lookback buffer, which will be prepended automatically to the returned
slice.
len_lookback_buffer: If not None, forces the returned slice to try to have
this number of timesteps in its lookback buffer (if available). If None
(default), tries to make the returned slice's lookback as large as the
current lookback buffer of this episode (`self`).
Returns:
The new MultiAgentEpisode representing the requested slice.
"""
if slice_.step not in [1, None]:
raise NotImplementedError(
"Slicing MultiAgentEnv with a step other than 1 (you used"
f" {slice_.step}) is not supported!"
)
# Translate `slice_` into one that only contains 0-or-positive ints and will
# NOT contain any None.
start = slice_.start
stop = slice_.stop
# Start is None -> 0.
if start is None:
start = 0
# Start is negative -> Interpret index as counting "from end".
elif start < 0:
start = max(len(self) + start, 0)
# Start is larger than len(self) -> Clip to len(self).
elif start > len(self):
start = len(self)
# Stop is None -> Set stop to our len (one ts past last valid index).
if stop is None:
stop = len(self)
# Stop is negative -> Interpret index as counting "from end".
elif stop < 0:
stop = max(len(self) + stop, 0)
# Stop is larger than len(self) -> Clip to len(self).
elif stop > len(self):
stop = len(self)
ref_lookback = None
try:
for aid, sa_episode in self.agent_episodes.items():
if ref_lookback is None:
ref_lookback = sa_episode.observations.lookback
assert sa_episode.observations.lookback == ref_lookback
assert sa_episode.actions.lookback == ref_lookback
assert sa_episode.rewards.lookback == ref_lookback
assert all(
ilb.lookback == ref_lookback
for ilb in sa_episode.extra_model_outputs.values()
)
except AssertionError:
raise ValueError(
"Can only slice a MultiAgentEpisode if all lookback buffers in this "
"episode have the exact same size!"
)
# Determine terminateds/truncateds and when (in agent timesteps) the
# single-agent episode slices start.
terminateds = {}
truncateds = {}
agent_t_started = {}
for aid, sa_episode in self.agent_episodes.items():
mapping = self.env_t_to_agent_t[aid]
# If the (agent) timestep directly at the slice stop boundary is equal to
# the length of the single-agent episode of this agent -> Use the
# single-agent episode's terminated/truncated flags.
# If `stop` is already beyond this agent's single-agent episode, then we
# don't have to keep track of this: The MultiAgentEpisode initializer will
# automatically determine that this agent must be done (b/c it has no action
# following its final observation).
if (
stop < len(mapping)
and mapping[stop] != self.SKIP_ENV_TS_TAG
and len(sa_episode) == mapping[stop]
):
terminateds[aid] = sa_episode.is_terminated
truncateds[aid] = sa_episode.is_truncated
# Determine this agent's t_started.
if start < len(mapping):
for i in range(start, len(mapping)):
if mapping[i] != self.SKIP_ENV_TS_TAG:
agent_t_started[aid] = mapping[i]
break
terminateds["__all__"] = all(
terminateds.get(aid) for aid in self.agent_episodes
)
truncateds["__all__"] = all(truncateds.get(aid) for aid in self.agent_episodes)
# Determine all other slice contents.
_lb = len_lookback_buffer if len_lookback_buffer is not None else ref_lookback
if start - _lb < 0 and ref_lookback < (_lb - start):
_lb = ref_lookback + start
observations = self.get_observations(
slice(start - _lb, stop + 1),
neg_index_as_lookback=True,
return_list=True,
)
actions = self.get_actions(
slice(start - _lb, stop),
neg_index_as_lookback=True,
return_list=True,
)
rewards = self.get_rewards(
slice(start - _lb, stop),
neg_index_as_lookback=True,
return_list=True,
)
extra_model_outputs = self.get_extra_model_outputs(
indices=slice(start - _lb, stop),
neg_index_as_lookback=True,
return_list=True,
)
# Create the actual slice to be returned.
ma_episode = MultiAgentEpisode(
id_=self.id_,
# In the following, offset `start`s automatically by lookbacks.
observations=observations,
observation_space=self.observation_space,
actions=actions,
action_space=self.action_space,
rewards=rewards,
extra_model_outputs=extra_model_outputs,
terminateds=terminateds,
truncateds=truncateds,
len_lookback_buffer=_lb,
env_t_started=self.env_t_started + start,
agent_episode_ids={
aid: eid.id_ for aid, eid in self.agent_episodes.items()
},
agent_t_started=agent_t_started,
agent_module_ids=self._agent_to_module_mapping,
agent_to_module_mapping_fn=self.agent_to_module_mapping_fn,
)
# Numpy'ize slice if `self` is also finalized.
if self.is_numpy:
ma_episode.to_numpy()
return ma_episode
def __len__(self):
"""Returns the length of an `MultiAgentEpisode`.
Note that the length of an episode is defined by the difference
between its actual timestep and the starting point.
Returns: An integer defining the length of the episode or an
error if the episode has not yet started.
"""
return self.env_t - self.env_t_started
def __repr__(self):
sa_eps_returns = {
aid: sa_eps.get_return() for aid, sa_eps in self.agent_episodes.items()
}
return (
f"MAEps(len={len(self)} done={self.is_done} "
f"Rs={sa_eps_returns} id_={self.id_})"
)
def print(self) -> None:
"""Prints this MultiAgentEpisode as a table of observations for the agents."""
# Find the maximum timestep across all agents to determine the grid width.
max_ts = max(ts.len_incl_lookback() for ts in self.env_t_to_agent_t.values())
lookback = next(iter(self.env_t_to_agent_t.values())).lookback
longest_agent = max(len(aid) for aid in self.agent_ids)
# Construct the header.
header = (
"ts"
+ (" " * longest_agent)
+ " ".join(str(i) for i in range(-lookback, max_ts - lookback))
+ "\n"
)
# Construct each agent's row.
rows = []
for agent, inf_buffer in self.env_t_to_agent_t.items():
row = f"{agent} " + (" " * (longest_agent - len(agent)))
for t in inf_buffer.data:
# Two spaces for alignment.
if t == "S":
row += " "
# Mark the step with an x.
else:
row += " x "
# Remove trailing space for alignment.
rows.append(row.rstrip())
# Join all components into a final string
print(header + "\n".join(rows))
def get_state(self) -> Dict[str, Any]:
"""Returns the state of a multi-agent episode.
Note that from an episode's state the episode itself can
be recreated.
Returns: A dicitonary containing pickable data for a
`MultiAgentEpisode`.
"""
return {
"id_": self.id_,
"agent_to_module_mapping_fn": self.agent_to_module_mapping_fn,
"_agent_to_module_mapping": self._agent_to_module_mapping,
"observation_space": self.observation_space,
"action_space": self.action_space,
"env_t_started": self.env_t_started,
"env_t": self.env_t,
"agent_t_started": self.agent_t_started,
# TODO (simon): Check, if we can store the `InfiniteLookbackBuffer`
"env_t_to_agent_t": self.env_t_to_agent_t,
"_hanging_actions_end": self._hanging_actions_end,
"_hanging_extra_model_outputs_end": self._hanging_extra_model_outputs_end,
"_hanging_rewards_end": self._hanging_rewards_end,
"_hanging_rewards_begin": self._hanging_rewards_begin,
"is_terminated": self.is_terminated,
"is_truncated": self.is_truncated,
"agent_episodes": list(
{
agent_id: agent_eps.get_state()
for agent_id, agent_eps in self.agent_episodes.items()
}.items()
),
"_start_time": self._start_time,
"_last_step_time": self._last_step_time,
"custom_data": self.custom_data,
}
@staticmethod
def from_state(state: Dict[str, Any]) -> "MultiAgentEpisode":
"""Creates a multi-agent episode from a state dictionary.
See `MultiAgentEpisode.get_state()` for creating a state for
a `MultiAgentEpisode` pickable state. For recreating a
`MultiAgentEpisode` from a state, this state has to be complete,
i.e. all data must have been stored in the state.
Args:
state: A dict containing all data required to recreate a MultiAgentEpisode`.
See `MultiAgentEpisode.get_state()`.
Returns:
A `MultiAgentEpisode` instance created from the state data.
"""
# Create an empty `MultiAgentEpisode` instance.
episode = MultiAgentEpisode(id_=state["id_"])
# Fill the instance with the state data.
episode.agent_to_module_mapping_fn = state["agent_to_module_mapping_fn"]
episode._agent_to_module_mapping = state["_agent_to_module_mapping"]
episode.observation_space = state["observation_space"]
episode.action_space = state["action_space"]
episode.env_t_started = state["env_t_started"]
episode.env_t = state["env_t"]
episode.agent_t_started = state["agent_t_started"]
episode.env_t_to_agent_t = state["env_t_to_agent_t"]
episode._hanging_actions_end = state["_hanging_actions_end"]
episode._hanging_extra_model_outputs_end = state[
"_hanging_extra_model_outputs_end"
]
episode._hanging_rewards_end = state["_hanging_rewards_end"]
episode._hanging_rewards_begin = state["_hanging_rewards_begin"]
episode.is_terminated = state["is_terminated"]
episode.is_truncated = state["is_truncated"]
episode.agent_episodes = {
agent_id: SingleAgentEpisode.from_state(agent_state)
for agent_id, agent_state in state["agent_episodes"]
}
episode._start_time = state["_start_time"]
episode._last_step_time = state["_last_step_time"]
episode._custom_data = state.get("custom_data", {})
# Validate the episode.
episode.validate()
return episode
def get_sample_batch(self) -> MultiAgentBatch:
"""Converts this `MultiAgentEpisode` into a `MultiAgentBatch`.
Each `SingleAgentEpisode` instances in `MultiAgentEpisode.agent_epiosdes`
will be converted into a `SampleBatch` and the environment timestep will be
passed as the returned MultiAgentBatch's `env_steps`.
Returns:
A MultiAgentBatch containing all of this episode's data.
"""
# TODO (simon): Check, if timesteps should be converted into global
# timesteps instead of agent steps.
# Note, only agents that have stepped are included into the batch.
return MultiAgentBatch(
policy_batches={
agent_id: agent_eps.get_sample_batch()
for agent_id, agent_eps in self.agent_episodes.items()
if agent_eps.t - agent_eps.t_started > 0
},
env_steps=self.env_t - self.env_t_started,
)
def get_return(
self,
include_hanging_rewards: bool = False,
) -> float:
"""Returns all-agent return.
Args:
include_hanging_rewards: Whether we should also consider
hanging rewards wehn calculating the overall return. Agents might
have received partial rewards, i.e. rewards without an
observation. These are stored in the "hanging" caches (begin and end)
for each agent and added up until the next observation is received by
that agent.
Returns:
The sum of all single-agents' returns (maybe including the hanging
rewards per agent).
"""
env_return = sum(
agent_eps.get_return() for agent_eps in self.agent_episodes.values()
)
if include_hanging_rewards:
for hanging_r in self._hanging_rewards_begin.values():
env_return += hanging_r
for hanging_r in self._hanging_rewards_end.values():
env_return += hanging_r
return env_return
def get_agents_to_act(self) -> Set[AgentID]:
"""Returns a set of agent IDs required to send an action to `env.step()` next.
Those are generally the agents that received an observation in the most recent
`env.step()` call.
Returns:
A set of AgentIDs that are supposed to send actions to the next `env.step()`
call.
"""
return {
aid
for aid in self.get_observations(-1).keys()
if not self.agent_episodes[aid].is_done
}
def get_agents_that_stepped(self) -> Set[AgentID]:
"""Returns a set of agent IDs of those agents that just finished stepping.
These are all the agents that have an observation logged at the last env
timestep, which may include agents, whose single agent episode just terminated
or truncated.
Returns:
A set of AgentIDs of those agents that just finished stepping (that have a
most recent observation on the env timestep scale), regardless of whether
their single agent episodes are done or not.
"""
return set(self.get_observations(-1).keys())
def get_duration_s(self) -> float:
"""Returns the duration of this Episode (chunk) in seconds."""
if self._last_step_time is None:
return 0.0
return self._last_step_time - self._start_time
def set_observations(
self,
*,
new_data: MultiAgentDict,
at_indices: Optional[Union[int, List[int], slice]] = None,
neg_index_as_lookback: bool = False,
) -> None:
"""Overwrites all or some single-agent Episode's observations with the provided data.
This is a helper method to batch `SingleAgentEpisode.set_observations`.
For more detail, see `SingleAgentEpisode.set_observations`.
Args:
new_data: A dict mapping agent IDs to new observation data.
Each value in the dict is the new observation data to overwrite existing data with.
This may be a list of individual observation(s) in case this episode
is still not numpy'ized yet. In case this episode has already been
numpy'ized, this should be (possibly complex) struct matching the
observation space and with a batch size of its leafs exactly the size
of the to-be-overwritten slice or segment (provided by `at_indices`).
at_indices: A single int is interpreted as one index, which to overwrite
with `new_data` (which is expected to be a single observation).
A list of ints is interpreted as a list of indices, all of which to
overwrite with `new_data` (which is expected to be of the same size
as `len(at_indices)`).
A slice object is interpreted as a range of indices to be overwritten
with `new_data` (which is expected to be of the same size as the
provided slice).
Thereby, negative indices by default are interpreted as "before the end"
unless the `neg_index_as_lookback=True` option is used, in which case
negative indices are interpreted as "before ts=0", meaning going back
into the lookback buffer.
neg_index_as_lookback: If True, negative values in `at_indices` are
interpreted as "before ts=0", meaning going back into the lookback
buffer. For example, an episode with
observations = [4, 5, 6, 7, 8, 9], where [4, 5, 6] is the
lookback buffer range (ts=0 item is 7), will handle a call to
`set_observations(individual_observation, -1,
neg_index_as_lookback=True)` by overwriting the value of 6 in our
observations buffer with the provided "individual_observation".
Raises:
IndexError: If the provided `at_indices` do not match the size of
`new_data`.
"""
for agent_id, new_agent_data in new_data.items():
if agent_id not in self.agent_episodes:
raise KeyError(f"AgentID '{agent_id}' not found in this episode.")
self.agent_episodes[agent_id].set_observations(
new_data=new_agent_data,
at_indices=at_indices,
neg_index_as_lookback=neg_index_as_lookback,
)
def set_actions(
self,
*,
new_data: MultiAgentDict,
at_indices: Optional[Union[int, List[int], slice]] = None,
neg_index_as_lookback: bool = False,
) -> None:
"""Overwrites all or some of this Episode's actions with the provided data.
This is a helper method to batch `SingleAgentEpisode.set_actions`.
For more detail, see `SingleAgentEpisode.set_actions`.
Args:
new_data: A dict mapping agent IDs to new action data.
Each value in the dict is the new action data to overwrite existing data with.
This may be a list of individual action(s) in case this episode
is still not numpy'ized yet. In case this episode has already been
numpy'ized, this should be (possibly complex) struct matching the
action space and with a batch size of its leafs exactly the size
of the to-be-overwritten slice or segment (provided by `at_indices`).
at_indices: A single int is interpreted as one index, which to overwrite
with `new_data` (which is expected to be a single observation).
A list of ints is interpreted as a list of indices, all of which to
overwrite with `new_data` (which is expected to be of the same size
as `len(at_indices)`).
A slice object is interpreted as a range of indices to be overwritten
with `new_data` (which is expected to be of the same size as the
provided slice).
Thereby, negative indices by default are interpreted as "before the end"
unless the `neg_index_as_lookback=True` option is used, in which case
negative indices are interpreted as "before ts=0", meaning going back
into the lookback buffer.
neg_index_as_lookback: If True, negative values in `at_indices` are
interpreted as "before ts=0", meaning going back into the lookback
buffer. For example, an episode with
actions = [4, 5, 6, 7, 8, 9], where [4, 5, 6] is the
lookback buffer range (ts=0 item is 7), will handle a call to
`set_actions(individual_action, -1,
neg_index_as_lookback=True)` by overwriting the value of 6 in our
actions buffer with the provided "individual_action".
Raises:
IndexError: If the provided `at_indices` do not match the size of
`new_data`.
"""
for agent_id, new_agent_data in new_data.items():
if agent_id not in self.agent_episodes:
raise KeyError(f"AgentID '{agent_id}' not found in this episode.")
self.agent_episodes[agent_id].set_actions(
new_data=new_agent_data,
at_indices=at_indices,
neg_index_as_lookback=neg_index_as_lookback,
)
def set_rewards(
self,
*,
new_data: MultiAgentDict,
at_indices: Optional[Union[int, List[int], slice]] = None,
neg_index_as_lookback: bool = False,
) -> None:
"""Overwrites all or some of this Episode's rewards with the provided data.
This is a helper method to batch `SingleAgentEpisode.set_rewards`.
For more detail, see `SingleAgentEpisode.set_rewards`.
Args:
new_data: A dict mapping agent IDs to new reward data.
Each value in the dict is the new reward data to overwrite existing data with.
This may be a list of individual reward(s) in case this episode
is still not numpy'ized yet. In case this episode has already been
numpy'ized, this should be a np.ndarray with a length exactly
the size of the to-be-overwritten slice or segment (provided by
`at_indices`).
at_indices: A single int is interpreted as one index, which to overwrite
with `new_data` (which is expected to be a single reward).
A list of ints is interpreted as a list of indices, all of which to
overwrite with `new_data` (which is expected to be of the same size
as `len(at_indices)`).
A slice object is interpreted as a range of indices to be overwritten
with `new_data` (which is expected to be of the same size as the
provided slice).
Thereby, negative indices by default are interpreted as "before the end"
unless the `neg_index_as_lookback=True` option is used, in which case
negative indices are interpreted as "before ts=0", meaning going back
into the lookback buffer.
neg_index_as_lookback: If True, negative values in `at_indices` are
interpreted as "before ts=0", meaning going back into the lookback
buffer. For example, an episode with
rewards = [4, 5, 6, 7, 8, 9], where [4, 5, 6] is the
lookback buffer range (ts=0 item is 7), will handle a call to
`set_rewards(individual_reward, -1,
neg_index_as_lookback=True)` by overwriting the value of 6 in our
rewards buffer with the provided "individual_reward".
Raises:
IndexError: If the provided `at_indices` do not match the size of
`new_data`.
"""
for agent_id, new_agent_data in new_data.items():
if agent_id not in self.agent_episodes:
raise KeyError(f"AgentID '{agent_id}' not found in this episode.")
self.agent_episodes[agent_id].set_rewards(
new_data=new_agent_data,
at_indices=at_indices,
neg_index_as_lookback=neg_index_as_lookback,
)
def set_extra_model_outputs(
self,
*,
key,
new_data: MultiAgentDict,
at_indices: Optional[Union[int, List[int], slice]] = None,
neg_index_as_lookback: bool = False,
) -> None:
"""Overwrites all or some of this Episode's extra model outputs with `new_data`.
This batches `SingleAgentEpisode.set_extra_model_outputs` calls for the
agent IDs provided in `new_data`. For more detail, see
`SingleAgentEpisode.set_extra_model_outputs`.
Args:
key: Existing extra model output key to override on each target agent
episode.
new_data: A dict mapping agent IDs to new extra model outputs data.
Each value in the dict is the new extra model outputs data to overwrite existing data with.
This may be a list of individual reward(s) in case this episode
is still not numpy'ized yet. In case this episode has already been
numpy'ized, this should be a np.ndarray with a length exactly
the size of the to-be-overwritten slice or segment (provided by
`at_indices`).
at_indices: A single int is interpreted as one index, which to overwrite
with `new_data` (which is expected to be a single extra model output).
A list of ints is interpreted as a list of indices, all of which to
overwrite with `new_data` (which is expected to be of the same size
as `len(at_indices)`).
A slice object is interpreted as a range of indices to be overwritten
with `new_data` (which is expected to be of the same size as the
provided slice).
Thereby, negative indices by default are interpreted as "before the end"
unless the `neg_index_as_lookback=True` option is used, in which case
negative indices are interpreted as "before ts=0", meaning going back
into the lookback buffer.
neg_index_as_lookback: If True, negative values in `at_indices` are
interpreted as "before ts=0", meaning going back into the lookback
buffer. For example, an episode with
extra_model_outputs[key][agent_id] = [4, 5, 6, 7, 8, 9], where [4, 5, 6] is the
lookback buffer range (ts=0 item is 7), will handle a call to
`set_extra_model_outputs(key, individual_output, -1,
neg_index_as_lookback=True)` by overwriting the value of 6 in our
extra_model_outputs[key][agent_id] buffer with the provided "individual_output".
Raises:
IndexError: If the provided `at_indices` do not match the size of
`new_data`.
"""
for agent_id, new_agent_data in new_data.items():
if agent_id not in self.agent_episodes:
raise KeyError(f"AgentID '{agent_id}' not found in this episode.")
self.agent_episodes[agent_id].set_extra_model_outputs(
key=key,
new_data=new_agent_data,
at_indices=at_indices,
neg_index_as_lookback=neg_index_as_lookback,
)
def env_steps(self) -> int:
"""Returns the number of environment steps.
Note, this episode instance could be a chunk of an actual episode.
Returns:
An integer that counts the number of environment steps this episode instance
has seen.
"""
return len(self)
def agent_steps(self) -> int:
"""Number of agent steps.
Note, there are >= 1 agent steps per environment step.
Returns:
An integer counting the number of agent steps executed during the time this
episode instance records.
"""
return sum(len(eps) for eps in self.agent_episodes.values())
def __getitem__(self, item: slice) -> "MultiAgentEpisode":
"""Enable squared bracket indexing- and slicing syntax, e.g. episode[-4:]."""
if isinstance(item, slice):
return self.slice(slice_=item)
else:
raise NotImplementedError(
f"MultiAgentEpisode does not support getting item '{item}'! "
"Only slice objects allowed with the syntax: `episode[a:b]`."
)
def _init_single_agent_episodes(
self,
*,
agent_module_ids: Optional[Dict[AgentID, ModuleID]] = None,
agent_episode_ids: Optional[Dict[AgentID, str]] = None,
observations: Optional[List[MultiAgentDict]] = None,
actions: Optional[List[MultiAgentDict]] = None,
rewards: Optional[List[MultiAgentDict]] = None,
infos: Optional[List[MultiAgentDict]] = None,
terminateds: Union[MultiAgentDict, bool] = False,
truncateds: Union[MultiAgentDict, bool] = False,
extra_model_outputs: Optional[List[MultiAgentDict]] = None,
):
if observations is None:
return
if actions is None:
assert not rewards
assert not extra_model_outputs
actions = []
rewards = []
extra_model_outputs = []
# Infos and `extra_model_outputs` are allowed to be None -> Fill them with
# proper dummy values, if so.
if infos is None:
infos = [{} for _ in range(len(observations))]
if extra_model_outputs is None:
extra_model_outputs = [{} for _ in range(len(actions))]
observations_per_agent = defaultdict(list)
infos_per_agent = defaultdict(list)
actions_per_agent = defaultdict(list)
rewards_per_agent = defaultdict(list)
extra_model_outputs_per_agent = defaultdict(list)
done_per_agent = defaultdict(bool)
len_lookback_buffer_per_agent = defaultdict(lambda: self._len_lookback_buffers)
all_agent_ids = set(
agent_episode_ids.keys() if agent_episode_ids is not None else []
)
agent_module_ids = agent_module_ids or {}
# First pass: count observations per agent in lookback AND total.
# This allows us to recover the correct env_t_to_agent_t mapping.
lookback_obs_count_per_agent = defaultdict(int)
total_obs_count_per_agent = defaultdict(int)
for data_idx, obs in enumerate(observations):
for agent_id in obs:
total_obs_count_per_agent[agent_id] += 1
if data_idx < self._len_lookback_buffers:
lookback_obs_count_per_agent[agent_id] += 1
# Compute the starting agent_t for each agent.
# The formula depends on whether there are observations after the lookback:
# - If new_chunk_obs > 0: first_agent_t = agent_t_started - lookback_count
# - If new_chunk_obs == 0: first_agent_t = agent_t_started - lookback_count + 1
# This is because agent_t_started = len(completed_actions), which equals the
# observation_index of the NEXT observation if there is one, or the LAST
# observation if the action is still hanging.
current_agent_t = {}
for agent_id, lookback_count in lookback_obs_count_per_agent.items():
total_count = total_obs_count_per_agent[agent_id]
new_chunk_obs = total_count - lookback_count
if new_chunk_obs > 0:
current_agent_t[agent_id] = (
self.agent_t_started[agent_id] - lookback_count
)
else:
current_agent_t[agent_id] = (
self.agent_t_started[agent_id] - lookback_count + 1
)
# Step through all observations and interpret these as the (global) env steps.
for data_idx, (obs, inf) in enumerate(zip(observations, infos)):
# If we do have actions/extra outs/rewards for this timestep, use the data.
# It may be that these lists have the same length as the observations list,
# in which case the data will be cached (agent did step/send an action,
# but the step has not been concluded yet by the env).
act = actions[data_idx] if len(actions) > data_idx else {}
extra_outs = (
extra_model_outputs[data_idx]
if len(extra_model_outputs) > data_idx
else {}
)
rew = rewards[data_idx] if len(rewards) > data_idx else {}
for agent_id, agent_obs in obs.items():
all_agent_ids.add(agent_id)
observations_per_agent[agent_id].append(agent_obs)
infos_per_agent[agent_id].append(inf.get(agent_id, {}))
# Pull out hanging action (if not first obs for this agent) and
# complete step for agent.
if len(observations_per_agent[agent_id]) > 1:
actions_per_agent[agent_id].append(
self._hanging_actions_end.pop(agent_id)
)
extra_model_outputs_per_agent[agent_id].append(
self._hanging_extra_model_outputs_end.pop(agent_id)
)
rewards_per_agent[agent_id].append(
self._hanging_rewards_end.pop(agent_id)
)
# First obs for this agent. Make sure the agent's mapping is
# appropriately prepended with self.SKIP_ENV_TS_TAG tags.
else:
if agent_id not in self.env_t_to_agent_t:
self.env_t_to_agent_t[agent_id].extend(
[self.SKIP_ENV_TS_TAG] * data_idx
)
len_lookback_buffer_per_agent[agent_id] -= data_idx
# Agent is still continuing (has an action for the next step).
if agent_id in act:
# Always push actions/extra outputs into cache, then remove them
# from there, once the next observation comes in. Same for rewards.
self._hanging_actions_end[agent_id] = act[agent_id]
self._hanging_extra_model_outputs_end[agent_id] = extra_outs.get(
agent_id, {}
)
self._hanging_rewards_end[agent_id] += rew.get(agent_id, 0.0)
# Agent is done (has no action for the next step).
elif terminateds.get(agent_id) or truncateds.get(agent_id):
done_per_agent[agent_id] = True
# There is more (global) action/reward data. This agent must therefore
# be done. Automatically add it to `done_per_agent` and `terminateds`.
elif data_idx < len(observations) - 1:
done_per_agent[agent_id] = terminateds[agent_id] = True
# Update env_t_to_agent_t mapping using the recovered agent_t.
# For agents in the lookback, current_agent_t was computed earlier as:
# agent_t_started - lookback_obs_count
# For agents not in lookback but with prior history, use agent_t_started.
# For truly new agents (no prior history), start at 0.
if agent_id not in current_agent_t:
current_agent_t[agent_id] = self.agent_t_started.get(agent_id, 0)
self.env_t_to_agent_t[agent_id].append(current_agent_t[agent_id])
current_agent_t[agent_id] += 1
# Those agents that did NOT step:
# - Get self.SKIP_ENV_TS_TAG added to their env_t_to_agent_t mapping.
# - Get their reward (if any) added up.
for agent_id in all_agent_ids:
if agent_id not in obs and agent_id not in done_per_agent:
self.env_t_to_agent_t[agent_id].append(self.SKIP_ENV_TS_TAG)
# If we are still in the global lookback buffer segment, deduct 1
# from this agents' lookback buffer, b/c we don't want the agent
# to use this (missing) obs/data in its single-agent lookback.
if (
len(self.env_t_to_agent_t[agent_id])
- self._len_lookback_buffers
<= 0
):
len_lookback_buffer_per_agent[agent_id] -= 1
self._hanging_rewards_end[agent_id] += rew.get(agent_id, 0.0)
# - Validate per-agent data.
# - Fix lookback buffers of env_t_to_agent_t mappings.
for agent_id in list(self.env_t_to_agent_t.keys()):
# Skip agent if it doesn't seem to have any data.
if agent_id not in observations_per_agent:
del self.env_t_to_agent_t[agent_id]
continue
assert (
len(observations_per_agent[agent_id])
== len(infos_per_agent[agent_id])
== len(actions_per_agent[agent_id]) + 1
== len(extra_model_outputs_per_agent[agent_id]) + 1
== len(rewards_per_agent[agent_id]) + 1
)
self.env_t_to_agent_t[agent_id].lookback = self._len_lookback_buffers
# Now create the individual episodes from the collected per-agent data.
for agent_id, agent_obs in observations_per_agent.items():
# If agent only has a single obs AND is already done, remove all its traces
# from this MultiAgentEpisode.
if len(agent_obs) == 1 and done_per_agent.get(agent_id):
self._del_agent(agent_id)
continue
# Try to figure out the module ID for this agent.
# If not provided explicitly by the user that initializes this episode
# object, try our mapping function.
module_id = agent_module_ids.get(
agent_id, self.agent_to_module_mapping_fn(agent_id, self)
)
# Create this agent's SingleAgentEpisode.
sa_episode = SingleAgentEpisode(
id_=(
agent_episode_ids.get(agent_id)
if agent_episode_ids is not None
else None
),
agent_id=agent_id,
module_id=module_id,
multi_agent_episode_id=self.id_,
observations=agent_obs,
observation_space=self.observation_space.get(agent_id),
infos=infos_per_agent[agent_id],
actions=actions_per_agent[agent_id],
action_space=self.action_space.get(agent_id),
rewards=rewards_per_agent[agent_id],
extra_model_outputs=(
{
k: [i[k] for i in extra_model_outputs_per_agent[agent_id]]
for k in extra_model_outputs_per_agent[agent_id][0].keys()
}
if extra_model_outputs_per_agent[agent_id]
else None
),
terminated=terminateds.get(agent_id, False),
truncated=truncateds.get(agent_id, False),
t_started=self.agent_t_started[agent_id],
len_lookback_buffer=max(len_lookback_buffer_per_agent[agent_id], 0),
)
# and store it.
self.agent_episodes[agent_id] = sa_episode
def _get(
self,
*,
what,
indices,
agent_ids=None,
env_steps=True,
neg_index_as_lookback=False,
fill=None,
one_hot_discrete=False,
return_list=False,
extra_model_outputs_key=None,
):
agent_ids = set(force_list(agent_ids)) or self.agent_ids
kwargs = dict(
what=what,
indices=indices,
agent_ids=agent_ids,
neg_index_as_lookback=neg_index_as_lookback,
fill=fill,
# Rewards and infos do not support one_hot_discrete option.
one_hot_discrete=dict(
{} if not one_hot_discrete else {"one_hot_discrete": one_hot_discrete}
),
extra_model_outputs_key=extra_model_outputs_key,
)
# User specified agent timesteps (indices) -> Simply delegate everything
# to the individual agents' SingleAgentEpisodes.
if env_steps is False:
if return_list:
raise ValueError(
f"`MultiAgentEpisode.get_{what}()` can't be called with both "
"`env_steps=False` and `return_list=True`!"
)
return self._get_data_by_agent_steps(**kwargs)
# User specified env timesteps (indices) -> We need to translate them for each
# agent into agent-timesteps.
# Return a list of individual per-env-timestep multi-agent dicts.
elif return_list:
return self._get_data_by_env_steps_as_list(**kwargs)
# Return a single multi-agent dict with lists/arrays as leafs.
else:
return self._get_data_by_env_steps(**kwargs)
def _get_data_by_agent_steps(
self,
*,
what,
indices,
agent_ids,
neg_index_as_lookback,
fill,
one_hot_discrete,
extra_model_outputs_key,
):
# Return requested data by agent-steps.
ret = {}
# For each agent, we retrieve the data through passing the given indices into
# the SingleAgentEpisode of that agent.
for agent_id, sa_episode in self.agent_episodes.items():
if agent_id not in agent_ids:
continue
inf_lookback_buffer = getattr(sa_episode, what)
hanging_val = self._get_hanging_value(what, agent_id)
# User wants a specific `extra_model_outputs` key.
if extra_model_outputs_key is not None:
inf_lookback_buffer = inf_lookback_buffer[extra_model_outputs_key]
hanging_val = hanging_val[extra_model_outputs_key]
agent_value = inf_lookback_buffer.get(
indices=indices,
neg_index_as_lookback=neg_index_as_lookback,
fill=fill,
_add_last_ts_value=hanging_val,
**one_hot_discrete,
)
if agent_value is None or (
isinstance(agent_value, list) and agent_value == []
):
continue
ret[agent_id] = agent_value
return ret
def _get_data_by_env_steps_as_list(
self,
*,
what: str,
indices: Union[int, slice, List[int]],
agent_ids: Collection[AgentID],
neg_index_as_lookback: bool,
fill: Any,
one_hot_discrete,
extra_model_outputs_key: str,
) -> List[MultiAgentDict]:
# Collect indices for each agent first, so we can construct the list in
# the next step.
agent_indices = {}
for agent_id in self.agent_episodes.keys():
if agent_id not in agent_ids:
continue
agent_t_indices = self.env_t_to_agent_t[agent_id].get(
indices,
neg_index_as_lookback=neg_index_as_lookback,
fill=self.SKIP_ENV_TS_TAG,
# For those records where there is no "hanging" last timestep (all
# other than obs and infos), we have to ignore the last entry in
# the env_t_to_agent_t mappings.
_ignore_last_ts=what not in ["observations", "infos"],
)
# Convert absolute agent_t to buffer position (including lookback offset).
# Formula: buffer_pos = agent_t - agent_t_started + lookback
sa_episode = self.agent_episodes[agent_id]
lookback = sa_episode.observations.lookback
if isinstance(agent_t_indices, int):
if agent_t_indices != self.SKIP_ENV_TS_TAG:
agent_t_indices = (
agent_t_indices - self.agent_t_started[agent_id] + lookback
)
else:
assert isinstance(agent_t_indices, list)
agent_t_indices = [
index - self.agent_t_started[agent_id] + lookback
if index != self.SKIP_ENV_TS_TAG
else index
for index in agent_t_indices
]
agent_indices[agent_id] = agent_t_indices
if not agent_indices:
return []
ret = []
for i in range(len(next(iter(agent_indices.values())))):
ret2 = {}
for agent_id, idxes in agent_indices.items():
hanging_val = self._get_hanging_value(what, agent_id)
(
inf_lookback_buffer,
indices_to_use,
) = self._get_inf_lookback_buffer_or_dict(
agent_id,
what,
extra_model_outputs_key,
hanging_val,
filter_for_skip_indices=idxes[i],
)
if (
what == "extra_model_outputs"
and not inf_lookback_buffer
and not hanging_val
):
continue
agent_value = self._get_single_agent_data_by_index(
what=what,
inf_lookback_buffer=inf_lookback_buffer,
agent_id=agent_id,
index_incl_lookback=indices_to_use,
fill=fill,
one_hot_discrete=one_hot_discrete,
extra_model_outputs_key=extra_model_outputs_key,
hanging_val=hanging_val,
)
if agent_value is not None:
ret2[agent_id] = agent_value
ret.append(ret2)
return ret
def _get_data_by_env_steps(
self,
*,
what: str,
indices: Union[int, slice, List[int]],
agent_ids: Collection[AgentID],
neg_index_as_lookback: bool,
fill: Any,
one_hot_discrete: bool,
extra_model_outputs_key: str,
) -> MultiAgentDict:
ignore_last_ts = what not in ["observations", "infos"]
ret = {}
for agent_id, sa_episode in self.agent_episodes.items():
if agent_id not in agent_ids:
continue
hanging_val = self._get_hanging_value(what, agent_id)
agent_indices = self.env_t_to_agent_t[agent_id].get(
indices,
neg_index_as_lookback=neg_index_as_lookback,
fill=self.SKIP_ENV_TS_TAG if fill is not None else None,
# For those records where there is no "hanging" last timestep (all
# other than obs and infos), we have to ignore the last entry in
# the env_t_to_agent_t mappings.
_ignore_last_ts=ignore_last_ts,
)
inf_lookback_buffer, agent_indices = self._get_inf_lookback_buffer_or_dict(
agent_id,
what,
extra_model_outputs_key,
hanging_val,
filter_for_skip_indices=agent_indices,
)
# Convert absolute agent_t to buffer position (including lookback offset).
# Formula: buffer_pos = agent_t - agent_t_started + lookback
lookback = sa_episode.observations.lookback
if isinstance(agent_indices, list):
agent_indices = [
index - self.agent_t_started[agent_id] + lookback
if index != self.SKIP_ENV_TS_TAG
else index
for index in agent_indices
]
agent_values = self._get_single_agent_data_by_env_step_indices(
what=what,
agent_id=agent_id,
indices_incl_lookback=agent_indices,
fill=fill,
one_hot_discrete=one_hot_discrete,
hanging_val=hanging_val,
extra_model_outputs_key=extra_model_outputs_key,
)
if len(agent_values) > 0:
ret[agent_id] = agent_values
else:
if agent_indices != self.SKIP_ENV_TS_TAG:
agent_indices = (
agent_indices - self.agent_t_started[agent_id] + lookback
)
agent_values = self._get_single_agent_data_by_index(
what=what,
inf_lookback_buffer=inf_lookback_buffer,
agent_id=agent_id,
index_incl_lookback=agent_indices,
fill=fill,
one_hot_discrete=one_hot_discrete,
extra_model_outputs_key=extra_model_outputs_key,
hanging_val=hanging_val,
)
if agent_values is not None:
ret[agent_id] = agent_values
return ret
def _get_single_agent_data_by_index(
self,
*,
what: str,
inf_lookback_buffer: InfiniteLookbackBuffer,
agent_id: AgentID,
index_incl_lookback: Union[int, str],
fill: Any,
one_hot_discrete: dict,
extra_model_outputs_key: str,
hanging_val: Any,
) -> Any:
sa_episode = self.agent_episodes[agent_id]
if index_incl_lookback == self.SKIP_ENV_TS_TAG:
# We don't want to fill -> Skip this agent.
if fill is None:
return None
# Provide filled value for this agent.
return getattr(sa_episode, f"get_{what}")(
indices=1000000000000,
neg_index_as_lookback=False,
fill=fill,
**dict(
{}
if extra_model_outputs_key is None
else {"key": extra_model_outputs_key}
),
**one_hot_discrete,
)
# No skip timestep -> Provide value at given index for this agent.
# Special case: extra_model_outputs and key=None (return all keys as
# a dict). Note that `inf_lookback_buffer` is NOT an infinite lookback
# buffer, but a dict mapping keys to individual infinite lookback
# buffers.
elif what == "extra_model_outputs" and extra_model_outputs_key is None:
assert hanging_val is None or isinstance(hanging_val, dict)
ret = {}
if inf_lookback_buffer:
for key, sub_buffer in inf_lookback_buffer.items():
ret[key] = sub_buffer.get(
indices=index_incl_lookback - sub_buffer.lookback,
neg_index_as_lookback=True,
fill=fill,
_add_last_ts_value=(
None if hanging_val is None else hanging_val[key]
),
**one_hot_discrete,
)
else:
for key in hanging_val.keys():
ret[key] = InfiniteLookbackBuffer().get(
indices=index_incl_lookback,
neg_index_as_lookback=True,
fill=fill,
_add_last_ts_value=hanging_val[key],
**one_hot_discrete,
)
return ret
# Extract data directly from the infinite lookback buffer object.
else:
# When a specific extra_model_outputs key was requested, the callers
# may pass the full extra_model_outputs dict (instead of the indexed
# sub-buffer/value). Ensure both are properly indexed here.
if what == "extra_model_outputs" and extra_model_outputs_key is not None:
if isinstance(inf_lookback_buffer, dict):
inf_lookback_buffer = inf_lookback_buffer[extra_model_outputs_key]
if hanging_val is not None and isinstance(hanging_val, dict):
hanging_val = hanging_val[extra_model_outputs_key]
return inf_lookback_buffer.get(
indices=index_incl_lookback - inf_lookback_buffer.lookback,
neg_index_as_lookback=True,
fill=fill,
_add_last_ts_value=hanging_val,
**one_hot_discrete,
)
def _get_single_agent_data_by_env_step_indices(
self,
*,
what: str,
agent_id: AgentID,
indices_incl_lookback: Union[int, str],
fill: Optional[Any] = None,
one_hot_discrete: bool = False,
extra_model_outputs_key: Optional[str] = None,
hanging_val: Optional[Any] = None,
) -> Any:
"""Returns single data item from the episode based on given (env step) indices.
The returned data item will have a batch size that matches the env timesteps
defined via `indices_incl_lookback`.
Args:
what: A (str) descriptor of what data to collect. Must be one of
"observations", "infos", "actions", "rewards", or "extra_model_outputs".
indices_incl_lookback: A list of ints specifying, which indices
to pull from the InfiniteLookbackBuffer defined by `agent_id` and `what`
(and maybe `extra_model_outputs_key`). Note that these indices
disregard the special logic of the lookback buffer. Meaning if one
index in `indices_incl_lookback` is 0, then the first value in the
lookback buffer should be returned, not the first value after the
lookback buffer (which would be normal behavior for pulling items from
an `InfiniteLookbackBuffer` object).
agent_id: The individual agent ID to pull data for. Used to look up the
`SingleAgentEpisode` object for this agent in `self`.
fill: An optional float value to use for filling up the returned results at
the boundaries. This filling only happens if the requested index range's
start/stop boundaries exceed the buffer's boundaries (including the
lookback buffer on the left side). This comes in very handy, if users
don't want to worry about reaching such boundaries and want to zero-pad.
For example, a buffer with data [10, 11, 12, 13, 14] and lookback
buffer size of 2 (meaning `10` and `11` are part of the lookback buffer)
will respond to `indices_incl_lookback=[-1, -2, 0]` and `fill=0.0`
with `[0.0, 0.0, 10]`.
one_hot_discrete: If True, will return one-hot vectors (instead of
int-values) for those sub-components of a (possibly complex) space
that are Discrete or MultiDiscrete. Note that if `fill=0` and the
requested `indices_incl_lookback` are out of the range of our data, the
returned one-hot vectors will actually be zero-hot (all slots zero).
extra_model_outputs_key: Only if what is "extra_model_outputs", this
specifies the sub-key (str) inside the extra_model_outputs dict, e.g.
STATE_OUT or ACTION_DIST_INPUTS.
hanging_val: In case we are pulling actions, rewards, or extra_model_outputs
data, there might be information "hanging" (cached). For example,
if an agent receives an observation o0 and then immediately sends an
action a0 back, but then does NOT immediately retrieve the next
observation, a0 is now cached (not fully logged yet with this
episode). The currently cached value must be provided here to be able
to return it in case the index is -1 (most recent timestep).
Returns:
A data item corresponding to the provided args.
"""
sa_episode = self.agent_episodes[agent_id]
inf_lookback_buffer = getattr(sa_episode, what)
if extra_model_outputs_key is not None:
inf_lookback_buffer = inf_lookback_buffer[extra_model_outputs_key]
if hanging_val is not None:
hanging_val = hanging_val[extra_model_outputs_key]
# If there are self.SKIP_ENV_TS_TAG items in `indices_incl_lookback` and user
# wants to fill these (together with outside-episode-bounds indices) ->
# Provide these skipped timesteps as filled values.
if self.SKIP_ENV_TS_TAG in indices_incl_lookback and fill is not None:
single_fill_value = inf_lookback_buffer.get(
indices=1000000000000,
neg_index_as_lookback=False,
fill=fill,
**one_hot_discrete,
)
ret = []
for i in indices_incl_lookback:
if i == self.SKIP_ENV_TS_TAG:
ret.append(single_fill_value)
else:
ret.append(
inf_lookback_buffer.get(
indices=i - getattr(sa_episode, what).lookback,
neg_index_as_lookback=True,
fill=fill,
_add_last_ts_value=hanging_val,
**one_hot_discrete,
)
)
if self.is_numpy:
ret = batch(ret)
else:
# Filter these indices out up front.
indices = [
i - inf_lookback_buffer.lookback
for i in indices_incl_lookback
if i != self.SKIP_ENV_TS_TAG
]
# If all indices were SKIP_ENV_TS_TAG (agent was inactive for all
# requested env steps), return an empty list. The caller already
# checks `if len(agent_values) > 0` before using this result.
if not indices:
ret = []
else:
ret = inf_lookback_buffer.get(
indices=indices,
neg_index_as_lookback=True,
fill=fill,
_add_last_ts_value=hanging_val,
**one_hot_discrete,
)
return ret
def _get_hanging_value(self, what: str, agent_id: AgentID) -> Any:
"""Returns the hanging action/reward/extra_model_outputs for given agent."""
if what == "actions":
return self._hanging_actions_end.get(agent_id)
elif what == "extra_model_outputs":
return self._hanging_extra_model_outputs_end.get(agent_id)
elif what == "rewards":
return self._hanging_rewards_end.get(agent_id)
def _copy_hanging(self, agent_id: AgentID, other: "MultiAgentEpisode") -> None:
"""Copies hanging action, reward, extra_model_outputs from `other` to `self."""
if agent_id in other._hanging_rewards_begin:
self._hanging_rewards_begin[agent_id] = other._hanging_rewards_begin[
agent_id
]
if agent_id in other._hanging_rewards_end:
self._hanging_actions_end[agent_id] = copy.deepcopy(
other._hanging_actions_end[agent_id]
)
self._hanging_rewards_end[agent_id] = other._hanging_rewards_end[agent_id]
self._hanging_extra_model_outputs_end[agent_id] = copy.deepcopy(
other._hanging_extra_model_outputs_end[agent_id]
)
def _del_hanging(self, agent_id: AgentID) -> None:
"""Deletes all hanging action, reward, extra_model_outputs of given agent."""
self._hanging_rewards_begin.pop(agent_id, None)
self._hanging_actions_end.pop(agent_id, None)
self._hanging_extra_model_outputs_end.pop(agent_id, None)
self._hanging_rewards_end.pop(agent_id, None)
def _del_agent(self, agent_id: AgentID) -> None:
"""Deletes all data of given agent from this episode."""
self._del_hanging(agent_id)
self.agent_episodes.pop(agent_id, None)
self.agent_ids.discard(agent_id)
self.env_t_to_agent_t.pop(agent_id, None)
self._agent_to_module_mapping.pop(agent_id, None)
self.agent_t_started.pop(agent_id, None)
def _get_inf_lookback_buffer_or_dict(
self,
agent_id: AgentID,
what: str,
extra_model_outputs_key: Optional[str] = None,
hanging_val: Optional[Any] = None,
filter_for_skip_indices=None,
):
"""Returns a single InfiniteLookbackBuffer or a dict of such.
In case `what` is "extra_model_outputs" AND `extra_model_outputs_key` is None,
a dict is returned. In all other cases, a single InfiniteLookbackBuffer is
returned.
"""
inf_lookback_buffer_or_dict = inf_lookback_buffer = getattr(
self.agent_episodes[agent_id], what
)
if what == "extra_model_outputs":
if extra_model_outputs_key is not None:
inf_lookback_buffer = inf_lookback_buffer_or_dict[
extra_model_outputs_key
]
elif inf_lookback_buffer_or_dict:
inf_lookback_buffer = next(iter(inf_lookback_buffer_or_dict.values()))
elif filter_for_skip_indices is not None:
return inf_lookback_buffer_or_dict, filter_for_skip_indices
else:
return inf_lookback_buffer_or_dict
if filter_for_skip_indices is not None:
inf_lookback_buffer_len = (
len(inf_lookback_buffer)
+ inf_lookback_buffer.lookback
+ (hanging_val is not None)
)
ignore_last_ts = what not in ["observations", "infos"]
if isinstance(filter_for_skip_indices, list):
filter_for_skip_indices = [
"S" if ignore_last_ts and i == inf_lookback_buffer_len else i
for i in filter_for_skip_indices
]
elif ignore_last_ts and filter_for_skip_indices == inf_lookback_buffer_len:
filter_for_skip_indices = "S"
return inf_lookback_buffer_or_dict, filter_for_skip_indices
else:
return inf_lookback_buffer_or_dict
@Deprecated(new="MultiAgentEpisode.custom_data[some-key] = ...", error=True)
def add_temporary_timestep_data(self):
pass
@Deprecated(new="MultiAgentEpisode.custom_data[some-key]", error=True)
def get_temporary_timestep_data(self):
pass
Reference in New Issue View Git Blame Copy Permalink