wehub-resource-sync
767 lines
30 KiB
Python
767 lines
30 KiB
Python
from typing import Any, Dict, List, Optional, Union
|
|
|
|
import gymnasium as gym
|
|
import numpy as np
|
|
import tree # pip install dm_tree
|
|
from gymnasium.utils.env_checker import data_equivalence
|
|
|
|
from ray.rllib.utils.numpy import LARGE_INTEGER, one_hot, one_hot_multidiscrete
|
|
from ray.rllib.utils.serialization import gym_space_from_dict, gym_space_to_dict
|
|
from ray.rllib.utils.spaces.space_utils import (
|
|
batch,
|
|
get_base_struct_from_space,
|
|
get_dummy_batch_for_space,
|
|
)
|
|
from ray.util.annotations import DeveloperAPI
|
|
|
|
|
|
@DeveloperAPI
|
|
class InfiniteLookbackBuffer:
|
|
def __init__(
|
|
self,
|
|
data: Optional[Union[List, np.ndarray]] = None,
|
|
lookback: int = 0,
|
|
space: Optional[gym.Space] = None,
|
|
):
|
|
self.data = data if data is not None else []
|
|
self.lookback = min(lookback, len(self.data))
|
|
self.finalized = not isinstance(self.data, list)
|
|
|
|
self.space = space
|
|
|
|
def __eq__(
|
|
self,
|
|
other: "InfiniteLookbackBuffer",
|
|
) -> bool:
|
|
"""Compares two `InfiniteLookbackBuffers.
|
|
|
|
Args:
|
|
other: Another object. If another `LookbackBuffer` instance all
|
|
their attributes are compared.
|
|
|
|
Returns:
|
|
`True`, if `other` is an `InfiniteLookbackBuffer` instance and all
|
|
attributes are identical. Otherwise, returns `False`.
|
|
"""
|
|
return (
|
|
isinstance(other, InfiniteLookbackBuffer)
|
|
# Todo (mark): Replace `data_equivalence` with ray / rllib implementation similar to `check` without asserts
|
|
and data_equivalence(self.data, other.data)
|
|
and self.lookback == other.lookback
|
|
and self.finalized == other.finalized
|
|
and self.space_struct == other.space_struct
|
|
and self.space == other.space
|
|
)
|
|
|
|
@property
|
|
def space(self):
|
|
return self._space
|
|
|
|
@space.setter
|
|
def space(self, value):
|
|
self._space = value
|
|
self._space_struct = get_base_struct_from_space(value)
|
|
|
|
@property
|
|
def space_struct(self):
|
|
return self._space_struct
|
|
|
|
def get_state(self) -> Dict[str, Any]:
|
|
"""Returns the pickable state of a buffer.
|
|
|
|
The data in the buffer is stored into a dictionary. Note that
|
|
buffers can also be generated from pickable states (see
|
|
`InfiniteLookbackBuffer.from_state`)
|
|
|
|
Returns:
|
|
A dict containing all the data and metadata from the buffer.
|
|
"""
|
|
return {
|
|
"data": self.data,
|
|
"lookback": self.lookback,
|
|
"finalized": self.finalized,
|
|
"space": gym_space_to_dict(self.space) if self.space else None,
|
|
}
|
|
|
|
@staticmethod
|
|
def from_state(state: Dict[str, Any]) -> "InfiniteLookbackBuffer":
|
|
"""Creates a new `InfiniteLookbackBuffer` from a state dict.
|
|
|
|
Args:
|
|
state: The state dict, as returned by `self.get_state`.
|
|
|
|
Returns:
|
|
A new `InfiniteLookbackBuffer` instance with the data and metadata
|
|
from the state dict.
|
|
"""
|
|
buffer = InfiniteLookbackBuffer()
|
|
buffer.lookback = state["lookback"]
|
|
buffer.finalized = state["finalized"]
|
|
buffer.space = gym_space_from_dict(state["space"]) if state["space"] else None
|
|
# space_struct is set when space is assigned
|
|
buffer.data = state["data"]
|
|
|
|
return buffer
|
|
|
|
def append(self, item) -> None:
|
|
"""Appends the given item to the end of this buffer."""
|
|
if self.finalized:
|
|
if isinstance(self.data, np.ndarray):
|
|
self.data = np.concatenate(
|
|
[self.data, np.asarray(item)[np.newaxis]], axis=0
|
|
)
|
|
else:
|
|
self.data = tree.map_structure(
|
|
lambda d, i: np.concatenate([d, [i]], axis=0), self.data, item
|
|
)
|
|
else:
|
|
self.data.append(item)
|
|
|
|
def extend(self, items) -> None:
|
|
"""Appends all items in `items` to the end of this buffer."""
|
|
if self.finalized:
|
|
# TODO (sven): When extending with a list of structs, we should
|
|
# probably rather do: `tree.map_structure(..., self.data,
|
|
# tree.map_structure(lambda *s: np.array(*s), *items)`)??
|
|
self.data = tree.map_structure(
|
|
lambda d, i: np.concatenate([d, i], axis=0),
|
|
self.data,
|
|
# Note, we could have dictionaries here.
|
|
np.array(items) if isinstance(items, list) else items,
|
|
)
|
|
else:
|
|
for item in items:
|
|
self.append(item)
|
|
|
|
def concat(self, other: "InfiniteLookbackBuffer") -> None:
|
|
"""Concatenates the data of `other` (w/o its lookback) to `self`.
|
|
|
|
Args:
|
|
other: The other InfiniteLookbackBuffer to be concatenated to self.
|
|
"""
|
|
self.data.extend(other.get())
|
|
|
|
def pop(self, index: int = -1) -> None:
|
|
"""Removes the item at `index` from this buffer, but does NOT return it.
|
|
|
|
Args:
|
|
index: The index to pop out of this buffer (w/o returning it from this
|
|
method).
|
|
"""
|
|
if self.finalized:
|
|
self.data = tree.map_structure(
|
|
lambda s: np.delete(s, index, axis=0), self.data
|
|
)
|
|
else:
|
|
self.data.pop(index)
|
|
|
|
def finalize(self) -> None:
|
|
"""Finalizes this buffer by converting internal data lists into numpy arrays.
|
|
|
|
Thereby, if the individual items in the list are nested structures, the
|
|
resulting buffer content will be a nested struct of np.ndarrays (leafs).
|
|
"""
|
|
if not self.finalized:
|
|
self.data = batch(self.data)
|
|
self.finalized = True
|
|
|
|
def get(
|
|
self,
|
|
indices: Optional[Union[int, slice, List[int]]] = None,
|
|
*,
|
|
neg_index_as_lookback: bool = False,
|
|
fill: Optional[Any] = None,
|
|
one_hot_discrete: bool = False,
|
|
_ignore_last_ts: bool = False,
|
|
_add_last_ts_value: Optional[Any] = None,
|
|
) -> Any:
|
|
"""Returns data, based on the given args, from this buffer.
|
|
|
|
Args:
|
|
indices: A single int is interpreted as an index, from which to return the
|
|
individual data stored at this index.
|
|
A list of ints is interpreted as a list of indices from which to gather
|
|
individual data in a batch of size len(indices).
|
|
A slice object is interpreted as a range of data 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.
|
|
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, a buffer with data [4, 5, 6, 7, 8, 9],
|
|
where [4, 5, 6] is the lookback buffer range (ts=0 item is 7), will
|
|
respond to `get(-1, neg_index_as_lookback=True)` with `6` and to
|
|
`get(slice(-2, 1), neg_index_as_lookback=True)` with `[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 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 `get(slice(-7, -2), fill=0.0)`
|
|
with `[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) 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).
|
|
_ignore_last_ts: Whether to ignore the last record in our internal
|
|
`self.data` when getting the provided indices.
|
|
_add_last_ts_value: Whether to add the value of this arg to the end of
|
|
the internal `self.data` buffer (just for the duration of this get
|
|
operation, not permanently).
|
|
"""
|
|
if indices is None:
|
|
data = self._get_all_data(
|
|
one_hot_discrete=one_hot_discrete,
|
|
_ignore_last_ts=_ignore_last_ts,
|
|
)
|
|
elif isinstance(indices, slice):
|
|
data = self._get_slice(
|
|
indices,
|
|
fill=fill,
|
|
neg_index_as_lookback=neg_index_as_lookback,
|
|
one_hot_discrete=one_hot_discrete,
|
|
_ignore_last_ts=_ignore_last_ts,
|
|
_add_last_ts_value=_add_last_ts_value,
|
|
)
|
|
elif isinstance(indices, list):
|
|
data = [
|
|
self._get_int_index(
|
|
idx,
|
|
fill=fill,
|
|
neg_index_as_lookback=neg_index_as_lookback,
|
|
one_hot_discrete=one_hot_discrete,
|
|
_ignore_last_ts=_ignore_last_ts,
|
|
_add_last_ts_value=_add_last_ts_value,
|
|
)
|
|
for idx in indices
|
|
]
|
|
if self.finalized:
|
|
data = batch(data)
|
|
else:
|
|
assert isinstance(indices, int)
|
|
data = self._get_int_index(
|
|
indices,
|
|
fill=fill,
|
|
neg_index_as_lookback=neg_index_as_lookback,
|
|
one_hot_discrete=one_hot_discrete,
|
|
_ignore_last_ts=_ignore_last_ts,
|
|
_add_last_ts_value=_add_last_ts_value,
|
|
)
|
|
|
|
return data
|
|
|
|
def __add__(
|
|
self, other: Union[List, "InfiniteLookbackBuffer", int, float, complex]
|
|
) -> "InfiniteLookbackBuffer":
|
|
"""Adds another InfiniteLookbackBuffer object or list to the end of this one.
|
|
|
|
Args:
|
|
other: Another `InfiniteLookbackBuffer` or a `list` or a number.
|
|
If a `InfiniteLookbackBuffer` its data (w/o its lookback buffer) gets
|
|
concatenated to self's data. If a `list`, we concat it to self's data.
|
|
If a number, we add this number to each element of self (if possible).
|
|
|
|
Returns:
|
|
A new `InfiniteLookbackBuffer` instance `self.data` containing
|
|
concatenated data from `self` and `other` (or adding `other` to each element
|
|
in self's data).
|
|
"""
|
|
|
|
if self.finalized:
|
|
raise RuntimeError(f"Cannot `add` to a finalized {type(self).__name__}.")
|
|
else:
|
|
# If `other` is an int, simply add it to all our values (if possible) and
|
|
# use the result as the underlying data for the returned buffer.
|
|
if isinstance(other, (int, float, complex)):
|
|
data = [
|
|
(d + other) if isinstance(d, (int, float, complex)) else d
|
|
for d in self.data
|
|
]
|
|
# If `other` is a InfiniteLookbackBuffer itself, do NOT include its
|
|
# lookback buffer anymore. We assume that `other`'s lookback buffer i
|
|
# already at the end of `self`.
|
|
elif isinstance(other, InfiniteLookbackBuffer):
|
|
data = self.data + other.data[other.lookback :]
|
|
# `other` is a list, simply concat the two lists and use the result as
|
|
# the underlying data for the returned buffer.
|
|
else:
|
|
data = self.data + other
|
|
|
|
return InfiniteLookbackBuffer(
|
|
data=data,
|
|
lookback=self.lookback,
|
|
space=self.space,
|
|
)
|
|
|
|
def __getitem__(self, item):
|
|
"""Support squared bracket syntax, e.g. buffer[:5]."""
|
|
return self.get(item)
|
|
|
|
def __setitem__(self, key, value):
|
|
self.set(new_data=value, at_indices=key)
|
|
|
|
def set(
|
|
self,
|
|
new_data,
|
|
*,
|
|
at_indices: Optional[Union[int, slice, List[int]]] = None,
|
|
neg_index_as_lookback: bool = False,
|
|
) -> None:
|
|
"""Overwrites all or some of the data in this buffer with the provided data.
|
|
|
|
Args:
|
|
new_data: The new data to overwrite existing records with.
|
|
at_indices: A single int is interpreted as an index, at which to overwrite
|
|
the individual record stored at this index with `new_data`.
|
|
A list of ints is interpreted as a list of indices, which to overwrite
|
|
with `new_data`, which must be a batch of size `len(at_indices)`.
|
|
A slice object is interpreted as a range, which to overwrite with
|
|
`new_data`. 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, a buffer with data [4, 5, 6, 7, 8, 9],
|
|
where [4, 5, 6] is the lookback buffer range (ts=0 item is 7), will
|
|
handle a call `set(99, at_indices=-1, neg_index_as_lookback=True)`
|
|
with `6` being replaced by 99 and to `set([98, 99, 100],
|
|
at_indices=slice(-2, 1), neg_index_as_lookback=True)` with
|
|
`[5, 6, 7]` being replaced by `[98, 99, 100]`.
|
|
"""
|
|
# `at_indices` is None -> Override all our data (excluding the lookback buffer).
|
|
if at_indices is None:
|
|
self._set_all_data(new_data)
|
|
|
|
elif isinstance(at_indices, slice):
|
|
self._set_slice(
|
|
new_data,
|
|
slice_=at_indices,
|
|
neg_index_as_lookback=neg_index_as_lookback,
|
|
)
|
|
elif isinstance(at_indices, list):
|
|
for i, idx in enumerate(at_indices):
|
|
self._set_int_index(
|
|
new_data[i],
|
|
idx=idx,
|
|
neg_index_as_lookback=neg_index_as_lookback,
|
|
)
|
|
else:
|
|
assert isinstance(at_indices, int)
|
|
self._set_int_index(
|
|
new_data,
|
|
idx=at_indices,
|
|
neg_index_as_lookback=neg_index_as_lookback,
|
|
)
|
|
|
|
def __len__(self):
|
|
"""Return the length of our data, excluding the lookback buffer."""
|
|
len_ = self.len_incl_lookback()
|
|
# Only count the data after the lookback.
|
|
return max(len_ - self.lookback, 0)
|
|
|
|
def len_incl_lookback(self):
|
|
if self.finalized:
|
|
if isinstance(self.data, np.ndarray):
|
|
return len(self.data)
|
|
return len(tree.flatten(self.data)[0])
|
|
else:
|
|
return len(self.data)
|
|
|
|
def __repr__(self):
|
|
return (
|
|
f"{type(self).__name__}({self.data[:self.lookback]} <- "
|
|
f"lookback({self.lookback}) | {self.data[self.lookback:]})"
|
|
)
|
|
|
|
def _get_all_data(self, one_hot_discrete=False, _ignore_last_ts=False):
|
|
data = self[: (None if not _ignore_last_ts else -1)]
|
|
if one_hot_discrete:
|
|
data = self._one_hot(data, space_struct=self.space_struct)
|
|
return data
|
|
|
|
def _set_all_data(self, new_data):
|
|
self._set_slice(new_data, slice(0, None))
|
|
|
|
def _get_slice(
|
|
self,
|
|
slice_,
|
|
fill=None,
|
|
neg_index_as_lookback=False,
|
|
one_hot_discrete=False,
|
|
_ignore_last_ts=False,
|
|
_add_last_ts_value=None,
|
|
):
|
|
# Fast path: finalized simple numpy array with no special options.
|
|
if (
|
|
self.finalized
|
|
and isinstance(self.data, np.ndarray)
|
|
and fill is None
|
|
and not one_hot_discrete
|
|
and not _ignore_last_ts
|
|
and _add_last_ts_value is None
|
|
):
|
|
start = slice_.start
|
|
stop = slice_.stop
|
|
step = slice_.step
|
|
lb = self.lookback
|
|
# Ultra-fast inline for the common case: positive-only, no step.
|
|
if (
|
|
not neg_index_as_lookback
|
|
and step is None
|
|
and (start is None or start >= 0)
|
|
and (stop is None or stop >= 0)
|
|
):
|
|
abs_start = lb if start is None else lb + start
|
|
abs_stop = len(self.data) if stop is None else lb + stop
|
|
return self.data[abs_start:abs_stop]
|
|
adj_slice, _, _, _ = self._interpret_slice(
|
|
slice_,
|
|
neg_index_as_lookback,
|
|
len_self_plus_lookback=len(self.data),
|
|
)
|
|
return self.data[adj_slice]
|
|
|
|
data_to_use = self.data
|
|
if _ignore_last_ts:
|
|
if self.finalized:
|
|
data_to_use = tree.map_structure(lambda s: s[:-1], self.data)
|
|
else:
|
|
data_to_use = self.data[:-1]
|
|
if _add_last_ts_value is not None:
|
|
if self.finalized:
|
|
data_to_use = tree.map_structure(
|
|
lambda s, t: np.append(s, t),
|
|
data_to_use.copy(),
|
|
_add_last_ts_value,
|
|
)
|
|
else:
|
|
data_to_use = np.append(data_to_use.copy(), _add_last_ts_value)
|
|
|
|
slice_, slice_len, fill_left_count, fill_right_count = self._interpret_slice(
|
|
slice_,
|
|
neg_index_as_lookback,
|
|
len_self_plus_lookback=(
|
|
self.len_incl_lookback()
|
|
+ int(_add_last_ts_value is not None)
|
|
- int(_ignore_last_ts)
|
|
),
|
|
)
|
|
|
|
# Perform the actual slice.
|
|
data_slice = None
|
|
if slice_len > 0:
|
|
if self.finalized:
|
|
data_slice = tree.map_structure(lambda s: s[slice_], data_to_use)
|
|
else:
|
|
data_slice = data_to_use[slice_]
|
|
|
|
if one_hot_discrete:
|
|
data_slice = self._one_hot(data_slice, space_struct=self.space_struct)
|
|
|
|
# Data is shorter than the range requested -> Fill the rest with `fill` data.
|
|
if fill is not None and (fill_right_count > 0 or fill_left_count > 0):
|
|
if self.finalized:
|
|
if fill_left_count:
|
|
if self.space is None:
|
|
fill_batch = np.array([fill] * fill_left_count)
|
|
else:
|
|
fill_batch = get_dummy_batch_for_space(
|
|
self.space,
|
|
fill_value=fill,
|
|
batch_size=fill_left_count,
|
|
one_hot_discrete=one_hot_discrete,
|
|
)
|
|
if data_slice is not None:
|
|
data_slice = tree.map_structure(
|
|
lambda s0, s: np.concatenate([s0, s]),
|
|
fill_batch,
|
|
data_slice,
|
|
)
|
|
else:
|
|
data_slice = fill_batch
|
|
if fill_right_count:
|
|
if self.space is None:
|
|
fill_batch = np.array([fill] * fill_right_count)
|
|
else:
|
|
fill_batch = get_dummy_batch_for_space(
|
|
self.space,
|
|
fill_value=fill,
|
|
batch_size=fill_right_count,
|
|
one_hot_discrete=one_hot_discrete,
|
|
)
|
|
if data_slice is not None:
|
|
data_slice = tree.map_structure(
|
|
lambda s0, s: np.concatenate([s, s0]),
|
|
fill_batch,
|
|
data_slice,
|
|
)
|
|
else:
|
|
data_slice = fill_batch
|
|
|
|
else:
|
|
if self.space is None:
|
|
fill_batch = [fill]
|
|
else:
|
|
fill_batch = [
|
|
get_dummy_batch_for_space(
|
|
self.space,
|
|
fill_value=fill,
|
|
batch_size=0,
|
|
one_hot_discrete=one_hot_discrete,
|
|
)
|
|
]
|
|
data_slice = (
|
|
fill_batch * fill_left_count
|
|
+ (data_slice if data_slice is not None else [])
|
|
+ fill_batch * fill_right_count
|
|
)
|
|
|
|
if data_slice is None:
|
|
if self.finalized:
|
|
return tree.map_structure(lambda s: s[slice_], data_to_use)
|
|
else:
|
|
return data_to_use[slice_]
|
|
return data_slice
|
|
|
|
def _set_slice(
|
|
self,
|
|
new_data,
|
|
slice_,
|
|
neg_index_as_lookback=False,
|
|
):
|
|
slice_, _, _, _ = self._interpret_slice(slice_, neg_index_as_lookback)
|
|
|
|
# Check, whether the setting to new_data changes the length of self
|
|
# (it shouldn't). If it does, raise an error.
|
|
try:
|
|
if self.finalized:
|
|
|
|
def __set(s, n):
|
|
if self.space:
|
|
assert self.space.contains(n[0])
|
|
assert len(s[slice_]) == len(n)
|
|
s[slice_] = n
|
|
|
|
tree.map_structure(__set, self.data, new_data)
|
|
else:
|
|
assert len(self.data[slice_]) == len(new_data)
|
|
self.data[slice_] = new_data
|
|
except AssertionError:
|
|
raise IndexError(
|
|
f"Cannot `set()` value via at_indices={slice_} (option "
|
|
f"neg_index_as_lookback={neg_index_as_lookback})! Slice of data "
|
|
"does NOT have the same size as `new_data`."
|
|
)
|
|
|
|
def _get_int_index(
|
|
self,
|
|
idx: int,
|
|
fill=None,
|
|
neg_index_as_lookback=False,
|
|
one_hot_discrete=False,
|
|
_ignore_last_ts=False,
|
|
_add_last_ts_value=None,
|
|
):
|
|
# Fast path: finalized simple numpy array with no special options.
|
|
if (
|
|
self.finalized
|
|
and isinstance(self.data, np.ndarray)
|
|
and fill is None
|
|
and not one_hot_discrete
|
|
and not _ignore_last_ts
|
|
and _add_last_ts_value is None
|
|
):
|
|
actual_idx = (
|
|
(self.lookback + idx) if (idx >= 0 or neg_index_as_lookback) else idx
|
|
)
|
|
return self.data[actual_idx]
|
|
|
|
data_to_use = self.data
|
|
if _ignore_last_ts:
|
|
if self.finalized:
|
|
data_to_use = tree.map_structure(lambda s: s[:-1], self.data)
|
|
else:
|
|
data_to_use = self.data[:-1]
|
|
if _add_last_ts_value is not None:
|
|
if self.finalized:
|
|
data_to_use = tree.map_structure(
|
|
lambda s, last: np.append(s, last), data_to_use, _add_last_ts_value
|
|
)
|
|
else:
|
|
data_to_use = data_to_use.copy()
|
|
data_to_use.append(_add_last_ts_value)
|
|
|
|
# If index >= 0 -> Ignore lookback buffer.
|
|
# Otherwise, include lookback buffer.
|
|
if idx >= 0 or neg_index_as_lookback:
|
|
idx = self.lookback + idx
|
|
# Negative indices mean: Go to left into lookback buffer starting from idx=0.
|
|
# But if we pass the lookback buffer, the index should be invalid and we will
|
|
# have to fill, if required. Invalidate the index by setting it to one larger
|
|
# than max.
|
|
if neg_index_as_lookback and idx < 0:
|
|
idx = len(self) + self.lookback - (_ignore_last_ts is True)
|
|
|
|
try:
|
|
if self.finalized:
|
|
data = tree.map_structure(lambda s: s[idx], data_to_use)
|
|
else:
|
|
data = data_to_use[idx]
|
|
# Out of range index -> If `fill`, use a fill dummy (B=0), if not, error out.
|
|
except IndexError as e:
|
|
if fill is not None:
|
|
if self.space is None:
|
|
return fill
|
|
return get_dummy_batch_for_space(
|
|
self.space,
|
|
fill_value=fill,
|
|
batch_size=0,
|
|
one_hot_discrete=one_hot_discrete,
|
|
)
|
|
else:
|
|
raise e from ValueError(f"Trying to get index {idx} from {data_to_use}")
|
|
|
|
# Convert discrete/multi-discrete components to one-hot vectors, if required.
|
|
if one_hot_discrete:
|
|
data = self._one_hot(data, self.space_struct)
|
|
return data
|
|
|
|
def _set_int_index(self, new_data, idx, neg_index_as_lookback):
|
|
actual_idx = idx
|
|
# If index >= 0 -> Ignore lookback buffer.
|
|
# Otherwise, include lookback buffer.
|
|
if actual_idx >= 0 or neg_index_as_lookback:
|
|
actual_idx = self.lookback + actual_idx
|
|
# Negative indices mean: Go to left into lookback buffer starting from idx=0.
|
|
# But if we pass the lookback buffer, the index should be invalid and we will
|
|
# have to fill, if required. Invalidate the index by setting it to one larger
|
|
# than max.
|
|
if neg_index_as_lookback and actual_idx < 0:
|
|
actual_idx = len(self) + self.lookback
|
|
|
|
try:
|
|
if self.finalized:
|
|
|
|
def __set(s, n):
|
|
if self.space:
|
|
assert self.space.contains(n), n
|
|
s[actual_idx] = n
|
|
|
|
tree.map_structure(__set, self.data, new_data)
|
|
else:
|
|
self.data[actual_idx] = new_data
|
|
except IndexError:
|
|
raise IndexError(
|
|
f"Cannot `set()` value at index {idx} (option "
|
|
f"neg_index_as_lookback={neg_index_as_lookback})! Out of range "
|
|
f"of buffer data."
|
|
)
|
|
|
|
def _interpret_slice(
|
|
self,
|
|
slice_,
|
|
neg_index_as_lookback,
|
|
len_self_plus_lookback=None,
|
|
):
|
|
if len_self_plus_lookback is None:
|
|
len_self_plus_lookback = len(self) + self.lookback
|
|
|
|
# Re-interpret slice bounds as absolute positions (>=0) within our
|
|
# internal data.
|
|
start = slice_.start
|
|
stop = slice_.stop
|
|
|
|
# Start is None -> Exclude lookback buffer.
|
|
if start is None:
|
|
start = self.lookback
|
|
# Start is negative.
|
|
elif start < 0:
|
|
# `neg_index_as_lookback=True` -> User wants to index into the lookback
|
|
# range.
|
|
if neg_index_as_lookback:
|
|
start = self.lookback + start
|
|
# Interpret index as counting "from end".
|
|
else:
|
|
start = len_self_plus_lookback + start
|
|
# Start is 0 or positive -> timestep right after lookback is interpreted as 0.
|
|
else:
|
|
start = self.lookback + start
|
|
|
|
# Stop is None -> Set stop to very last index + 1 of our internal data.
|
|
if stop is None:
|
|
stop = len_self_plus_lookback
|
|
# Stop is negative.
|
|
elif stop < 0:
|
|
# `neg_index_as_lookback=True` -> User wants to index into the lookback
|
|
# range. Set to 0 (beginning of lookback buffer) if result is a negative
|
|
# index.
|
|
if neg_index_as_lookback:
|
|
stop = self.lookback + stop
|
|
# Interpret index as counting "from end". Set to 0 (beginning of actual
|
|
# episode) if result is a negative index.
|
|
else:
|
|
stop = len_self_plus_lookback + stop
|
|
# Stop is positive -> Add lookback range to it.
|
|
else:
|
|
stop = self.lookback + stop
|
|
|
|
fill_left_count = fill_right_count = 0
|
|
# Both start and stop are on left side.
|
|
if start < 0 and stop < 0:
|
|
fill_left_count = abs(start - stop)
|
|
fill_right_count = 0
|
|
start = stop = 0
|
|
# Both start and stop are on right side.
|
|
elif start >= len_self_plus_lookback and stop >= len_self_plus_lookback:
|
|
fill_right_count = abs(start - stop)
|
|
fill_left_count = 0
|
|
start = stop = len_self_plus_lookback
|
|
# Set to 0 (beginning of actual episode) if result is a negative index.
|
|
elif start < 0:
|
|
fill_left_count = -start
|
|
start = 0
|
|
elif stop >= len_self_plus_lookback:
|
|
fill_right_count = stop - len_self_plus_lookback
|
|
stop = len_self_plus_lookback
|
|
# Only `stop` might be < 0, when slice has negative step and start is > 0.
|
|
elif stop < 0:
|
|
if start >= len_self_plus_lookback:
|
|
fill_left_count = start - len_self_plus_lookback + 1
|
|
start = len_self_plus_lookback - 1
|
|
fill_right_count = -stop - 1
|
|
stop = -LARGE_INTEGER
|
|
|
|
assert start >= 0 and (stop >= 0 or stop == -LARGE_INTEGER), (start, stop)
|
|
|
|
step = slice_.step if slice_.step is not None else 1
|
|
slice_ = slice(start, stop, step)
|
|
slice_len = max(0, (stop - start + (step - (1 if step > 0 else -1))) // step)
|
|
return slice_, slice_len, fill_left_count, fill_right_count
|
|
|
|
def _one_hot(self, data, space_struct):
|
|
if space_struct is None:
|
|
raise ValueError(
|
|
f"Cannot `one_hot` data in `{type(self).__name__}` if a "
|
|
"gym.Space was NOT provided during construction!"
|
|
)
|
|
|
|
def _convert(dat_, space):
|
|
if isinstance(space, gym.spaces.Discrete):
|
|
return one_hot(dat_, depth=space.n)
|
|
elif isinstance(space, gym.spaces.MultiDiscrete):
|
|
return one_hot_multidiscrete(dat_, depths=space.nvec)
|
|
return dat_
|
|
|
|
if isinstance(data, list):
|
|
data = [
|
|
tree.map_structure(_convert, dslice, space_struct) for dslice in data
|
|
]
|
|
else:
|
|
data = tree.map_structure(_convert, data, space_struct)
|
|
return data
|