def add_states(self, parent_ids: List[NodeId], n_iter: int = None, **kwargs): """ Update the history of the tree adding the necessary data to recreate a \ the trajectories sampled by the :class:`Swarm`. Args: parent_ids: List of states hashes representing the parent nodes of \ the current states. n_iter: Number of iteration of the algorithm when the data was sampled. kwargs: Keyword arguments representing different :class:`States` instances. Returns: None """ ids, id_states = self.get_states_ids(**kwargs) leaf_ids, parent_ids = judo.to_numpy(ids), judo.to_numpy(parent_ids) # Keep track of nodes that are active to make sure that they are not pruned self.last_added = set(leaf_ids) | set(parent_ids) for i, (leaf, parent) in enumerate(zip(leaf_ids, parent_ids)): node_data, edge_data = self._extract_data_from_states(index=i, **kwargs) self.append_leaf( leaf_id=leaf, parent_id=parent, node_data=node_data, edge_data=edge_data, epoch=n_iter, ) if not hasher.uses_true_hash: id_states.update(ids=leaf_ids)
def minimize(self, x: typing.Tensor): """ Apply ``scipy.optimize.minimize`` to a single point. Args: x: Array representing a single point of the function to be minimized. Returns: Optimization result object returned by ``scipy.optimize.minimize``. """ def _optimize(_x): try: _x = _x.reshape((1, ) + _x.shape) y = self.function(_x) except (ZeroDivisionError, RuntimeError): y = numpy.inf return y bounds = ScipyBounds( ub=judo.to_numpy(self.bounds.high) if self.bounds is not None else None, lb=judo.to_numpy(self.bounds.low) if self.bounds is not None else None, ) return minimize(_optimize, x, bounds=bounds, *self.args, **self.kwargs)
def _update_lims(self, swarm, X: numpy.ndarray): backup_bounds = swarm.env.bounds if swarm is not None else Bounds.from_array( X) bounds = (swarm.critic.bounds if has_embedding(swarm) and self.use_embeddings else backup_bounds) self.xlim, self.ylim = bounds.to_tuples()[:2] self.xlim = judo.to_numpy(self.xlim[0]), judo.to_numpy(self.xlim[1]) self.ylim = judo.to_numpy(self.ylim[0]), judo.to_numpy(self.ylim[1])
def statistics_from_array(x: numpy.ndarray): """Return the (mean, std, max, min) of an array.""" try: return ( judo.to_numpy(x).mean(), judo.to_numpy(x).std(), judo.to_numpy(x).max(), judo.to_numpy(x).min(), ) except AttributeError: return numpy.nan, numpy.nan, numpy.nan, numpy.nan
def get_plot_data(self, data: Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]): """Create the meshgrid needed to interpolate the target data points.""" x, y, z = (data[:, 0], data[:, 1], data[:, 2]) if isinstance(data, numpy.ndarray) else data x, y, z = judo.to_numpy(x), judo.to_numpy(y), judo.to_numpy(z) # target grid to interpolate to xi = numpy.linspace(x.min(), x.max(), self.n_points) yi = numpy.linspace(x.min(), x.max(), self.n_points) xx, yy = numpy.meshgrid(xi, yi) return x, y, xx, yy, z, self.xlim, self.ylim
def iterate_path_data( self, node_ids: Union[Tuple[NodeId], List[NodeId]], batch_size: int = None, names: NamesData = None, ) -> NodeDataGenerator: """ Return a generator that yields the data of the nodes contained in the provided path. Args: node_ids: Ids of the nodes of the path that will be sampled. The \ nodes must be provided in order, starting with the first \ node of the path. batch_size: If it is not None, the generator will return batches of \ data with the target batch size. If Batch size is less than 1 \ it will return a single batch containing all the data. names: Names of the data attributes that will be yielded by the generator. Returns: Generator providing the data corresponding to a path in the internal tree. """ names = self._validate_names(names) return_children = any(name.startswith(self.next_prefix) for name in names) path_generator = self.path_data_generator( node_ids=judo.to_numpy(node_ids), return_children=return_children, ) return self._generate_batches(path_generator, names=names, batch_size=batch_size)
def minimize_batch( self, x: typing.Tensor) -> Tuple[typing.Tensor, typing.Tensor]: """ Minimize a batch of points. Args: x: Array representing a batch of points to be optimized, stacked \ across the first dimension. Returns: Tuple of arrays containing the local optimum found for each point, \ and an array with the values assigned to each of the points found. """ x = judo.to_numpy(judo.copy(x)) with Backend.use_backend("numpy"): result = judo.zeros_like(x) rewards = judo.zeros((x.shape[0], 1)) for i in range(x.shape[0]): new_x, reward = self.minimize_point(x[i, :]) result[i, :] = new_x rewards[i, :] = float(reward) self.bounds.high = tensor(self.bounds.high) self.bounds.low = tensor(self.bounds.low) result, rewards = tensor(result), tensor(rewards) return result, rewards
def get_plot_data(self, swarm: Swarm = None) -> numpy.ndarray: """Extract the frame from the :class:`AtariEnv` that the target \ :class:`Swarm` contains.""" if swarm is None: return numpy.zeros((210, 160, 3)) state = swarm.get("best_state") state = judo.to_numpy(state) return self.image_from_state(swarm=swarm, state=state)
def get_xy_coords(swarm: Swarm, use_embedding: False) -> numpy.ndarray: """ Get the x,y coordinates to represent observations values in 2D. If the :class:`Swarm` has a critic to calculate 2D embeddings, the \ two first dimensions of embedding values will be used. Otherwise return \ the first two dimensions of ``observs``. """ if use_embedding and has_embedding(swarm): X = swarm.critic.preprocess_input( env_states=swarm.walkers.env_states, walkers_states=swarm.walkers.states, model_states=swarm.walkers.model_states, batch_size=swarm.walkers.n, ) return judo.to_numpy(X[:, :2]) elif isinstance(swarm, numpy.ndarray): return swarm return judo.to_numpy(swarm.walkers.env_states.observs[:, :2])
def lennard_jones(x: np.ndarray) -> np.ndarray: result = np.zeros(x.shape[0]) x = judo.to_numpy(x) assert isinstance(x, np.ndarray) for i in range(x.shape[0]): try: result[i] = _lennard_fast(x[i]) except ZeroDivisionError: result[i] = np.inf result = judo.as_tensor(result) return result
def __repr__(self): with numpy.printoptions(linewidth=100, threshold=200, edgeitems=9): init_actions = judo.to_numpy( self.internal_swarm.walkers.states.init_actions.flatten()) with Backend.use_backend("numpy"): y = numpy.bincount(init_actions.astype(int)) ii = numpy.nonzero(y)[0] string = str(self.root_walker) string += "\n Init actions [action, count]: \n%s" % numpy.vstack( (ii, y[ii])).T return string
def update_tree(self, parent_ids: List[int]) -> None: """ Add a list of walker states represented by `states_ids` to the :class:`Tree`. Args: parent_ids: list containing the ids of the parents of the new states added. """ if self.tree is not None: self.tree.add_states( parent_ids=judo.to_numpy(parent_ids), env_states=self.walkers.env_states, model_states=self.walkers.model_states, walkers_states=self.walkers.states, n_iter=int(self.walkers.epoch), )
def resize_image(frame: Tensor, width: int, height: int, mode: str = "RGB") -> Tensor: """ Use PIL to resize an RGB frame to an specified height and width. Args: frame: Target numpy array representing the image that will be resized. width: Width of the resized image. height: Height of the resized image. mode: Passed to Image.convert. Returns: The resized frame that matches the provided width and height. """ from PIL import Image frame = judo.to_numpy(frame) with judo.Backend.use_backend(name="numpy"): frame = Image.fromarray(frame) frame = judo.to_numpy(frame.convert(mode).resize(size=(width, height))) return judo.tensor(frame)
def predict( self, root_env_states: StatesEnv, walkers: StepWalkers, ) -> StatesModel: """ Select the most frequent ``init_action`` assigned to the internal swarm's walkers. The selected ``dt`` will be equal to the minimum ``init_dts`` among all \ the walkers that sampled the selected ``init_action``. Args: root_env_states: :env-st:`StatesEnv` class containing the data \ corresponding to the root walker of a :class:`StepSwarm`. walkers: :walkers:`StepWalkers` used by the internal warm of a \ :class:`StepSwarm`. Returns: :class:`StatesModel` containing the ``actions`` and ``dt`` that the root walkers will use to step the :env:`Environment`. """ init_actions = judo.astype(walkers.states.init_actions.flatten(), judo.int) init_actions = judo.to_numpy(init_actions) with Backend.use_backend("numpy"): y = numpy.bincount(init_actions) most_used_action = numpy.nonzero(y)[0][0] most_used_action = tensor(most_used_action) root_model_states = StatesModel( batch_size=1, state_dict={ "actions": { "dtype": judo.int64 }, "dt": { "dtype": judo.int64 } }, ) root_model_states.actions[:] = most_used_action if hasattr(root_model_states, "dt"): init_dts = judo.astype(walkers.states.init_dts.flatten(), judo.int) index_dt = init_actions == most_used_action target_dt = init_dts[index_dt].min() root_model_states.dt[:] = target_dt return root_model_states
def image_from_state(swarm: Swarm, state: numpy.ndarray) -> numpy.ndarray: """ Return the frame corresponding to a given :class:`AtariEnv` state. Args: swarm: Swarm containing the target environment. state: States that will be used to extract the frame. Returns: Array of size (210, 160, 3) containing the RGB frame representing \ the target state. """ env = get_plangym_env(swarm) state = judo.to_numpy(state) env.set_state(state.astype(numpy.uint8)) env.step(0) return numpy.asarray(env.ale.getScreenRGB(), dtype=numpy.uint8)
def get_z_coords(self, swarm: Swarm, X: numpy.ndarray = None): """Get the values assigned by the :class:`Critic` to the regions of the state space.""" if swarm is None: return numpy.ones(self.n_points**self.n_points) if swarm.critic.bounds is None: swarm.critic.bounds = Bounds.from_array(X, scale=1.1) # target grid to interpolate to xi = numpy.linspace(swarm.critic.bounds.low[0], swarm.critic.bounds.high[0], self.n_points) yi = numpy.linspace(swarm.critic.bounds.low[1], swarm.critic.bounds.high[1], self.n_points) xx, yy = numpy.meshgrid(xi, yi) grid = numpy.c_[xx.ravel(), yy.ravel()] if swarm.swarm.critic.warmed: memory_values = swarm.swarm.critic.predict(grid) memory_values = relativize(-memory_values) else: memory_values = numpy.arange(grid.shape[0]) return judo.to_numpy(memory_values)
def get_plot_data(self, swarm: Swarm, attr: str): """ Extract the data of the attribute of the :class:`Swarm` that will be \ represented as a histogram. Args: swarm: Target :class:`Swarm`. attr: Attribute of the target :class:`States` that will be plotted. Returns: Histogram containing the target data. """ if swarm is None: return super(SwarmHistogram, self).get_plot_data(swarm) data = swarm.get(attr) if swarm is not None else numpy.arange(10) data = judo.to_numpy(data) self._update_lims(data) return super(SwarmHistogram, self).get_plot_data(data)
def get_plot_data( self, data: numpy.ndarray ) -> Tuple[Tuple[numpy.ndarray, numpy.ndarray], Tuple[Union[ Scalar, None], Union[Scalar, None]], ]: """ Calculate the histogram of the streamed data. Args: data: Values used to calculate the histogram. Returns: Tuple containing (values, bins), xlim. xlim is a tuple \ containing two typing_.Scalars that represent the limits of the x \ axis of the histogram. """ if data is None: data = numpy.arange(10) data = judo.to_numpy(data) data[numpy.isnan(data)] = 0.0 return numpy.histogram(data, self.n_bins), self.xlim
def get_z_coords(self, swarm: Swarm, X: numpy.ndarray = None) -> numpy.ndarray: """Extract the memory values of the :class:`Critic`'s grid.""" if swarm is None: return numpy.ones(self.n_points**self.n_points) if swarm.critic.bounds is None: swarm.critic.bounds = Bounds.from_array(X, scale=1.1) # target grid to interpolate to xi = numpy.linspace(swarm.critic.bounds.low[0], swarm.critic.bounds.high[0], self.n_points) yi = numpy.linspace(swarm.critic.bounds.low[1], swarm.critic.bounds.high[1], self.n_points) xx, yy = numpy.meshgrid(xi, yi) grid = numpy.c_[xx.ravel(), yy.ravel()] if swarm.swarm.critic.warmed: memory_values = swarm.swarm.critic.model.transform(grid) memory_values = numpy.array([ swarm.swarm.critic.memory[ix[0], ix[1]].astype(numpy.float32) for ix in memory_values.astype(int) ]) else: memory_values = numpy.arange(grid.shape[0]) return judo.to_numpy(memory_values)
def hash_torch(x): bytes = judo.to_numpy(x).tobytes() return xxhash.xxh32_intdigest(bytes)
def get_z_coords(self, swarm: Swarm, X: numpy.ndarray = None): """Return the normalized ``distances`` of the walkers.""" distances: numpy.ndarray = judo.to_numpy(swarm.get("distances")) return distances
def get_z_coords(self, swarm: Swarm, X: numpy.ndarray = None): """Return the normalized ``virtual_rewards`` of the walkers.""" virtual_rewards: numpy.ndarray = judo.to_numpy( swarm.get("virtual_rewards")) return virtual_rewards
def get_z_coords(self, swarm: Swarm, X: numpy.ndarray = None): """Return the normalized ``cum_rewards`` of the walkers.""" rewards: numpy.ndarray = judo.to_numpy( relativize(swarm.get("cum_rewards"))) return rewards
def get_plot_data(self, data): """Perform the necessary data wrangling for plotting the data.""" return judo.to_numpy(data)
def get_states_ids(self, walkers_states: StatesWalkers, **kwargs) -> Tuple[Tensor, StatesWalkers]: leaf_ids = judo.to_numpy(walkers_states.get("id_walkers")) return leaf_ids, walkers_states