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)