agents

import "github.com/umbralcalc/stochadex/pkg/agents"

Package agents provides decision-making agents that operate over a generic Environment[S, A] interface. The package is intended to host any agent built on the same environment framework. Currently it ships MCTS (UCT) as the only agent, with MAST as an optional rollout strategy on top.

Per-player terminal scores are []float64 in [0,1] (the established stochadex value convention). Codecs (encoder/decoder for S into the stochadex row’s []float64) are supplied by the caller as function fields on each partition. This package does not depend on any encoding protocol.

Key Features:

Usage Patterns:

Index

Constants

Default progressive-widening parameters. alpha = 0.5 is the usual choice: the outcome count grows as the square root of the visit count, which keeps enough visits per outcome for their averages to mean something.

const (
    MCTSDefaultChanceWideningFactor   = 1.0
    MCTSDefaultChanceWideningExponent = 0.5
)

Default hyperparameters used when MCTSConfig fields are zero.

const (
    MCTSDefaultSimulations     = 120
    MCTSDefaultRolloutMaxSteps = 220
    MCTSDefaultMaxTreeDepth    = 14
    MCTSDefaultExploration     = 1.41
)

Param key for direct (within-step) best-action input. The value should be a 1-element slice — the legal-action index. A negative value (e.g. the -1 sentinel) leaves the state unchanged.

const ApplyParamBestIdx = "best_legal_idx"

Param key for state-history (lag-1) best-action input. The value (set via ParamsAsPartitions) is a 1-element slice containing the partition index of the upstream partition whose row[BestIdxSlot] holds the best-action index.

const ApplyParamBestIdxPartition = "best_idx_partition"

MASTAggregationParamPartition is the params_as_partitions key used by downstream samplers to learn this partition’s index for state-history reads. The value is a 1-element slice containing the partition index.

const MASTAggregationParamPartition = "mast_aggregates_partition"

MASTAggregationParamUpdates is the params_from_upstream key used to read the variable-length update batch from an upstream rollout partition.

const MASTAggregationParamUpdates = "mast_updates"

MASTDefaultTau is the softmax temperature applied to MAST means when sampling. Smaller → more exploitative; larger → closer to uniform. 1.0 diverges from uniform within ~50 updates per key in practice.

const MASTDefaultTau = 1.0

MASTSamplePrior is the score assigned to actions whose key has not yet been observed, so they retain selection probability early in search. 0.5 sits at the midpoint of the [0, 1] reward range — neither encouraged nor discouraged.

const MASTSamplePrior = 0.5

Param key used by MCTSRolloutIteration.Iterate to read the leaf state from an upstream MCTSTreeIteration via params_from_upstream. The slice should be of length StateWidth followed by a single has_leaf flag (matching MCTSTreeIteration’s row layout: leaf_state then has_leaf). Use Indices on the NamedUpstreamConfig to slice out the leaf_state + has_leaf section of the tree’s row.

const MCTSRolloutParamLeaf = "leaf"

Param key used by MCTSTreeIteration.Iterate to read rollout scores from an upstream MCTSRolloutIteration via params_from_upstream (within-step). The slice should be of length Players + 1 (P scores + ok flag) — the layout produced by MCTSRolloutIteration.

This mode creates a within-step dependency that breaks if the rollout partition also depends on the tree (the standard MCTS pipeline does). Use MCTSTreeParamRolloutScoresPartition for the lag-1 state-history mode instead when wiring the tree + rollout pipeline.

const MCTSTreeParamRolloutScores = "rollout_scores"

Param key used by MCTSTreeIteration.Iterate to read rollout scores from an upstream MCTSRolloutIteration via params_as_partitions (state-history, lag-1). The value is a 1-element slice containing the rollout partition’s index; the tree reads stateHistories[idx].Values row 0 (= the previous step’s rollout output) at runtime.

This is the standard wiring used by NewMCTSSelfPlayPartitions: rollout reads tree’s leaf within-step (so the rollout sees the freshest leaf), and tree reads rollout’s scores lag-1 (so rollout doesn’t have to wait on tree for last step’s scores). The 1-step lag aligns correctly: at step N+1 tree backs up the path it selected at step N with scores from rollout at step N (which were for that very leaf).

State-history mode takes priority over within-step mode if both keys are present.

const MCTSTreeParamRolloutScoresPartition = "rollout_scores_partition"

Param key used by MCTSTreeIteration.Iterate to read the current search root state. Whenever this param’s value differs from the cached root encoding, the tree is reset to the decoded new root. The slice should be of length StateWidth.

Set this via the embedded simulation run’s outer params_from_upstream (e.g. an outer apply partition piping its current game state into the inner sim via the “<innerName>/root_state” forwarding mechanism). When MCTSTreeIteration is used standalone with no outer pipeline, the param is absent and the tree retains the root set at Configure time.

const MCTSTreeParamRootState = "root_state"

Row layout slot accessors. Use these to compute params_from_upstream indices when wiring downstream partitions to MCTSTreeIteration’s row.

const MCTSTreeRowBestRootIdx = 0

const MCTSTreeRowLeafStateOffset = 1

TTTWidth is the encoded row width: 9 cells + current player.

const TTTWidth = 10

Variables

WinLines is the eight three-in-a-row patterns.

var WinLines = [...][3]int{
    {0, 1, 2}, {3, 4, 5}, {6, 7, 8},
    {0, 3, 6}, {1, 4, 7}, {2, 5, 8},
    {0, 4, 8}, {2, 4, 6},
}

func MASTAggregationCountSlot

func MASTAggregationCountSlot(k int) int

MASTAggregationCountSlot returns the row offset of the count for key k.

func MASTAggregationRowWidth

func MASTAggregationRowWidth(maxKeys int) int

MASTAggregationRowWidth returns the required state_width for an MASTAggregationIteration with the given key bound.

func MASTAggregationSumSlot

func MASTAggregationSumSlot(k int) int

MASTAggregationSumSlot returns the row offset of the sum for key k.

func MASTMeanForKey

func MASTMeanForKey(row []float64, k int) (mean float64, count int)

MASTMeanForKey reads the running mean reward for key k from a row in the MASTAggregationIteration’s layout. Returns (0, 0) when the key has not been observed. Used by samplers that have read the partition’s row via params_as_partitions.

func MASTRolloutNumPathOffset

func MASTRolloutNumPathOffset(players int) int

MASTRolloutNumPathOffset returns the row offset of the num_path counter.

func MASTRolloutOkOffset

func MASTRolloutOkOffset(players int) int

MASTRolloutOkOffset returns the row offset of the ok flag.

func MASTRolloutPathOffset

func MASTRolloutPathOffset(players int) int

MASTRolloutPathOffset returns the row offset of the first (key_idx, reward) pair.

func MASTRolloutRowWidth

func MASTRolloutRowWidth(players, maxPath int) int

MASTRolloutRowWidth returns the required state_width for an MASTRolloutIteration with the given player count and path bound.

func MASTRolloutScoresOffset

func MASTRolloutScoresOffset(i int) int

MASTRolloutScoresOffset returns the row offset of score slot i.

func MCTSRolloutRowWidth

func MCTSRolloutRowWidth(players int) int

MCTSRolloutRowWidth returns the required InitStateValues / StateWidth for a MCTSRolloutIteration with the given player count.

func MCTSTreeRowHasLeafOffset

func MCTSTreeRowHasLeafOffset(stateWidth int) int

func MCTSTreeRowVisitsOffset

func MCTSTreeRowVisitsOffset(stateWidth int) int

func MCTSTreeRowWidth

func MCTSTreeRowWidth(stateWidth, maxLegalActions int) int

MCTSTreeRowWidth returns the required InitStateValues / StateWidth for a MCTSTreeIteration with the given encoded-state width and max legal-action count.

func MCTSTreeRowWinsOffset

func MCTSTreeRowWinsOffset(stateWidth, maxLegalActions int) int

func TTTEncode

func TTTEncode(s TTTState) []float64

TTTEncode produces the []float64 row representation of a TTTState. Done/Winner are derivable from the cells so they are not encoded.

func TTTKey

func TTTKey(a TTTAction) string

TTTKey is a stable string key for an action, useful for any aggregation or lookup that needs a string-typed action identifier.

func WinnerToTerminal

func WinnerToTerminal(winner, players int, done bool) []float64

WinnerToTerminal builds an Environment.Terminal-compatible result from a binary winner. For envs whose native terminal predicate is “winner int, done bool” rather than per-player scores, embed this in your Environment implementation:

func (g *MyGame) Terminal(s State) ([]float64, bool) {
    w, done := g.winnerOrDone(s)
    return agents.WinnerToTerminal(w, g.Players(s), done), done
}

type ApplyIteration

ApplyIteration advances the environment by one ply per stochadex outer step using a best-action signal supplied either via params_from_upstream (within-step) OR via params_as_partitions (lagged read of an upstream partition’s state-history row). The partition row is the encoded current game state; one outer step decodes the row, applies the chosen legal action, and writes the encoded post-move state back.

Row layout (width = StateWidth):

row[0 .. StateWidth-1]   encoded current game state.

Two read modes

  1. Direct param mode (ApplyParamBestIdx): set ParamsFromUpstream[ApplyParamBestIdx] to read the best-action index directly within the same step. Used when the upstream partition is not in a self-referential cycle with apply.

  2. State-history mode (ApplyParamBestIdxPartition + BestIdxSlot): set ParamsAsPartitions[ApplyParamBestIdxPartition] to the upstream partition’s name, and set BestIdxSlot to the offset of the best-action index within that partition’s row. Apply will read the PREVIOUS step’s row 0 of that partition each step. Used to break the apply ↔︎ search dependency cycle in NewMCTSSelfPlayPartitions: search reads apply within-step (so apply does not wait on search), and apply reads search lagged by one step (so search does not wait on apply). The 1-step lag aligns correctly because at step N apply applies the best_idx that search produced at step N-1 for apply’s state at step N-1 — which is the same state apply currently holds (apply only advances when it applies a move).

State-history mode takes priority if both are configured.

Warm fields (must be set before Configure):

Optional warm field:

type ApplyIteration[S any, A any] struct {
    Env         Environment[S, A]
    Decoder     func([]float64) (S, error)
    Encoder     func(S) []float64
    BestIdxSlot int
    // contains filtered or unexported fields
}

func (*ApplyIteration[S, A]) Configure

func (m *ApplyIteration[S, A]) Configure(partitionIndex int, settings *simulator.Settings)

Configure implements simulator.Iteration.

func (*ApplyIteration[S, A]) Iterate

func (m *ApplyIteration[S, A]) Iterate(params *simulator.Params, partitionIndex int, stateHistories []*simulator.StateHistory, timestepsHistory *simulator.CumulativeTimestepsHistory) []float64

Iterate implements simulator.Iteration.

type BeliefSpec

BeliefSpec configures in-tree belief updating.

type BeliefSpec struct {
    // ObservationPartition names the partition whose state[0] the decision-maker
    // observes each step. Its value under the trajectory's true parameters is
    // scored against what each sample predicts, and the belief is reweighted by
    // the result. An observation the samples all predict alike carries no
    // information and leaves the belief untouched, which is the correct
    // behaviour and what makes an uninformative action look uninformative.
    ObservationPartition string
    // Variance is the observation-noise variance of the Gaussian likelihood used
    // for that scoring. Larger values make the belief move more slowly.
    Variance float64
}

type Environment

Environment is the game/decision-process interface that MCTS searches over. Implementations are pure: Legal/Apply must not mutate s, and Apply must return a fresh value (the search clones aggressively).

Terminal returns the per-player [0,1] score vector and a done flag. The score vector lets the env represent draws (0.5/0.5), graded scoring (Catan-like), and standard winner-takes-all (one-hot) without faking a “winner” int. For binary games, see WinnerToTerminal in rollout.go.

Actor returns the index of the player whose decision creates the next edge from s; backups credit nodes by Actor at the parent. Players returns the total seat count and bounds the score vector length.

type Environment[S any, A any] interface {
    Legal(s S) []A
    Apply(s S, a A) (S, error)
    Terminal(s S) (scores []float64, done bool)
    Actor(s S) int
    Players(s S) int
}

type MASTAggregationIteration

MASTAggregationIteration is a stochadex iteration that maintains running (count, sum) pairs per action-key index for MAST (Move-Average Sampling Technique). Each step it reads a variable-length update batch from an upstream rollout partition via params_from_upstream and applies the (key_idx, reward) increments to its row.

Row layout (width = 2 * MaxKeys):

row[2*k]     count for key k
row[2*k+1]   sum   for key k

Use MASTAggregationRowWidth(K) to compute state_width.

Update batch format

The upstream batch is a single []float64 with the layout

[num_updates, key_idx_0, reward_0, key_idx_1, reward_1, ...]

where num_updates is the number of valid (key_idx, reward) pairs in the slice. MASTRolloutIteration emits exactly this layout in its row’s path-suffix slots; wire MASTAggregationIteration’s params_from_upstream (key MASTAggregationParamUpdates) to those slots.

Out-of-range key indices and updates beyond the slice’s declared num_updates are silently dropped.

Read access

Downstream samplers (e.g. MASTRolloutIteration) read the aggregates via state-history mode (lag-1) using params_as_partitions. See MASTAggregationParamPartition for the canonical key.

type MASTAggregationIteration[A any] struct {
    MaxKeys int
    // contains filtered or unexported fields
}

func (*MASTAggregationIteration[A]) Configure

func (m *MASTAggregationIteration[A]) Configure(partitionIndex int, settings *simulator.Settings)

Configure implements simulator.Iteration.

func (*MASTAggregationIteration[A]) Iterate

func (m *MASTAggregationIteration[A]) Iterate(params *simulator.Params, partitionIndex int, stateHistories []*simulator.StateHistory, timestepsHistory *simulator.CumulativeTimestepsHistory) []float64

Iterate implements simulator.Iteration.

type MASTRolloutIteration

MASTRolloutIteration is a stochadex iteration that runs one MAST-biased rollout per step. It reads the leaf state from an upstream MCTSTreeIteration (within-step), reads the running aggregates from a MASTAggregationIteration (lag-1, via params_as_partitions), runs a playout sampling each ply via softmax over the aggregates, and emits the per-player scores plus the (key_idx, reward) path that the MASTAggregationIteration will absorb on the next step.

Row layout (width = Players + 2 + 2 * MaxPath):

row[0 .. Players-1]                           per-player [0,1] scores
row[Players]                                  ok flag (1 = scores valid)
row[Players+1]                                num_path (length of valid pairs)
row[Players+2 .. Players+1+2*MaxPath]         (key_idx, reward) pairs,
                                              padded with zeros

The (num_path, pairs…) suffix matches MASTAggregationParamUpdates so the downstream MASTAggregationIteration can read it as a single slice.

Warm fields (must be set before Configure):

type MASTRolloutIteration[S any, A any] struct {
    Env      Environment[S, A]
    Cfg      MCTSConfig[S, A]
    Decoder  func([]float64) (S, error)
    KeyToIdx func(A) int
    MaxKeys  int
    MaxPath  int
    Players  int
    Tau      float64
    Progress func(s S, player int) (float64, bool)
    // contains filtered or unexported fields
}

func (*MASTRolloutIteration[S, A]) Configure

func (m *MASTRolloutIteration[S, A]) Configure(partitionIndex int, settings *simulator.Settings)

Configure implements simulator.Iteration.

func (*MASTRolloutIteration[S, A]) Iterate

func (m *MASTRolloutIteration[S, A]) Iterate(params *simulator.Params, partitionIndex int, stateHistories []*simulator.StateHistory, timestepsHistory *simulator.CumulativeTimestepsHistory) []float64

Iterate implements simulator.Iteration.

type MCTSChanceTree

MCTSChanceTree is the UCT tree for stochastic environments, where an action’s value is an average over sampled successors rather than the value of one.

Levels alternate. A decision node holds a state and one chance child per legal action. A chance node holds no state — it stands for “this action was taken, and nature has not resolved yet” — and its children are sampled successors. Action selection at a decision node is UCB1 over the chance children, whose statistics are therefore averages over outcomes.

A chance node cannot enumerate its outcomes, so their number grows with visits as ceil(k * visits^alpha), alpha in (0,1). Unbounded, every visit would draw a fresh successor and no outcome would collect enough visits for its average to mean anything.

Use MCTSTree for deterministic environments: one successor per edge is exact there, and outcome nodes whose samples never differ are pure cost.

type MCTSChanceTree[S any, A any] struct {
    // contains filtered or unexported fields
}

func NewMCTSChanceTree

func NewMCTSChanceTree[S any, A any](root S) *MCTSChanceTree[S, A]

NewMCTSChanceTree returns a tree with a single decision node at root.

func (*MCTSChanceTree[S, A]) BackupScores

func (t *MCTSChanceTree[S, A]) BackupScores(path []int, scores []float64)

BackupScores credits each node on path with its actor’s score.

func (*MCTSChanceTree[S, A]) BackupVisits

func (t *MCTSChanceTree[S, A]) BackupVisits(path []int, scores []float64)

BackupVisits increments visits along path, crediting wins only when scores are present — the no-signal-tolerant form, for the same reason as MCTSTree’s.

func (*MCTSChanceTree[S, A]) NodeCount

func (t *MCTSChanceTree[S, A]) NodeCount() int

NodeCount returns the number of nodes, decision and chance alike.

func (*MCTSChanceTree[S, A]) Reset

func (t *MCTSChanceTree[S, A]) Reset(root S)

Reset replaces the tree with a fresh root.

func (*MCTSChanceTree[S, A]) Root

func (t *MCTSChanceTree[S, A]) Root() S

Root returns the root state.

func (*MCTSChanceTree[S, A]) RootBestLegalIdx

func (t *MCTSChanceTree[S, A]) RootBestLegalIdx() (int, bool)

RootBestLegalIdx returns the most-visited action at the root.

func (*MCTSChanceTree[S, A]) RootStatsByLegalIdx

func (t *MCTSChanceTree[S, A]) RootStatsByLegalIdx(maxLegalActions int) (visits, wins []float64)

RootStatsByLegalIdx returns per-action visits and win sums at the root, padded to maxLegalActions. The statistics come from the chance children, so each is already an average over that action’s sampled outcomes.

func (*MCTSChanceTree[S, A]) SelectLeaf

func (t *MCTSChanceTree[S, A]) SelectLeaf(env Environment[S, A], cfg *MCTSConfig[S, A], rng *rand.Rand) (path []int, leafState S, leafIdx int, ok bool)

SelectLeaf is the boolean-shaped form, matching MCTSTree.SelectLeaf.

func (*MCTSChanceTree[S, A]) SelectLeafWithOutcome

func (t *MCTSChanceTree[S, A]) SelectLeafWithOutcome(env Environment[S, A], cfg *MCTSConfig[S, A], rng *rand.Rand) (path []int, leafState S, leafIdx int, outcome MCTSLeafOutcome)

SelectLeafWithOutcome walks from the root to a leaf, alternating UCB1 action choices at decision nodes with sampled outcomes at chance nodes, and reports why it stopped. It mirrors MCTSTree.SelectLeafWithOutcome so the same backup rules apply, and pairs with BackupScores / BackupVisits.

type MCTSConfig

MCTSConfig holds UCT hyperparameters and the rollout driver.

Rollout is the single graded rollout signature: it returns a per-player score vector in [0,1] and an ok flag (false signals “no signal” — caller counts the visit but skips the win credit). For binary winner games or games using a Progress proxy, build the rollout from helpers in rollout.go (UniformRandomRollout, OneHotFromWinner, FromProgress).

Progress is an optional per-state, per-player [0,1] value proxy used by the FromProgress rollout adapter to score truncated rollouts.

type MCTSConfig[S any, A any] struct {
    Simulations     int
    Exploration     float64
    MaxTreeDepth    int
    RolloutMaxSteps int
    Rollout         MCTSRolloutFn[S, A]
    Progress        func(s S, player int) (float64, bool)

    // ChanceWideningFactor and ChanceWideningExponent tune progressive widening
    // in MCTSChanceTree: a chance node holds up to
    // ceil(factor * visits^exponent) sampled outcomes. Ignored by MCTSTree,
    // which has no chance nodes. Zero selects the package defaults.
    //
    // A larger factor or exponent averages over more outcomes per node but
    // spreads the same visits more thinly, so each average is noisier.
    ChanceWideningFactor   float64
    ChanceWideningExponent float64
}

func (*MCTSConfig[S, A]) ApplyDefaults

func (c *MCTSConfig[S, A]) ApplyDefaults()

ApplyDefaults fills in zero-valued hyperparameters with the package defaults, mutating the receiver. Called once at the start of each search run; safe to call multiple times.

Exported so external packages building on MCTSConfig can share the defaults logic without duplicating it.

type MCTSEdgeStat

MCTSEdgeStat is per-action telemetry exposed at the root after a search. Useful for JSON reports or logging the search’s distribution over moves.

type MCTSEdgeStat[A any] struct {
    Action       A       `json:"action"`
    Visits       int     `json:"visits"`
    MeanForActor float64 `json:"mean_for_actor"`
}

func RunChanceMCTSSearch

func RunChanceMCTSSearch[S any, A any](env StochasticEnvironment[S, A], root S, cfg MCTSConfig[S, A], baseSeed uint64, sims int) (A, []MCTSEdgeStat[A], error)

RunChanceMCTSSearch is RunMCTSSearch over an MCTSChanceTree, so each action’s value averages over sampled outcomes. Prefer it whenever the environment has a real transition distribution and the value it reports will be believed.

func RunMCTSSearch

func RunMCTSSearch[S any, A any](env Environment[S, A], root S, cfg MCTSConfig[S, A], baseSeed uint64, sims int) (A, []MCTSEdgeStat[A], error)

RunMCTSSearch runs sims UCT simulations from root and returns the best legal action plus per-edge stats. Independent of stochadex partitions — useful for one-shot “what’s the best move?” queries.

Defaults are filled in from the package constants. If cfg.Rollout is nil, UniformRandomRollout is used.

type MCTSLeafOutcome

MCTSLeafOutcome says why selection stopped where it did. The distinction is load-bearing for callers driving the decomposed pipeline: the four non-expansion outcomes each need a different backup, and MCTSTree.RunOne applies exactly one of those per iteration. Collapsing them into a single “not ok” and dropping the iteration silently starves the tree of statistics — see MCTSTreeIteration.Iterate.

type MCTSLeafOutcome int

const (
    // MCTSLeafExpanded: a new child node was created. Roll out from it and back
    // the resulting scores up along the returned path.
    MCTSLeafExpanded MCTSLeafOutcome = iota
    // MCTSLeafTerminal: selection reached a finished position. The environment
    // scores it exactly, so back those scores up directly — no rollout needed.
    MCTSLeafTerminal
    // MCTSLeafDepthCapped: MaxTreeDepth was reached. RunOne rolls out from the
    // capped node (a Progress proxy is the usual way to score the truncation)
    // and backs the result up; a decomposed caller should do the same.
    MCTSLeafDepthCapped
    // MCTSLeafNoLegalActions: a non-terminal node offering no legal actions —
    // a stalled position. Handled exactly like a depth cap.
    MCTSLeafNoLegalActions
    // MCTSLeafApplyFailed: env.Apply errored during expansion, so no node was
    // created. Nothing is backed up: this is an environment fault, not a search
    // result, and crediting a visit for it would bias selection away from an
    // action whose only sin is a broken transition.
    MCTSLeafApplyFailed
)

type MCTSRolloutFn

MCTSRolloutFn is the single rollout signature used by the search. It runs a stochastic playout from s for at most maxSteps actions and returns:

The seed argument is the rollout’s full RNG seed; implementations should be deterministic given the same seed.

type MCTSRolloutFn[S any, A any] func(env Environment[S, A], s S, maxSteps int, seed uint64) (scores []float64, ok bool, err error)

func FromProgress

func FromProgress[S any, A any](inner MCTSRolloutFn[S, A], progress func(s S, player int) (float64, bool)) MCTSRolloutFn[S, A]

FromProgress wraps an inner rollout so that truncated rollouts (ok=false from inner) are rescued by scoring the final state via the supplied progress function. progress returns a per-player [0,1] proxy of “how close is this player to winning” and an ok flag (false = no proxy available for that player; treated as 0).

All-equal progress vectors carry no comparative signal and are treated as no signal (ok=false from the wrapper) so the search relies on UCB exploration alone — see the docstring on MCTSTree.backupVisits for the stall-move failure mode this avoids.

FromProgress needs to know the final state of the inner rollout, so it re-runs the playout itself rather than wrapping inner. The inner rollout is therefore only consulted for early termination.

func UniformRandomRollout

func UniformRandomRollout[S any, A any]() MCTSRolloutFn[S, A]

UniformRandomRollout returns a MCTSRolloutFn that plays uniformly random legal actions until either Terminal returns done or maxSteps is reached. On termination the env’s Terminal scores are returned. On truncation (maxSteps reached without termination) the rollout returns ok=false — compose with FromProgress if you have a progress proxy to score truncated rollouts.

type MCTSRolloutIteration

MCTSRolloutIteration runs one rollout per stochadex step. It reads the leaf state to roll out from via params_from_upstream (typically wired to a MCTSTreeIteration’s leaf_state slot via Indices) and outputs a per-player score vector plus an ok flag in its own row.

Row layout (width = Players + 1):

row[0 .. Players-1]   per-player [0,1] scores from the rollout
row[Players]          ok flag (1 if the rollout produced valid scores,
                      0 otherwise — the upstream MCTSTreeIteration uses
                      this to decide whether to apply backupScores or
                      backupVisits with no signal)

Use MCTSRolloutRowWidth(P) to size InitStateValues / state_width.

Stateless across steps — each Iterate is one independent rollout. Swap in FromProgress, WinnerToTerminal, or any custom MCTSRolloutFn via Cfg.Rollout without touching the partition wiring.

Warm fields (must be set before Configure):

type MCTSRolloutIteration[S any, A any] struct {
    Env     Environment[S, A]
    Cfg     MCTSConfig[S, A]
    Decoder func([]float64) (S, error)
    Players int
    // contains filtered or unexported fields
}

func (*MCTSRolloutIteration[S, A]) Configure

func (m *MCTSRolloutIteration[S, A]) Configure(partitionIndex int, settings *simulator.Settings)

Configure implements simulator.Iteration.

func (*MCTSRolloutIteration[S, A]) Iterate

func (m *MCTSRolloutIteration[S, A]) Iterate(params *simulator.Params, partitionIndex int, stateHistories []*simulator.StateHistory, timestepsHistory *simulator.CumulativeTimestepsHistory) []float64

Iterate implements simulator.Iteration.

type MCTSTree

MCTSTree is the in-memory UCT search tree for a fixed root state. Methods are not safe for concurrent use; one MCTSTree per goroutine.

The MCTSTree owns no environment or config — those are passed in to RunOne per-simulation. This makes MCTSTree easy to embed in iterations that may want to tweak config between simulations (e.g. adaptive exploration).

type MCTSTree[S any, A any] struct {
    // contains filtered or unexported fields
}

func NewMCTSTree

func NewMCTSTree[S any, A any](root S) *MCTSTree[S, A]

NewMCTSTree returns a MCTSTree with a single root node containing root.

func (*MCTSTree[S, A]) AdvanceRoot

func (t *MCTSTree[S, A]) AdvanceRoot(env Environment[S, A], legalIdx int)

AdvanceRoot promotes the root’s child at the given legal index to be the new root, preserving its subtree (classic MCTS tree reuse). If that child was never expanded, the tree is rebuilt fresh from the resulting state.

env is needed to compute the post-move state if the subtree is missing.

func (*MCTSTree[S, A]) BackupScores

func (t *MCTSTree[S, A]) BackupScores(path []int, scores []float64)

BackupScores credits each node along path with the score belonging to its actor (visits and wins both increment). Exported wrapper around the internal backupScores so iterations split across selection / rollout / backup partitions can apply the scores when they arrive.

func (*MCTSTree[S, A]) BackupVisits

func (t *MCTSTree[S, A]) BackupVisits(path []int, scores []float64)

BackupVisits is the no-signal-tolerant variant: visits always increment, but wins are only credited when scores is non-nil. See backupVisits docs for the engine-stall reasoning.

func (*MCTSTree[S, A]) NodeCount

func (t *MCTSTree[S, A]) NodeCount() int

NodeCount returns the number of nodes currently in the tree (including the root). Useful for telemetry and capacity tuning.

func (*MCTSTree[S, A]) Reset

func (t *MCTSTree[S, A]) Reset(root S)

Reset replaces the entire tree with a fresh root at the given state. Use when no usable subtree exists (opening move, or after an external change to the root).

func (*MCTSTree[S, A]) Root

func (t *MCTSTree[S, A]) Root() S

Root returns the root state.

func (*MCTSTree[S, A]) RootBestLegalIdx

func (t *MCTSTree[S, A]) RootBestLegalIdx() (int, bool)

RootBestLegalIdx returns the most-visited (then most-winning) child legal index. Ties are broken via reservoir sampling over equally-good children so the choice is not biased toward the first-listed action — important for engine-heavy games where the first listed legal action is often a stall (recycle, pass) and a deterministic first-tie pick would deadlock the agent. Reservoir randomness is seeded from the current tree shape so the result is reproducible without taking an external rng.

func (*MCTSTree[S, A]) RootEdgeStats

func (t *MCTSTree[S, A]) RootEdgeStats(legal []A) []MCTSEdgeStat[A]

RootEdgeStats reports per-action visit counts and mean-for-actor for each expanded child of the root. legal must be the same slice ordering used by the env’s Legal(root); pass env.Legal(tree.Root()) at the call site.

func (*MCTSTree[S, A]) RootStatsByLegalIdx

func (t *MCTSTree[S, A]) RootStatsByLegalIdx(maxLegalActions int) (visits, wins []float64)

RootStatsByLegalIdx returns per-legal-action visit counts and win sums at the root, padded with zeros up to maxLegalActions. Returns (visits, wins) each of length maxLegalActions. Used to expose root statistics in fixed-width row layouts.

func (*MCTSTree[S, A]) RunOne

func (t *MCTSTree[S, A]) RunOne(env Environment[S, A], cfg *MCTSConfig[S, A], rng *rand.Rand)

RunOne does one UCT iteration: selection → expansion → rollout → backup. rng must be seeded by the caller; one call uses one RNG.

Calls cfg.applyDefaults() so a fresh MCTSConfig with only Rollout set works out of the box. The mutation is idempotent (only zero values are filled).

RunOne is the all-in-one path used by RunMCTSSearch and by callers who don’t need the selection / rollout / backup phases as separate stochadex partitions. For the decomposed pipeline use SelectLeaf + BackupScores.

func (*MCTSTree[S, A]) SelectLeaf

func (t *MCTSTree[S, A]) SelectLeaf(env Environment[S, A], cfg *MCTSConfig[S, A], rng *rand.Rand) (path []int, leafState S, leafIdx int, ok bool)

SelectLeaf walks the tree from the root using UCB1 (with first-visit preference for unvisited children) until it reaches an unexpanded edge, then expands it by creating a new child node. Returns the path of node indices from the root’s child down to the new leaf, the leaf’s state, the leaf’s node index, and ok=true. Returns ok=false if the root is terminal, has no legal moves, MaxTreeDepth is reached, or env.Apply fails during expansion.

ok=false collapses four outcomes that want different backups, so prefer SelectLeafWithOutcome unless you genuinely only care whether the tree grew.

SelectLeaf does NOT roll out and does NOT back up — it is the (selection + expansion) half of one MCTS iteration. Pair it with MCTSTree.BackupScores or MCTSTree.BackupVisits to apply the scores when they arrive.

Calls cfg.applyDefaults() so a fresh MCTSConfig works out of the box. The mutation is idempotent (only zero values are filled).

func (*MCTSTree[S, A]) SelectLeafWithOutcome

func (t *MCTSTree[S, A]) SelectLeafWithOutcome(env Environment[S, A], cfg *MCTSConfig[S, A], rng *rand.Rand) (path []int, leafState S, leafIdx int, outcome MCTSLeafOutcome)

SelectLeafWithOutcome is SelectLeaf with the stop reason reported explicitly, so a caller splitting selection / rollout / backup across partitions can apply the same backup RunOne would for each case.

type MCTSTreeIteration

MCTSTreeIteration runs the (selection + expansion + backup) phase of a UCT MCTS search as a stochadex iteration. The tree itself lives on the struct (graph state, fundamentally not []float64-shaped); the partition row exposes a fixed-width summary that downstream partitions can consume via params_from_upstream:

row[MCTSTreeRowBestRootIdx]                        — most-visited root
                                                 legal-action index
                                                 after the most recent
                                                 update (-1 if not
                                                 yet decided)
row[MCTSTreeRowLeafStateOffset .. +StateWidth-1]   — encoded state of the
                                                 leaf the search just
                                                 selected (input to a
                                                 rollout partition)
row[MCTSTreeRowHasLeafOffset(W)]                   — 1 if the leaf_state
                                                 slot is real, 0
                                                 otherwise (used by
                                                 rollout partitions to
                                                 short-circuit when
                                                 nothing has been
                                                 selected yet)
row[MCTSTreeRowVisitsOffset(W) .. +MaxLegalActions-1] — per-legal-action
                                                    root visit counts,
                                                    padded with zeros
row[MCTSTreeRowWinsOffset(W,K) .. +MaxLegalActions-1] — per-legal-action
                                                    root win sums,
                                                    padded with zeros

Use MCTSTreeRowWidth(W, K) to compute the required state_width / init slice length.

Pipeline lag

MCTSTreeIteration is one half of a 2-step pipeline with a downstream rollout partition. The rollout partition reads (leaf_state, has_leaf) and outputs scores; MCTSTreeIteration then reads those scores via params_from_upstream (key MCTSTreeParamRolloutScores) and applies a backup to the path it selected two steps earlier. The 2-step lag is fundamental to expressing selection-then-backup as stochadex’s single-row dataflow.

In steady state each outer step does one selection + one backup, so the throughput is one MCTS iteration per stochadex step (after a 2-step fill).

Warm fields (must be set before Configure):

type MCTSTreeIteration[S any, A any] struct {
    Env             Environment[S, A]
    Cfg             MCTSConfig[S, A]
    Decoder         func([]float64) (S, error)
    Encoder         func(S) []float64
    MaxLegalActions int
    StateWidth      int
    Players         int
    // ChanceNodes searches an MCTSChanceTree instead, averaging each action's
    // value over sampled successors rather than committing to the first one
    // drawn. Requires Env to implement StochasticEnvironment.
    //
    // Leave it off for deterministic environments: they gain nothing and would
    // pay for outcome nodes whose samples never differ. Turn it on whenever the
    // value the search reports is going to be believed — against a stochastic
    // model the deterministic search plans as though it knew which way the dice
    // would fall, and its over-promise grows with the simulation budget.
    ChanceNodes bool
    // contains filtered or unexported fields
}

func (*MCTSTreeIteration[S, A]) Configure

func (m *MCTSTreeIteration[S, A]) Configure(partitionIndex int, settings *simulator.Settings)

Configure implements simulator.Iteration. Decodes the encoded root from is.InitStateValues[1 .. 1+StateWidth] and resets the tree. The first slot (MCTSTreeRowBestRootIdx) and the stats slots can be left zero in the init: best_root_idx is initialised to -1 to signal “not yet decided”.

func (*MCTSTreeIteration[S, A]) Iterate

func (m *MCTSTreeIteration[S, A]) Iterate(params *simulator.Params, partitionIndex int, stateHistories []*simulator.StateHistory, timestepsHistory *simulator.CumulativeTimestepsHistory) []float64

Iterate implements simulator.Iteration.

func (*MCTSTreeIteration[S, A]) MCTSTree

func (m *MCTSTreeIteration[S, A]) MCTSTree() *MCTSTree[S, A]

MCTSTree exposes the underlying deterministic search tree (typically for telemetry). It returns nil when ChanceNodes is set, since the partition is then driving an MCTSChanceTree instead — use SearchTree for the common surface.

func (*MCTSTreeIteration[S, A]) SearchTree

func (m *MCTSTreeIteration[S, A]) SearchTree() interface {
    Root() S
    NodeCount() int
    RootStatsByLegalIdx(maxLegalActions int) (visits, wins []float64)
    RootBestLegalIdx() (int, bool)
}

SearchTree exposes whichever tree this partition is driving, for the telemetry both kinds share (root statistics, node count).

type SimulationEnvironment

SimulationEnvironment adapts a stochadex sub-simulation into an Environment[[]float64, int], so MCTS can plan over the same forward model the rest of the engine simulates and calibrates — rather than over game rules written by hand.

This is the counterpart to the environment registry in pkg/api. That hook lets a config NAME decision rules a downstream module wrote in Go; this type needs no rules at all, because the dynamics are already stated as partitions. An action is a params injection, a transition is one step of the sub-simulation, and the reward is read off the resulting rows.

State encoding

s[0]                 step index within the episode
s[1]                 accumulated (discounted) reward so far
s[2 .. 2+totalWidth] every partition's state-history window, flattened
                     row-major with the latest row first, concatenated in
                     partition order

Carrying the accumulated reward in the state is what lets a finite-horizon, per-step-reward problem satisfy the existing Environment contract without changing it: Terminal fires at the horizon and normalises the accumulated reward into the [0,1] score the UCT backups expect. ReturnRange declares the normalisation bounds, since UCB1 is only meaningful on a bounded value scale.

Determinism, and the approximation it encodes

Apply must be pure — MCTS materialises a child state once per edge and reuses it on every later visit — but a stochastic sub-simulation is not. That is not a problem this type solves for itself: it delegates to simulator.ReentrantSimulation, the engine’s re-entrant evaluation tier, and supplies a seed derived from (ScenarioSeed, state, action). Two calls with the same arguments therefore return the same successor.

That is common random numbers: the noise is pinned as a function of where you are and what you do, turning the stochastic model into a deterministic surrogate that MCTSTree searches exactly. It is a real modelling choice, not a free win — a planner solving a pinned scenario can exploit the particular noise draw it is going to receive, which biases its values optimistic. On the battery test problem that over-promise reaches tens of percent and *grows* with the simulation budget, since more search means more exploitation of the one realisation it can see.

This type therefore also implements StochasticEnvironment, via ApplySample. A search that uses it (MCTSChanceTree) builds chance nodes and averages over sampled successors, which removes most of the bias: TestChanceNodesReduceOptimism measures both on the same problem.

Requirements and limits

type SimulationEnvironment struct {
    // contains filtered or unexported fields
}

func NewSimulationEnvironment

func NewSimulationEnvironment(settings *simulator.Settings, implementations *simulator.Implementations, spec SimulationEnvironmentSpec) *SimulationEnvironment

NewSimulationEnvironment validates the spec against the sub-simulation and returns the environment. It panics on a misconfiguration, matching the rest of the config-assembly tier: these are programming errors caught before any search starts, and a silently misshapen environment would produce plausible but meaningless recommendations.

func (*SimulationEnvironment) Actor

func (e *SimulationEnvironment) Actor([]float64) int

Actor implements Environment. A planning problem has a single decision maker.

func (*SimulationEnvironment) Apply

func (e *SimulationEnvironment) Apply(s []float64, a int) ([]float64, error)

Apply implements Environment: one step of the sub-simulation under action a.

The transition is pure. Its seed is derived from the scenario seed, the encoded state and the action, so the same arguments always give the same successor — see the type docs for what that pins and what it costs.

func (*SimulationEnvironment) ApplySample

func (e *SimulationEnvironment) ApplySample(s []float64, a int, seed uint64) ([]float64, error)

ApplySample implements StochasticEnvironment: one transition under an explicit sample seed, so a search can draw more than one successor of the same (state, action) and average over them instead of committing to the first.

Apply is ApplySample with the sample seed fixed at zero. That is the whole difference between planning against a pinned scenario and planning against the distribution.

func (*SimulationEnvironment) InitialState

func (e *SimulationEnvironment) InitialState() []float64

InitialState encodes the sub-simulation’s configured initial rows as the episode’s starting state.

func (e *SimulationEnvironment) Legal(s []float64) []int

Legal implements Environment.

func (*SimulationEnvironment) Players

func (e *SimulationEnvironment) Players([]float64) int

Players implements Environment.

func (*SimulationEnvironment) Progress

func (e *SimulationEnvironment) Progress(s []float64, player int) (float64, bool)

Progress scores an unfinished state on the same [0,1] scale Terminal uses, so a rollout that runs out of steps still contributes a value instead of no signal. Compose it with FromProgress.

The proxy is the reward banked so far. That is deliberately conservative: it credits nothing for a position that is merely promising, so it under-rates a leaf whose payoff comes later. It beats the alternative of discarding the rollout, which leaves the search exploring on visit counts alone — the failure mode that bites when the horizon is longer than RolloutMaxSteps.

Supply SimulationEnvironmentSpec.Progress to score positions instead of banked reward, which is what a domain proxy (a win probability, a margin) is for.

func (*SimulationEnvironment) Return

func (e *SimulationEnvironment) Return(s []float64) float64

Return reads the accumulated (discounted) reward out of an encoded state, which is the quantity a caller actually wants to report — the [0,1] score is an artefact of the UCB1 value scale.

func (*SimulationEnvironment) StateWidth

func (e *SimulationEnvironment) StateWidth() int

StateWidth is the encoded state width: the step index, the accumulated reward, and every partition’s row.

func (*SimulationEnvironment) Terminal

func (e *SimulationEnvironment) Terminal(s []float64) ([]float64, bool)

Terminal implements Environment: the episode ends at the horizon, scored by the accumulated reward normalised into [0,1].

type SimulationEnvironmentSpec

SimulationEnvironmentSpec configures a SimulationEnvironment.

type SimulationEnvironmentSpec struct {
    // Actions is the discrete action set. Each entry is the params value
    // injected under ActionParam on ActionPartition, so actions are whatever
    // the model already reads as parameters — a dispatch rate, a policy
    // threshold, an intervention size.
    Actions [][]float64
    // ActionPartition names the partition receiving the action, and ActionParam
    // the params key it is written to.
    ActionPartition string
    ActionParam     string
    // Horizon is the episode length in steps.
    Horizon int
    // Reward scores one transition from the post-step rows, addressed by
    // partition name. Returning the per-step reward (not the cumulative one)
    // keeps the discounting in one place.
    Reward func(rows map[string][]float64) float64
    // Discount is the per-step discount factor applied to rewards. Zero means
    // undiscounted (treated as 1).
    Discount float64
    // MinReturn and MaxReturn bound the achievable accumulated reward and are
    // used to normalise it into the [0,1] score UCB1 needs. A return outside
    // the range is clamped — widen the range rather than leaving it clamped,
    // since a saturated score carries no gradient for the search.
    MinReturn, MaxReturn float64
    // ScenarioSeed pins the noise realisation this environment plans against.
    ScenarioSeed uint64
    // Legal optionally restricts the action set given the decoded rows. Nil
    // means every action is always legal.
    Legal func(rows map[string][]float64) []int
    // Progress optionally scores an unfinished position in [0,1], for rollouts
    // that hit their step limit before the horizon. Nil falls back to the reward
    // banked so far — see Progress.
    Progress func(rows map[string][]float64) (float64, bool)
    // ParameterSamples turns the run into posterior-predictive planning: each
    // entry is one draw of the model's uncertain parameters, and the search
    // averages over them instead of committing to a point estimate.
    //
    // A sample set rather than a fitted distribution, because that is what the
    // inference tier already produces — SMC particle parameters, or draws from a
    // posterior_estimation partition — and it carries the posterior's actual
    // shape rather than a Gaussian summary of it.
    //
    // Each sample is routed into the model by ParameterTargets. Leave nil to
    // plan at whatever parameters the model was configured with.
    ParameterSamples [][]float64
    // ParameterTargets says where each sample's values go. Together the targets'
    // Indices must cover the sample vector.
    ParameterTargets []SimulationParamTarget
    // ParameterWeights is how much credence each sample carries, which is what an
    // inference tier hands a planner: SMC produces particles AND their weights,
    // and the weights are the posterior. Need not sum to 1; it is normalised.
    //
    // Nil means equal credence, which is right when the samples were themselves
    // drawn from the posterior (posterior_estimation's sampler does this, so its
    // draws already carry the posterior's shape) and wrong when they were drawn
    // from something else and reweighted.
    ParameterWeights []float64
    // Belief turns on in-tree belief updating: the planner carries a distribution
    // over ParameterSamples and reweights it from what each step reveals, so it
    // can value an action for what it teaches rather than only for what it pays.
    //
    // This is the expensive option. Every transition evaluates the model once per
    // sample to score the observation, so a step costs (len(ParameterSamples)+1)
    // model steps. Keep the sample set small.
    //
    // Nil plans on a fixed draw instead: still posterior-predictive across
    // trajectories, but blind to information — see the type docs.
    Belief *BeliefSpec
}

type SimulationParamTarget

SimulationParamTarget routes part of a parameter sample into the model, the same shape SMC uses to route a particle’s parameters into its own model.

type SimulationParamTarget struct {
    // Partition names the model partition receiving the values, and Param the
    // params key written on it.
    Partition string
    Param     string
    // Indices selects which of the sample's entries to send, in order.
    Indices []int
}

type StochasticEnvironment

StochasticEnvironment is an Environment whose transitions can be resampled.

Environment.Apply is deterministic by contract, because the search materialises a successor once per edge and reuses it. For a genuinely stochastic process that pins the noise: the planner ends up solving one realisation, and can exploit the particular draw it is going to receive, so the value it computes is optimistic relative to the true expectation.

Implementing this interface says “there is a distribution here, and here is how to draw from it”. A search that recognises it can build chance nodes and average over outcomes rather than committing to the first one — see MCTSChanceTree. Environments that are actually deterministic (game rules) should not implement it; they gain nothing and would pay for outcome nodes that never differ.

type StochasticEnvironment[S any, A any] interface {

    // ApplySample draws one successor of (s, a) under the given sample seed.
    // Distinct seeds must give draws from the transition distribution, and the
    // same seed must reproduce its draw — the search caches outcomes and has to
    // be able to revisit them.
    ApplySample(s S, a A, seed uint64) (S, error)
    // contains filtered or unexported methods
}

type TTTAction

TTTAction is a cell index 0..8.

type TTTAction int

type TTTGame

TTTGame implements Environment[TTTState, TTTAction]. Tic-tac-toe is the canonical fixture for testing this package and downstream consumers (e.g. pkg/macros): small enough to be obvious, large enough that random play loses, and deterministic at endgame — from a “win in one” position MCTS must pick the winning move; from a “block in one” position it must block.

type TTTGame struct{}

func (*TTTGame) Actor

func (g *TTTGame) Actor(s TTTState) int

func (*TTTGame) Apply

func (g *TTTGame) Apply(s TTTState, a TTTAction) (TTTState, error)

Apply places the current player’s mark on cell a and updates Done/Winner.

func (g *TTTGame) Legal(s TTTState) []TTTAction

Legal returns the empty cells; nil if the game is over.

func (*TTTGame) Players

func (g *TTTGame) Players(s TTTState) int

func (*TTTGame) Terminal

func (g *TTTGame) Terminal(s TTTState) ([]float64, bool)

Terminal returns the per-player [0,1] score vector and a done flag. Draw splits 0.5/0.5; a winner takes 1, the other 0.

type TTTState

TTTState is a tic-tac-toe position. cells holds 0=empty, 1=X, 2=O.

type TTTState struct {
    Cells   [9]int8
    Current int // 0=X to move, 1=O to move
    Done    bool
    Winner  int // -1=draw or in-progress, 0=X won, 1=O won
}

func TTTDecode

func TTTDecode(v []float64) (TTTState, error)

TTTDecode rebuilds a TTTState from its row representation, recomputing Done/Winner so callers can spell positions declaratively.

func TTTFromGrid

func TTTFromGrid(grid [9]int8, currentPlayer int) TTTState

TTTFromGrid builds a TTTState from a literal cell grid plus the player to move. Done/Winner are derived. Useful for spelling test positions inline.

Generated by gomarkdoc