mixle.inference.conditional_jit_controller module

D5: the LEARNED controller over the D1-D4/D6 estimator-tree IR (workstream ConditionalJIT track).

Frame (see the ConditionalJIT track, D1-D6): the estimator tree is an IR. D1 (mixle.inference.node_report) instruments every node with a per-round residual/Q-gain/cost report. D2 (mixle.inference.freeze_rollup) spends that report on one fixed rule (freeze once converged). D3 (mixle.inference.block_em) spends it on a fixed GREEDY rule (gain-per-cost ranking within a fixed budget_fraction). D4 (mixle.inference.leaf_hotswap) spends it on a fixed rule for gradient leaves specifically (swap once plateaued). D6 (origin/backend-respecialization, a sibling branch stacked on D3 that this branch does not directly contain – see this module’s PR description) spends it on a fixed compile-economics rule for execution backends. D5 REPLACES the “fixed rule” part with a LEARNED one: the SAME per-round D1 features drive a policy that is trained – online, from realized gain/cost, or offline, from logged history – rather than hand-tuned.

Correctness backbone (unchanged from the rest of the D-track): a learned scheduling POLICY can only ever change WHICH blocks get how much budget THIS round – never what a block’s update computes, never the accept/reject monotone-F gate D2/D3 already enforce. Any interleaving of partial E-steps and per-block conditional M-steps is coordinate ascent on the same Neal-Hinton free energy F regardless of which policy chose the interleaving, so F remains the audit receipt; a bad learned decision costs speed (a wasted round, a slow warmup), never correctness. This module therefore never runs its own EM and never touches accepted/objective bookkeeping – it only chooses the budget_fraction knob mixle.inference.block_em.run_block_em()’s existing greedy ranking (mixle.inference.block_em._select_active()) already accepts, so “learned” and “greedy” are two policies over the exact same scheduling loop, not two different loops.

Two learning modes:

  • Online bandits (BanditController) – reuses mixle.task.bandit (this codebase’s existing multi-armed-bandit module, built for exactly this “pick an arm/observe-a-reward, no offline data needed” loop) rather than reimplementing UCB1/Thompson sampling. The action space is a small discretized set of budget_fraction levels (an “arm”); the reward is the round’s REALIZED Q-gain-per-cost. Learns round-by-round DURING a fit, and carries over (the same policy object, still adapting) across multiple fits if the caller reuses it – see the offline-vs-warm-start framing in mixle.tests.conditional_jit_controller_test.

  • Offline DesignModel (DesignModelController) – reuses mixle.task.edge.DesignModel (this codebase’s existing design-space surrogate: a GP fitted on logged (point, quality, fingerprint) rows, warm-startable across different but related tasks via its fingerprint machinery – see mixle.task.edge’s own docstring) rather than reinventing an offline contextual bandit. budget_fraction is treated as the (continuous, 1-D) design point; the current round’s aggregated D1 features (ControllerState) are the fingerprint DesignModel already conditions proposals on, so a DesignModel trained on logged rounds from OTHER fit problems can propose a good budget for a brand-new, held-out problem with ZERO online exploration.

Action-type registry, not a closed enum: ActionType lists every action kind the D-track roadmap names for D5 (BLOCK_SELECTION, BUDGET_ALLOCATION, plus the future STRUCTURE_EDIT – H3’s structure-edit schedule / G2’s projections / the evolve-ops population search – and BACKEND_CHOICE – D6’s compile-economics RespecializationDecision). Only BLOCK_SELECTION and BUDGET_ALLOCATION are REAL, implemented action types today (and in this module they are the SAME knob: choosing a budget_fraction is exactly what determines which blocks D3’s existing ranking selects). STRUCTURE_EDIT/BACKEND_CHOICE are documented extension points (see ACTION_TYPE_REGISTRY), not wired to real H3/G2/D6 machinery here – that wiring is future work, explicitly out of scope for D5 per the roadmap item.

Reusable “brain” note (the roadmap’s own phrase: D5 “shares its brain with F5 and the I1/J1 method pickers”): LearnedController is deliberately generic over its own state/action/reward types (Generic[StateT, ActionT], a plain select_action(state) -> action / update(state, action, realized_gain, realized_cost) -> None surface with no block-EM-specific assumptions baked into the base class). A future F5 (backend/method-picker for some other stage) or I1/J1 (other “which method” pickers named in the roadmap, not built here) could instantiate BanditController or DesignModelController directly against THEIR OWN state/action types by subclassing LearnedController the same way this module’s two concrete controllers do – reusing the bandit/DesignModel wiring pattern without needing block-EM’s ControllerState/ControllerAction dataclasses at all. Nothing here builds F5/I1/J1 themselves.

class ActionType(*values)[source]

Bases: str, Enum

The D5 action-type registry (see module docstring): a registry, not a closed enum – future track items are expected to add new members and new LearnedController subclasses that consume them, without needing to touch this module’s two REAL, implemented action types.

class BanditController(*, budget_levels=_DEFAULT_BUDGET_LEVELS, algorithm='ucb1', ucb_c=1.0, seed=None)[source]

Bases: LearnedController[ControllerState, ControllerAction]

Online contextual-free multi-armed bandit over discretized budget_fraction levels.

Reuses mixle.task.bandit (UCB1 by default, or ThompsonGaussian) rather than reimplementing a bandit algorithm. Needs NO logged/offline data: it can start learning cold, round 1 of the very first fit it sees, from nothing but realized gain/cost – the “online bandits … no offline training data needed” mode from the roadmap item. The state (ControllerState) is accepted by select_action()/update() for interface symmetry with DesignModelController and any future contextual variant, but neither UCB1 nor ThompsonGaussian actually conditions on it (both are the classic context-FREE multi-armed setting) – a genuinely contextual variant (e.g. LinUCB keyed on ControllerState.as_vector()) is a natural extension of this class that this module does not need to build to satisfy “online bandits … given the state is a real feature vector” (the DesignModel mode already covers the contextual case, see its docstring).

Parameters:
select_action(state)[source]

Return this round’s action given state.

Parameters:

state (ControllerState)

Return type:

ControllerAction

update(state, action, realized_gain, realized_cost)[source]

Feed back the REALIZED gain/cost of action taken in state – the online-learning signal every concrete controller trains from.

Parameters:

  • state (ControllerState)

  • action (ControllerAction)

  • realized_gain (float)

  • realized_cost (float)

Return type:

None

class ControllerAction(action_type, budget_fraction, payload=<factory>)[source]

Bases: object

One controller decision for the current round.

budget_fraction is the real, implemented knob (see ACTION_TYPE_REGISTRY); payload is reserved for future STRUCTURE_EDIT/BACKEND_CHOICE action data and is unused by both controllers in this module.

Parameters:

  • action_type (ActionType)

  • budget_fraction (float)

  • payload (dict)

class ControllerState(round_index, n_eligible, mean_residual, mean_q_gain, mean_cost, score_spread)[source]

Bases: object

One round’s controller-visible state: D1 features aggregated over the eligible blocks.

This IS “features = D1 reports” from the roadmap item, aggregated to a fixed-size vector (one row per ROUND, not per node) because the eligible block set can change size round to round (freezing, zero-weight collapse) while both the bandit arm space and DesignModel’s fingerprint need a fixed dimension.

Parameters:
as_vector()[source]

The fixed-length (STATE_FEATURE_DIM) feature vector, e.g. for a DesignModel fingerprint or any future contextual (LinUCB-style) bandit.

Return type:

tuple[float, float, float, float, float]

classmethod from_scores(round_index, eligible, gain_per_cost, cost, residual, q_gain)[source]

Build a ControllerState from the same per-block dicts D3’s greedy scheduler already computes each round (mixle.inference.block_em._block_scores()) – no separate feature-extraction pass over the tree is needed.

Parameters:
Return type:

ControllerState

class DesignModelController(*, bounds=(0.05, 1.0), default_budget=0.5, design=None, seed=None)[source]

Bases: LearnedController[ControllerState, ControllerAction]

Offline controller wrapping mixle.task.edge.DesignModel.

Fit on logged (state, action, gain, cost) tuples collected across MANY prior fits (pass a pre-populated design= to warm-start from history, or let one accumulate rows via repeated update() calls across several fits before relying on select_action()). Treats budget_fraction as a 1-D continuous design point and ControllerState.as_vector() as the DesignModel fingerprint – DesignModel’s own cross-task warm-start machinery (see mixle.task.edge) is exactly “condition the proposal on which task this is”, which is what makes this mode work on a NEW, held-out fit problem’s state vector with no online exploration of that new problem at all, given a DesignModel already trained on OTHER problems.

Before at least two logged rows exist, select_action() cannot fit a GP and falls back to default_budget (an honest cold-start, not a fabricated proposal).

Parameters:
select_action(state)[source]

Return this round’s action given state.

Parameters:

state (ControllerState)

Return type:

ControllerAction

update(state, action, realized_gain, realized_cost)[source]

Feed back the REALIZED gain/cost of action taken in state – the online-learning signal every concrete controller trains from.

Parameters:

  • state (ControllerState)

  • action (ControllerAction)

  • realized_gain (float)

  • realized_cost (float)

Return type:

None

class LearnedController[source]

Bases: Generic[StateT, ActionT]

Generic state -> action learned scheduling policy (D5).

A drop-in alternative to a hand-written greedy heuristic: anywhere a scheduler currently reads some per-round state and applies a FIXED rule to pick an action, a LearnedController can read the same state and apply a TRAINED rule instead, updating from the REALIZED outcome after the fact. See the module docstring’s “reusable brain” note – this base class carries no block-EM-specific assumptions, so a future F5/I1/J1 method-picker item could subclass it directly for its own StateT/ActionT.

select_action(state)[source]

Return this round’s action given state.

Parameters:

state (StateT)

Return type:

ActionT

update(state, action, realized_gain, realized_cost)[source]

Feed back the REALIZED gain/cost of action taken in state – the online-learning signal every concrete controller trains from.

Parameters:

  • state (StateT)

  • action (ActionT)

  • realized_gain (float)

  • realized_cost (float)

Return type:

None