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
