def test_init(self, bounds_fixture):
assert bounds_fixture.dtype is not None
assert judo.is_tensor(bounds_fixture.low)
assert judo.is_tensor(bounds_fixture.high)
assert isinstance(bounds_fixture.shape, tuple)
assert bounds_fixture.low.shape == bounds_fixture.shape
assert bounds_fixture.high.shape == bounds_fixture.shape
def copy(self) -> "States":
"""Crete a copy of the current instance."""
param_dict = {
str(name): judo.copy(val) if judo.is_tensor(val) else copy.deepcopy(val)
for name, val in self.items()
}
return States(batch_size=self.n, **param_dict)
def get_params_dict(self) -> StateDict:
"""Return a dictionary describing the data stored in the :class:`States`."""
return {
k: {"shape": v.shape, "dtype": v.dtype}
for k, v in self.__dict__.items()
if judo.is_tensor(v)
}
def _ix(self, index: int):
# TODO(guillemdb): Allow slicing
data = {
k: API.unsqueeze(v[index], 0) if judo.is_tensor(v) else v
for k, v in self.items()
}
return self.__class__(batch_size=1, **data)
def merge_one_name(states_list, name):
vals = []
for state in states_list:
data = state[name]
# Attributes that are not tensors are not stacked.
if not judo.is_tensor(data):
return data
state_len = len(state)
if len(data.shape) == 0 and state_len == 1:
# Name is scalar vector. Data is typing.Scalar value. Transform to array first
value = tensor([data]).flatten()
elif len(data.shape) == 1 and state_len == 1:
if data.shape[0] == 1:
# Name is typing.Scalar vector. Data already transformed to an array
value = data
else:
# Name is a matrix of vectors. Data needs an additional dimension
value = tensor([data])
elif len(data.shape) == 1 and state_len > 1:
# Name is a typing.Scalar vector. Data already has is a one dimensional array
value = data
elif (len(data.shape) > 1 and state_len > 1
or len(data.shape) > 1 and len(state) == 1):
# Name is a matrix of vectors. Data has the correct shape
value = data
else:
raise ValueError(
"Could not infer data concatenation for attribute %s with shape %s"
% (name, data.shape), )
vals.append(value)
return API.concatenate(vals)
def from_bounds_params(
cls,
function: Callable,
shape: tuple = None,
high: Union[int, float, typing.Tensor] = numpy.inf,
low: Union[int, float, typing.Tensor] = numpy.NINF,
custom_domain_check: Callable[[typing.Tensor], typing.Tensor] = None,
) -> "Function":
"""
Initialize a function defining its shape and bounds without using a :class:`Bounds`.
Args:
function: Callable that takes a batch of vectors (batched across \
the first dimension of the array) and returns a vector of \
typing.Scalar. This function is applied to a batch of walker \
observations.
shape: Input shape of the solution vector without taking into account \
the batch dimension. For example, a two dimensional function \
applied to a batch of 5 walkers will have shape=(2,), even though
the observations will have shape (5, 2)
high: Upper bound of the function domain. If it's an typing.Scalar it will \
be the same for all dimensions. If its a numpy array it will \
be the upper bound for each dimension.
low: Lower bound of the function domain. If it's an typing.Scalar it will \
be the same for all dimensions. If its a numpy array it will \
be the lower bound for each dimension.
custom_domain_check: Callable that checks points inside the bounds \
to know if they are in a custom domain when it is not just \
a set of rectangular bounds.
Returns:
:class:`Function` with its :class:`Bounds` created from the provided arguments.
"""
if (not (judo.is_tensor(high) or isinstance(high, (list, tuple)))
and not (judo.is_tensor(low) or isinstance(low, (list, tuple)))
and shape is None):
raise TypeError(
"Need to specify shape or high or low must be a numpy array.")
bounds = Bounds(high=high, low=low, shape=shape)
return Function(function=function,
bounds=bounds,
custom_domain_check=custom_domain_check)
def split_kwargs_in_chunks(
kwargs: Dict[str, Union[list, Tensor]],
n_chunks: int,
allow_size_1: bool = True
) -> Generator[Dict[str, Union[list, Tensor]], None, None]:
"""Split the kwargs passed to ``make_transitions`` in similar batches."""
n_values = len(next(iter(
kwargs.values()))) # Assumes all data have the same len
chunk_size = int(numpy.ceil(n_values / n_chunks))
for start, end in similiar_chunks_indexes(n_values,
n_chunks,
allow_size_1=allow_size_1):
if start + chunk_size >= n_values - 2: # Do not allow the last chunk to have size 1
yield {
k: v[start:n_values] if judo.is_tensor(v) else v
for k, v in kwargs.items()
}
break
else:
yield {
k: v[start:end] if judo.is_tensor(v) else v
for k, v in kwargs.items()
}
def iteritems(self):
"""
Iterate the states attributes by walker.
Returns:
Tuple containing all the names of the attributes, and the values that
correspond to a given walker.
"""
if self.n < 1:
return self.vals()
for i in range(self.n):
values = (v[i] if judo.is_tensor(v) else v for v in self.vals())
yield tuple(self._names), tuple(values)
def __init__(self, batch_size: int, state_dict: Optional[StateDict] = None, **kwargs):
"""
Initialize a :class:`States`.
Args:
batch_size: The number of items in the first dimension of the tensors.
state_dict: Dictionary defining the attributes of the tensors.
**kwargs: The name-tensor pairs can also be specified as kwargs.
"""
attr_dict = self.params_to_arrays(state_dict, batch_size) if state_dict is not None else {}
attr_dict.update(kwargs)
self._tensor_names = [k for k, v in attr_dict.items() if judo.is_tensor(v)]
self._names = list(attr_dict.keys())
self._attr_dict = attr_dict
self.update(**self._attr_dict)
self._batch_size = batch_size
def split_states(self, n_chunks: int) -> Generator["States", None, None]:
"""
Return a generator for n_chunks different states, where each one \
contain only the data corresponding to one walker.
"""
def get_chunck_size(state, start, end):
for name in state._names:
attr = state[name]
if judo.is_tensor(attr):
return len(attr[start:end])
return int(numpy.ceil(self.n / n_chunks))
for start, end in judo.similiar_chunks_indexes(self.n, n_chunks):
chunk_size = get_chunck_size(self, start, end)
data = {k: val[start:end] if judo.is_tensor(val) else val for k, val in self.items()}
new_state = self.__class__(batch_size=chunk_size, **data)
yield new_state
def clone(
self,
will_clone: Tensor,
compas_ix: Tensor,
ignore: Optional[Set[str]] = None,
):
"""
Clone all the stored data according to the provided arrays.
Args:
will_clone: Array of shape (n_walkers,) of booleans indicating the \
index of the walkers that will clone to a random companion.
compas_ix: Array of integers of shape (n_walkers,). Contains the \
indexes of the walkers that will be copied.
ignore: set containing the names of the attributes that will not be \
cloned.
"""
ignore = set() if ignore is None else ignore
for name in self.keys():
if judo.is_tensor(self[name]) and name not in ignore:
self[name][will_clone] = self[name][compas_ix][will_clone]
def get_chunck_size(state, start, end):
for name in state._names:
attr = state[name]
if judo.is_tensor(attr):
return len(attr[start:end])
return int(numpy.ceil(self.n / n_chunks))
def __init__(
self,
high: Union[Tensor, Scalar] = numpy.inf,
low: Union[Tensor, Scalar] = numpy.NINF,
shape: Optional[tuple] = None,
dtype: Optional[type] = None,
):
"""
Initialize a :class:`Bounds`.
Args:
high: Higher value for the bound interval. If it is an typing_.Scalar \
it will be applied to all the coordinates of a target vector. \
If it is a vector, the bounds will be checked coordinate-wise. \
It defines and closed interval.
low: Lower value for the bound interval. If it is a typing_.Scalar it \
will be applied to all the coordinates of a target vector. \
If it is a vector, the bounds will be checked coordinate-wise. \
It defines and closed interval.
shape: Shape of the array that will be bounded. Only needed if `high` and `low` are \
vectors and it is used to define the dimensions that will be bounded.
dtype: Data type of the array that will be bounded. It can be inferred from `high` \
or `low` (the type of `high` takes priority).
Examples:
Initializing :class:`Bounds` using numpy arrays:
>>> import numpy
>>> high, low = numpy.ones(3, dtype=float), -1 * numpy.ones(3, dtype=int)
>>> bounds = Bounds(high=high, low=low)
>>> print(bounds)
Bounds shape float64 dtype (3,) low [-1 -1 -1] high [1. 1. 1.]
Initializing :class:`Bounds` using typing_.Scalars:
>>> import numpy
>>> high, low, shape = 4, 2.1, (5,)
>>> bounds = Bounds(high=high, low=low, shape=shape)
>>> print(bounds)
Bounds shape float64 dtype (5,) low [2.1 2.1 2.1 2.1 2.1] high [4. 4. 4. 4. 4.]
"""
# Infer shape if not specified
if shape is None and hasattr(high, "shape"):
shape = high.shape
elif shape is None and hasattr(low, "shape"):
shape = low.shape
elif shape is None:
raise TypeError(
"If shape is None high or low need to have .shape attribute.")
# High and low will be arrays of target shape
if not judo.is_tensor(high):
high = tensor(high) if isinstance(
high, _Iterable) else API.ones(shape) * high
if not judo.is_tensor(low):
low = tensor(low) if isinstance(
low, _Iterable) else API.ones(shape) * low
self.high = judo.astype(high, dtype)
self.low = judo.astype(low, dtype)
if dtype is not None:
self.dtype = dtype
elif hasattr(high, "dtype"):
self.dtype = high.dtype
elif hasattr(low, "dtype"):
self.dtype = low.dtype
else:
self.dtype = type(high) if high is not None else type(low)
def __init__(self,
state: Tensor,
observ: Tensor,
reward: Scalar,
id_walker=None,
time=0.0,
state_dict: StateDict = None,
**kwargs):
"""
Initialize a :class:`OneWalker`.
Args:
state: Non batched numpy array defining the state of the walker.
observ: Non batched numpy array defining the observation of the walker.
reward: typing.Scalar value representing the reward of the walker.
id_walker: Hash of the provided State. If None it will be calculated when the
the :class:`OneWalker` is initialized.
state_dict: External :class:`typing.StateDict` that overrides the default values.
time: Time step of the current walker. Measures the length of the path followed \
by the walker.
**kwargs: Additional data needed to define the walker. Its structure \
needs to be defined in the provided ``state_dict``. These attributes
will be assigned to the :class:`EnvStates` of the :class:`Swarm`.
"""
self.id_walkers = None
self.rewards = None
self.observs = None
self.states = None
self.times = None
self._observs_size = observ.shape
self._observs_dtype = observ.dtype
self._states_size = state.shape
self._states_dtype = state.dtype
self._rewards_dtype = tensor(reward).dtype
# Accept external definition of param_dict values
walkers_dict = self.get_params_dict()
if state_dict is not None:
for k, v in state_dict.items():
if k in ["observs", "states"
]: # These two are parsed from the provided opts
continue
if k in walkers_dict:
walkers_dict[k] = v
super(OneWalker, self).__init__(batch_size=1, state_dict=walkers_dict)
# Keyword arguments must be defined in state_dict
if state_dict is not None:
for k in kwargs.keys():
if k not in state_dict:
raise ValueError(
"The provided attributes must be defined in state_dict."
"param_dict: %s\n kwargs: %s" % (state_dict, kwargs))
self.observs[:] = judo.copy(observ)
self.states[:] = judo.copy(state)
self.rewards[:] = judo.copy(reward) if judo.is_tensor(
reward) else copy.deepcopy(reward)
self.times[:] = judo.copy(time) if judo.is_tensor(
time) else copy.deepcopy(time)
self.id_walkers[:] = (judo.copy(id_walker.squeeze()) if id_walker
is not None else hasher.hash_tensor(state))
self.update(**kwargs)
