zea.agent.selection

Action selection strategies

These selection strategies implement a variety of policies for choosing which focused transmit to fire next, potentially given some beliefs about the state of tissue, as represented by particles.

For a comprehensive example usage, see: Active Perception for Focused Transmit Steering

All strategies are stateless, meaning that they do not maintain any internal state.

Classes

CovarianceSamplingLines(n_actions, ...[, ...])

Covariance sampling action selection.

EquispacedLines(n_actions, ...[, ...])

Equispaced lines action selection.

GreedyEntropy(n_actions, n_possible_actions, ...)

Greedy entropy action selection.

LinesActionModel(n_actions, ...)

Base class for action selection methods that select lines.

MaskActionModel()

Base class for any action selection method that does masking.

UniformRandomLines(n_actions, ...)

Uniform random lines action selection.

class zea.agent.selection.CovarianceSamplingLines(n_actions, n_possible_actions, img_width, img_height, seed=42, n_masks=200)[source]

Bases: LinesActionModel

Covariance sampling action selection.

This class models the line-to-line correlation to select the mask with the highest entropy.

random_uniform_lines(batch_size, seed=None)[source]

Wrapper around random_uniform_lines function to use attributes from class.

sample(particles, seed=None)[source]

Sample the action using the covariance sampling method.

Parameters:
  • particles (Tensor) – Particles of shape (batch_size, n_particles, h, w)

  • seed (int | SeedGenerator | jax.random.key, optional) – Seed for random number generation. Defaults to None.

Returns:

  • Newly selected lines as k-hot vectors, shaped (batch_size, n_possible_actions)

  • Masks of shape (batch_size, img_height, img_width)

Return type:

Tuple[Tensor, Tensor]

class zea.agent.selection.EquispacedLines(n_actions, n_possible_actions, img_width, img_height, assert_equal_spacing=True)[source]

Bases: LinesActionModel

Equispaced lines action selection.

Creates masks with equispaced lines that sweep across the image.

initial_sample_stateless(batch_size=1)[source]

Initial sample stateless.

Generates a batch of initial equispaced line masks.

Returns:

  • Selected lines as k-hot vectors, shaped (batch_size, n_possible_actions)

  • Masks of shape (batch_size, img_height, img_width)

Return type:

Tuple[Tensor, Tensor]

sample(current_lines=None, batch_size=1)[source]

Sample the action using the equispaced method.

Generates or updates an equispaced mask to sweep rightwards by one step across the image.

Returns:

The mask of shape (batch_size, img_size, img_size)

Return type:

Tensor

sample_stateless(current_lines)[source]

Sample stateless.

Updates an existing equispaced mask to sweep rightwards by one step across the image.

Parameters:

current_lines – Currently selected lines as k-hot vectors, shaped (batch_size, n_possible_actions)

Returns:

  • Newly selected lines as k-hot vectors, shaped (batch_size, n_possible_actions)

  • Masks of shape (batch_size, img_height, img_width)

Return type:

Tuple[Tensor, Tensor]

class zea.agent.selection.GreedyEntropy(n_actions, n_possible_actions, img_width, img_height, mean=0, std_dev=1, num_lines_to_update=5, entropy_sigma=1.0)[source]

Bases: LinesActionModel

Greedy entropy action selection.

Selects the max entropy line and reweights the entropy values around it, approximating the decrease in entropy that would occur from observing that line.

The neighbouring values are decreased by a Gaussian function centered at the selected line.

static compute_pairwise_pixel_gaussian_error(particles, stack_n_cols=1, n_possible_actions=None, entropy_sigma=1)[source]

Compute the pairwise pixelwise Gaussian error.

This function computes the Gaussian error between each pair of pixels in the set of particles provided. This can be used to approximate the entropy of a Gaussian mixture model, where the particles are the means of the Gaussians. For more details see Section 4 here: https://arxiv.org/abs/2406.14388

Parameters:

particles (Tensor) – Particles of shape (batch_size, n_particles, height, width)

Returns:

batch of pixelwise pairwise Gaussian errors, of shape (n_particles, n_particles, batch, height, width)

Return type:

Tensor

compute_pixelwise_entropy(particles)[source]

This function computes the entropy for each line using a Gaussian Mixture Model approximation of the posterior distribution. For more details see Section VI. B here: https://arxiv.org/pdf/2410.13310

Parameters:

particles (Tensor) – Particles of shape (batch_size, n_particles, height, width)

Returns:

batch of entropies per pixel, of shape (batch, height, width)

Return type:

Tensor

sample(particles)[source]

Sample the action using the greedy entropy method.

Parameters:

particles (Tensor) – Particles of shape (batch_size, n_particles, height, width)

Returns:

  • Newly selected lines as k-hot vectors, shaped (batch_size, n_possible_actions)

  • Masks of shape (batch_size, img_height, img_width)

Return type:

Tuple[Tensor, Tensor]

select_line_and_reweight_entropy(entropy_per_line)[source]

Select the line with maximum entropy and reweight the entropies.

Selected the max entropy line and reweights the entropy values around it, approximating the decrease in entropy that would occur from observing that line.

Parameters:

entropy_per_line (Tensor) – Entropy per line of shape (batch_size, n_possible_actions)

Returns:

The selected line index and the updated entropies per line

Return type:

Tuple

class zea.agent.selection.LinesActionModel(n_actions, n_possible_actions, img_width, img_height)[source]

Bases: MaskActionModel

Base class for action selection methods that select lines.

lines_to_im_size(lines)[source]

Convert k-hot-encoded line vectors to image size.

Parameters:

lines (Tensor) – shape is (n_masks, n_possible_actions)

Returns:

Masks of shape (n_masks, img_height, img_width)

Return type:

Tensor

class zea.agent.selection.MaskActionModel[source]

Bases: object

Base class for any action selection method that does masking.

apply(action, observation)[source]

Apply the action to the observation.

Parameters:
  • action (Tensor) – The mask to be applied.

  • observation (Tensor) – The observation to which the action is applied.

Returns:

The masked tensor

Return type:

Tensor

class zea.agent.selection.UniformRandomLines(n_actions, n_possible_actions, img_width, img_height)[source]

Bases: LinesActionModel

Uniform random lines action selection.

Creates masks with uniformly randomly sampled lines.

sample(batch_size=1, seed=None)[source]

Sample the action using the uniform random method.

Generates or updates an equispaced mask to sweep rightwards by one step across the image.

Parameters:

seed (int | SeedGenerator | jax.random.key, optional) – Seed for random number generation. Defaults to None.

Returns:

  • Newly selected lines as k-hot vectors, shaped (batch_size, n_possible_actions)

  • Masks of shape (batch_size, img_height, img_width)

Return type:

Tuple[Tensor, Tensor]