multi_step_lookahead

bojaxns.gaussian_process_formulation.multi_step_lookahead

Module Contents

convert_tree_to_graph(tree, action_labels=None, batch_index=0)

Converts a search tree into a Graphviz graph.

Parameters:

• tree (mctx.Tree) – A Tree containing a batch of search data.

• action_labels (Optional[Sequence[str]]) – Optional labels for edges, defaults to the action index.

• batch_index (int) – Index of the batch element to plot.

Returns:

A Graphviz graph representation of tree.

class MultiLookAheadState

Bases: NamedTuple

data: bojaxns.gaussian_process_formulation.distribution_math.GaussianProcessData

depth: chex.Array

RewardFnType

run_multi_lookahead(rng_key, data, ns_results, batch_size, max_depth, num_actions, num_simulations, S)

Parameters:

• rng_key (chex.PRNGKey) –

• data (bojaxns.gaussian_process_formulation.distribution_math.GaussianProcessData) –

• ns_results (bojaxns.base.MarginalisationData) –

• batch_size (int) –

• max_depth (int) –

• num_actions (int) –

• num_simulations (int) –

• S (int) –

Return type:

Tuple[chex.Array, mctx.PolicyOutput[mctx.GumbelMuZeroExtraData]]

static_fori_loop(lower, upper, body_fun, init_val)