Sampling oracles

banditry.sampling_oracles.sampler.Sampler

Bases: ABC

Posterior-sampling contract used by TSAgent for Thompson sampling.

A Sampler takes a template ValueFunction together with the data observed so far and returns a copy of the model whose weights are a single draw from the (approximate) posterior over weights; TSAgent then optimises the sampled model's predictions to pick the next action. Instances are constructed from TSConfig.sampler_config keyword arguments (see :class:~banditry.sampling_oracles.langevin_sampler.LangevinSampler and :class:~banditry.sampling_oracles.nuts_sampler.NUTSSampler for the accepted keys).

Attributes:

Name	Type	Description
generator	Generator | None	Optional torch.Generator driving all stochastic choices, for reproducible sampling.

sample(model, Xc, Xe, y, nll=None) abstractmethod

Draw one posterior sample of the model weights.

Parameters:

Name	Type	Description	Default
model	ValueFunction	Template ValueFunction; implementations work on a deep copy and leave the original unmodified.	required
Xc	FloatTensor | None	Continuous features of shape (n, model.num_cont), or None when the model has no continuous inputs.	required
Xe	LongTensor | None	Categorical (enum) features of shape (n, model.num_enum), or None when the model has no categorical inputs.	required
y	FloatTensor	Observed targets of shape (n, model.num_out).	required
nll	Callable[..., FloatTensor] | None	Optional negative-log-likelihood callable (pred, target, obs_std, **kwargs) -> FloatTensor; when None the Gaussian NLL is used.	None

Returns:

Type	Description
ValueFunction	A new ValueFunction with weights set to a single posterior draw, in eval
ValueFunction	mode and with its y-scaler fitted to y.

Source code in src/banditry/sampling_oracles/sampler.py

@abstractmethod
def sample(
    self,
    model: ValueFunction,
    Xc: FloatTensor | None,
    Xe: LongTensor | None,
    y: FloatTensor,
    nll: Callable[..., FloatTensor] | None = None,
) -> ValueFunction:
    """Draw one posterior sample of the model weights.

    Args:
        model: Template ``ValueFunction``; implementations work on a deep copy and
            leave the original unmodified.
        Xc: Continuous features of shape ``(n, model.num_cont)``, or ``None`` when the
            model has no continuous inputs.
        Xe: Categorical (enum) features of shape ``(n, model.num_enum)``, or ``None``
            when the model has no categorical inputs.
        y: Observed targets of shape ``(n, model.num_out)``.
        nll: Optional negative-log-likelihood callable
            ``(pred, target, obs_std, **kwargs) -> FloatTensor``; when ``None`` the
            Gaussian NLL is used.

    Returns:
        A new ``ValueFunction`` with weights set to a single posterior draw, in eval
        mode and with its y-scaler fitted to ``y``.
    """

__call__(model, Xc, Xe, y, nll=None)

Alias for :meth:sample.

Source code in src/banditry/sampling_oracles/sampler.py

def __call__(
    self,
    model: ValueFunction,
    Xc: FloatTensor | None,
    Xe: LongTensor | None,
    y: FloatTensor,
    nll: Callable[..., FloatTensor] | None = None,
) -> ValueFunction:
    """Alias for :meth:`sample`."""
    return self.sample(model, Xc, Xe, y, nll=nll)

total_variation(**kwargs)

Diagnostic hook for sampler-quality metrics; the base implementation returns 0.

Source code in src/banditry/sampling_oracles/sampler.py

def total_variation(self, **kwargs):
    """Diagnostic hook for sampler-quality metrics; the base implementation returns 0."""
    return 0

banditry.sampling_oracles.langevin_sampler

LangevinSampler(step_size=welling_teh_schedule(a=0.001), prior_precision=0.0001, temperature=1.0, precondition=False, precond_alpha=0.99, precond_eps=1e-05, batch_size=64, min_batches=10, num_epochs=10, burn_in=10, init_obs_noise=1.0, min_obs_noise=1e-06, max_obs_noise=1.0, generator=None)

Bases: Sampler

Sample TS model weights from an approximate posterior using SGLD.

Runs stochastic gradient Langevin dynamics (Welling & Teh, 2011) on a deep copy of the model: minibatch gradient steps on the negative log posterior with Gaussian noise injected by the :class:~banditry.optimisation_oracles.sgld.SGLD optimizer. The observation-noise standard deviation is learned jointly as a per-output parameter (excluded from the weight prior). The total number of updates is the maximum of num_epochs epochs, min_batches updates, and burn_in + 1 epochs, so at least one post-burn-in epoch always runs; one post-burn-in state is kept via streaming uniform (reservoir) selection and returned as the posterior draw.

The constructor keyword arguments below are exactly the valid keys of TSConfig.sampler_config when sampler="langevin".

Parameters:

Name	Type	Description	Default
step_size	float | Callable[[int], float]	Learning-rate schedule: either a constant float or a callable (t, n_burn_in) -> lr such as the one returned by :func:welling_teh_schedule (the default).	welling_teh_schedule(a=0.001)
prior_precision	float	Precision of the isotropic zero-mean Gaussian prior over model weights, applied as weight decay on the gradients.	0.0001
temperature	float	Scales the variance of the injected Langevin noise; 1.0 targets the true posterior, smaller values sharpen it, and 0 reduces the update to plain SGD.	1.0
precondition	bool	If True, use RMSProp-style diagonal preconditioning (pSGLD): gradient step and noise are rescaled by an EMA of squared gradients.	False
precond_alpha	float	EMA decay rate of the squared-gradient accumulator used by the preconditioner.	0.99
precond_eps	float	Numerical-stability constant added to the preconditioner denominator.	1e-05
batch_size	int	Minibatch size (capped at the dataset size).	64
min_batches	int	Minimum total number of gradient updates to perform.	10
num_epochs	int	Target number of passes over the data.	10
burn_in	int	Number of initial epochs discarded before model states become eligible for selection as the returned sample.	10
init_obs_noise	float	Initial observation-noise standard deviation (in standardised-y units).	1.0
min_obs_noise	float	Lower clamp for the learned observation noise.	1e-06
max_obs_noise	float	Upper clamp for the learned observation noise.	1.0
generator	Generator | None	Optional torch.Generator making noise injection and sample selection reproducible.	None

Source code in src/banditry/sampling_oracles/langevin_sampler.py

def __init__(
    self,
    step_size: float | Callable[[int], float] = welling_teh_schedule(a=1e-3),  # noqa: B008 — deterministic closure
    prior_precision: float = 1e-4,
    temperature: float = 1.0,
    precondition: bool = False,
    precond_alpha: float = 0.99,
    precond_eps: float = 1e-5,
    batch_size: int = 64,
    min_batches: int = 10,
    num_epochs: int = 10,
    burn_in: int = 10,
    init_obs_noise: float = 1.0,
    min_obs_noise: float = 1e-6,
    max_obs_noise: float = 1.0,
    generator: torch.Generator | None = None,
):

    if callable(step_size):
        self._step_size_fn: Callable[[int], float] = step_size
    else:
        _c = float(step_size)

        def _f(*args, **kwargs):
            return _c

        self._step_size_fn = _f
    self.prior_precision = float(prior_precision)
    self.temperature = float(temperature)
    self.precondition = bool(precondition)
    self.precond_alpha = float(precond_alpha)
    self.precond_eps = float(precond_eps)
    self.batch_size = int(batch_size)
    self.min_batches = int(min_batches)
    self.num_epochs = int(num_epochs)
    self.burn_in = int(burn_in)
    self.init_obs_noise = float(init_obs_noise)
    self.min_obs_noise = float(min_obs_noise)
    self.max_obs_noise = float(max_obs_noise)
    self.generator = generator

sample(model, Xc, Xe, y, nll=None)

Draw one posterior weight sample via SGLD.

Standardises y (and fits the model's x-scaler), then runs SGLD on a deep copy of model for the configured number of updates, keeping one post-burn-in state chosen uniformly at random as the posterior draw.

Parameters:

Name	Type	Description	Default
model	ValueFunction	Template ValueFunction; left unmodified.	required
Xc	FloatTensor | None	Continuous features of shape (n, model.num_cont), or None.	required
Xe	LongTensor | None	Categorical features of shape (n, model.num_enum), or None.	required
y	FloatTensor	Observed targets of shape (n, model.num_out).	required
nll	Callable[..., FloatTensor] | None	Optional NLL callable (pred, target, obs_std, **kwargs); it may return per-element values (reduced internally to a dataset total) or a scalar total loss. Defaults to the Gaussian NLL.	None

Returns:

Type	Description
ValueFunction	A new ValueFunction in eval mode carrying the sampled weights and the
ValueFunction	fitted y-scaler.

Source code in src/banditry/sampling_oracles/langevin_sampler.py

def sample(
    self,
    model: ValueFunction,
    Xc: FloatTensor | None,
    Xe: LongTensor | None,
    y: FloatTensor,
    nll: Callable[..., FloatTensor] | None = None,
) -> ValueFunction:
    """Draw one posterior weight sample via SGLD.

    Standardises ``y`` (and fits the model's x-scaler), then runs SGLD on a deep copy
    of ``model`` for the configured number of updates, keeping one post-burn-in state
    chosen uniformly at random as the posterior draw.

    Args:
        model: Template ``ValueFunction``; left unmodified.
        Xc: Continuous features of shape ``(n, model.num_cont)``, or ``None``.
        Xe: Categorical features of shape ``(n, model.num_enum)``, or ``None``.
        y: Observed targets of shape ``(n, model.num_out)``.
        nll: Optional NLL callable ``(pred, target, obs_std, **kwargs)``; it may
            return per-element values (reduced internally to a dataset total) or a
            scalar total loss. Defaults to the Gaussian NLL.

    Returns:
        A new ``ValueFunction`` in eval mode carrying the sampled weights and the
        fitted y-scaler.
    """

    nll_fn = _gaussian_nll if nll is None else nll

    Xc_t, Xe_t, y_t = self._prepare_xy(model, Xc, Xe, y)
    yscaler, y_scaled = self._fit_y_scaler(y_t)

    working_model = deepcopy(model)
    working_model.clear_y_scaler()
    if working_model.num_cont > 0:
        working_model.fit_x_scaler(Xc_t)
    working_model.train()

    Xc_scaled, Xe_ready = working_model.xtrans(Xc_t, Xe_t)

    num_data = y_scaled.shape[0]
    batch_size = min(self.batch_size, num_data)
    batch_size = max(1, batch_size)

    ds = TensorDataset(Xc_scaled, Xe_ready, y_scaled)
    dl = DataLoader(
        ds,
        batch_size=batch_size,
        shuffle=True,
        drop_last=False,
    )
    if len(dl) == 0:
        dl = DataLoader(ds, batch_size=num_data, shuffle=True, drop_last=False)

    updates_per_epoch = max(1, len(dl))
    min_updates_for_sampling = (self.burn_in + 1) * updates_per_epoch
    target_updates = max(
        self.num_epochs * updates_per_epoch,
        self.min_batches,
        min_updates_for_sampling,
    )
    burn_in_updates = self.burn_in * updates_per_epoch
    import banditry.logging_utils as log

    log.debug(f"Learning rate after burn in: {self._step_size_fn(burn_in_updates, n_burn_in=burn_in_updates)}")

    model_params = [p for p in working_model.parameters() if p.requires_grad]

    log_obs_noise = nn.Parameter(
        torch.full(
            (working_model.num_out,),
            math.log(self.init_obs_noise),
            dtype=y_scaled.dtype,
            device=y_scaled.device,
        )
    )
    optimizer = SGLD(
        [
            {"params": model_params, "prior_precision": self.prior_precision},
            {"params": [log_obs_noise], "prior_precision": 0.0},
        ],
        lr=self._step_size_fn(0, n_burn_in=burn_in_updates),
        temperature=self.temperature,
        precondition=self.precondition,
        precond_alpha=self.precond_alpha,
        precond_eps=self.precond_eps,
        generator=self.generator,
    )

    selected_state: dict[str, torch.Tensor] | None = None
    post_burn_seen = 0
    updates = 0

    while updates < target_updates:
        for bxc, bxe, by in dl:
            optimizer.zero_grad()

            pred = working_model.forward(bxc, bxe, _pre_scaled=True)

            obs_noise = torch.exp(log_obs_noise).clamp(min=self.min_obs_noise, max=self.max_obs_noise).view(1, -1)

            nll_value = nll_fn(
                pred,
                by,
                obs_noise,
                num_data=num_data,
                model=working_model,
            )
            if not torch.is_tensor(nll_value):
                nll_value = torch.as_tensor(nll_value, device=pred.device, dtype=pred.dtype)

            if nll_value.dim() == 0:
                loss = nll_value
            else:
                per_datapoint = nll_value.reshape(nll_value.shape[0], -1).sum(dim=1)
                loss = per_datapoint.mean() * num_data

            loss.backward()

            lr = self._step_size_fn(updates, n_burn_in=burn_in_updates)
            for group in optimizer.param_groups:
                group["lr"] = lr

            optimizer.step(noise=True)

            with torch.no_grad():
                log_obs_noise.clamp_(math.log(self.min_obs_noise), math.log(self.max_obs_noise))
            updates += 1

        if updates >= burn_in_updates:
            post_burn_seen += 1
            if selected_state is None:
                selected_state = deepcopy(working_model.state_dict())
            else:
                # replace with prob 1/post_burn_seen, stream uniform
                if self._draw_sample_id(post_burn_seen) == 0:
                    selected_state = deepcopy(working_model.state_dict())

    sampled_model = working_model
    if selected_state is not None:
        sampled_model.load_state_dict(selected_state, strict=True)

    sampled_model.set_y_scaler(yscaler)
    sampled_model.eval()
    return sampled_model

welling_teh_schedule(a, b=1.0, gamma=0.55, lr_floor=1e-05)

Build the polynomial step-size schedule of Welling & Teh (2011).

Implements the decaying schedule from Welling & Teh (2011), "Bayesian Learning via Stochastic Gradient Langevin Dynamics": eps_t = a * (b + t) ** (-gamma) during burn-in. After burn-in the rate is frozen at its final burn-in value and clamped from below by lr_floor, i.e. max(a * (b + n_burn_in) ** (-gamma), lr_floor).

Parameters:

Name	Type	Description	Default
a	float	Initial scale of the schedule; must be positive.	required
b	float	Offset delaying the start of the decay; must be non-negative.	1.0
gamma	float	Decay exponent; must lie in (0.5, 1] for the SGLD convergence conditions of Welling & Teh to hold.	0.55
lr_floor	float	Lower bound applied to the post-burn-in learning rate.	1e-05

Returns:

Type	Description
Callable[[int], float]	A callable (t, n_burn_in) -> lr mapping the update step t (given the total
Callable[[int], float]	number of burn-in updates n_burn_in) to a learning rate.

Raises:

Type	Description
ValueError	If a <= 0, b < 0, or gamma is outside (0.5, 1].

Source code in src/banditry/sampling_oracles/langevin_sampler.py

def welling_teh_schedule(
    a: float, b: float = 1.0, gamma: float = 0.55, lr_floor: float = 1e-5
) -> Callable[[int], float]:
    """Build the polynomial step-size schedule of Welling & Teh (2011).

    Implements the decaying schedule from Welling & Teh (2011), "Bayesian Learning via
    Stochastic Gradient Langevin Dynamics": ``eps_t = a * (b + t) ** (-gamma)`` during
    burn-in. After burn-in the rate is frozen at its final burn-in value and clamped from
    below by ``lr_floor``, i.e. ``max(a * (b + n_burn_in) ** (-gamma), lr_floor)``.

    Args:
        a: Initial scale of the schedule; must be positive.
        b: Offset delaying the start of the decay; must be non-negative.
        gamma: Decay exponent; must lie in ``(0.5, 1]`` for the SGLD convergence
            conditions of Welling & Teh to hold.
        lr_floor: Lower bound applied to the post-burn-in learning rate.

    Returns:
        A callable ``(t, n_burn_in) -> lr`` mapping the update step ``t`` (given the total
        number of burn-in updates ``n_burn_in``) to a learning rate.

    Raises:
        ValueError: If ``a <= 0``, ``b < 0``, or ``gamma`` is outside ``(0.5, 1]``.
    """
    # Polynomial step-size schedule from Welling & Teh (2011).

    if a <= 0:
        raise ValueError(f"a must be positive, got {a}")
    if b < 0:
        raise ValueError(f"b must be non-negative, got {b}")
    if not (0.5 < gamma <= 1.0):
        raise ValueError(f"gamma must be in (0.5, 1], got {gamma}")

    def _schedule(t: int, n_burn_in: bool) -> float:
        if t < n_burn_in:
            return a * (b + t) ** (-gamma)
        else:
            return max(a * (b + n_burn_in) ** (-gamma), lr_floor)

    return _schedule

banditry.sampling_oracles.nuts_sampler.NUTSSampler(num_samples=64, warmup_steps=256, target_accept_prob=0.8, max_tree_depth=10, adapt_step_size=True, adapt_mass_matrix=True, use_multinomial_sampling=True, prior_std=1.0, obs_noise_prior_loc=0.0, obs_noise_prior_scale=1.0, init_obs_noise=1.0, min_obs_noise=1e-06, max_obs_noise=1000.0, jit_compile=False, ignore_jit_warnings=True, disable_progbar=True, generator=None)

Bases: Sampler

Sample TS model weights from a posterior using Pyro NUTS.

Runs the No-U-Turn Sampler (Hoffman & Gelman, 2014) via Pyro's MCMC over all trainable model parameters plus a per-output observation-noise parameter, then returns the model with one randomly chosen posterior draw loaded. Weights receive an isotropic zero-mean Gaussian prior with standard deviation prior_std; the observation noise is parameterised through a sigmoid squashed into (min_obs_noise, max_obs_noise) with a Gaussian prior on the unconstrained value.

Requires the banditry[nuts] extra (installs pyro-ppl).

The constructor keyword arguments below are exactly the valid keys of TSConfig.sampler_config when sampler="nuts".

Parameters:

Name	Type	Description	Default
num_samples	int	Number of posterior draws collected after warmup; one is picked uniformly at random as the returned sample.	64
warmup_steps	int	Number of warmup (adaptation) iterations discarded before sampling.	256
target_accept_prob	float	Target acceptance probability for step-size adaptation.	0.8
max_tree_depth	int	Maximum doubling depth of the NUTS trajectory tree.	10
adapt_step_size	bool	Whether to adapt the leapfrog step size during warmup.	True
adapt_mass_matrix	bool	Whether to adapt the mass matrix during warmup.	True
use_multinomial_sampling	bool	Whether to draw states multinomially from the trajectory instead of using slice sampling.	True
prior_std	float	Standard deviation of the zero-mean Gaussian prior over model weights.	1.0
obs_noise_prior_loc	float	Prior mean of the log observation noise; clamped into [log(min_obs_noise), log(max_obs_noise)] and mapped to the unconstrained parameterisation.	0.0
obs_noise_prior_scale	float	Prior standard deviation of the unconstrained observation-noise parameter.	1.0
init_obs_noise	float	Initial observation-noise standard deviation (in standardised-y units), clamped into the allowed range.	1.0
min_obs_noise	float	Lower bound of the observation-noise range.	1e-06
max_obs_noise	float	Upper bound of the observation-noise range.	1000.0
jit_compile	bool	Whether Pyro should JIT-compile the potential function.	False
ignore_jit_warnings	bool	Whether to silence JIT tracer warnings.	True
disable_progbar	bool	Whether to hide the MCMC progress bar.	True
generator	Generator | None	Optional torch.Generator; seeds Pyro's RNG and drives the choice of the returned draw, for reproducibility.	None

Source code in src/banditry/sampling_oracles/nuts_sampler.py

def __init__(
    self,
    num_samples: int = 64,
    warmup_steps: int = 256,
    target_accept_prob: float = 0.8,
    max_tree_depth: int = 10,
    adapt_step_size: bool = True,
    adapt_mass_matrix: bool = True,
    use_multinomial_sampling: bool = True,
    prior_std: float = 1.0,
    obs_noise_prior_loc: float = 0.0,
    obs_noise_prior_scale: float = 1.0,
    init_obs_noise: float = 1.0,
    min_obs_noise: float = 1e-6,
    max_obs_noise: float = 1e3,
    jit_compile: bool = False,
    ignore_jit_warnings: bool = True,
    disable_progbar: bool = True,
    generator: torch.Generator | None = None,
):

    self.num_samples = int(num_samples)
    self.warmup_steps = int(warmup_steps)
    self.target_accept_prob = float(target_accept_prob)
    self.max_tree_depth = int(max_tree_depth)
    self.adapt_step_size = bool(adapt_step_size)
    self.adapt_mass_matrix = bool(adapt_mass_matrix)
    self.use_multinomial_sampling = bool(use_multinomial_sampling)
    self.prior_std = float(prior_std)
    self.obs_noise_prior_loc = float(obs_noise_prior_loc)
    self.obs_noise_prior_scale = float(obs_noise_prior_scale)
    self.init_obs_noise = float(init_obs_noise)
    self.min_obs_noise = float(min_obs_noise)
    self.max_obs_noise = float(max_obs_noise)
    self.jit_compile = bool(jit_compile)
    self.ignore_jit_warnings = bool(ignore_jit_warnings)
    self.disable_progbar = bool(disable_progbar)
    self.generator = generator

sample(model, Xc, Xe, y, nll=None)

Draw one posterior weight sample via NUTS MCMC.

Standardises y (and fits the model's x-scaler), builds a Pyro model whose log-density combines the weight and noise priors with -nll as the likelihood factor, runs NUTS for warmup_steps plus num_samples iterations, and loads one uniformly chosen posterior draw into a deep copy of model.

Parameters:

Name	Type	Description	Default
model	ValueFunction	Template ValueFunction; left unmodified.	required
Xc	FloatTensor | None	Continuous features of shape (n, model.num_cont), or None.	required
Xe	LongTensor | None	Categorical features of shape (n, model.num_enum), or None.	required
y	FloatTensor	Observed targets of shape (n, model.num_out).	required
nll	Callable[..., FloatTensor] | None	Optional NLL callable (pred, target, obs_std, **kwargs); its sum is used as the negative log-likelihood factor. Defaults to the Gaussian NLL.	None

Returns:

Type	Description
ValueFunction	A new ValueFunction in eval mode carrying the sampled weights and the
ValueFunction	fitted y-scaler.

Raises:

Type	Description
ValueError	If the model has no trainable parameters.
RuntimeError	If MCMC returns no posterior draws.

Source code in src/banditry/sampling_oracles/nuts_sampler.py

def sample(
    self,
    model: ValueFunction,
    Xc: FloatTensor | None,
    Xe: LongTensor | None,
    y: FloatTensor,
    nll: Callable[..., FloatTensor] | None = None,
) -> ValueFunction:
    """Draw one posterior weight sample via NUTS MCMC.

    Standardises ``y`` (and fits the model's x-scaler), builds a Pyro model whose
    log-density combines the weight and noise priors with ``-nll`` as the likelihood
    factor, runs NUTS for ``warmup_steps`` plus ``num_samples`` iterations, and loads
    one uniformly chosen posterior draw into a deep copy of ``model``.

    Args:
        model: Template ``ValueFunction``; left unmodified.
        Xc: Continuous features of shape ``(n, model.num_cont)``, or ``None``.
        Xe: Categorical features of shape ``(n, model.num_enum)``, or ``None``.
        y: Observed targets of shape ``(n, model.num_out)``.
        nll: Optional NLL callable ``(pred, target, obs_std, **kwargs)``; its sum is
            used as the negative log-likelihood factor. Defaults to the Gaussian NLL.

    Returns:
        A new ``ValueFunction`` in eval mode carrying the sampled weights and the
        fitted y-scaler.

    Raises:
        ValueError: If the model has no trainable parameters.
        RuntimeError: If MCMC returns no posterior draws.
    """

    nll_fn = _gaussian_nll if nll is None else nll

    Xc_t, Xe_t, y_t = self._prepare_xy(model, Xc, Xe, y)
    yscaler, y_scaled = self._fit_y_scaler(y_t)

    working_model = deepcopy(model)
    if working_model.num_cont > 0:
        working_model.fit_x_scaler(Xc_t)
    Xc_prescaled, Xe_prescaled = working_model.xtrans(Xc_t, Xe_t)
    working_model.clear_y_scaler()
    working_model.eval()

    param_specs: list[tuple[str, str, torch.Tensor]] = []
    init_values: dict[str, torch.Tensor] = {}
    for idx, (param_name, param_value) in enumerate(working_model.named_parameters()):
        if not param_value.requires_grad:
            continue
        site_name = self._site_name(idx)
        detached = param_value.detach()
        param_specs.append((site_name, param_name, detached))
        init_values[site_name] = detached.clone()
    if len(param_specs) == 0:
        raise ValueError("No trainable parameters found for posterior sampling.")

    num_data = y_scaled.shape[0]

    init_noise = min(max(self.init_obs_noise, self.min_obs_noise), self.max_obs_noise)
    init_values["raw_obs_noise"] = torch.full(
        (working_model.num_out,),
        self._noise_to_raw(init_noise),
        device=y_scaled.device,
        dtype=y_scaled.dtype,
    )

    pyro.clear_param_store()

    def _model():
        sampled_params: dict[str, torch.Tensor] = {}
        for site_name, param_name, param_value in param_specs:
            sampled_params[param_name] = pyro.sample(
                site_name,
                dist.Normal(
                    torch.zeros_like(param_value),
                    torch.full_like(param_value, self.prior_std),
                ).to_event(param_value.dim()),
            )

        raw_obs_noise = pyro.sample(
            "raw_obs_noise",
            dist.Normal(
                torch.full(
                    (working_model.num_out,),
                    self._raw_noise_prior_loc(),
                    device=y_scaled.device,
                    dtype=y_scaled.dtype,
                ),
                torch.full(
                    (working_model.num_out,),
                    self.obs_noise_prior_scale,
                    device=y_scaled.device,
                    dtype=y_scaled.dtype,
                ),
            ).to_event(1),
        )
        obs_noise = self._raw_to_noise(raw_obs_noise).view(1, -1)

        pred = functional_call(
            working_model,
            sampled_params,
            (Xc_prescaled, Xe_prescaled),
            {"_pre_scaled": True},
        )
        nll_value = nll_fn(
            pred,
            y_scaled,
            obs_noise,
            num_data=num_data,
            model=working_model,
            model_params=sampled_params,
        )
        if not torch.is_tensor(nll_value):
            nll_value = torch.as_tensor(nll_value, device=pred.device, dtype=pred.dtype)
        else:
            nll_value = nll_value.to(device=pred.device, dtype=pred.dtype)
        total_nll = nll_value.sum() if nll_value.dim() > 0 else nll_value
        pyro.factor("likelihood", -total_nll)

    if self.generator is not None:
        seed = int(torch.randint(0, 2**31 - 1, (1,), generator=self.generator).item())
        pyro.set_rng_seed(seed)

    kernel = NUTS(
        _model,
        target_accept_prob=self.target_accept_prob,
        max_tree_depth=self.max_tree_depth,
        adapt_step_size=self.adapt_step_size,
        adapt_mass_matrix=self.adapt_mass_matrix,
        use_multinomial_sampling=self.use_multinomial_sampling,
        jit_compile=self.jit_compile,
        ignore_jit_warnings=self.ignore_jit_warnings,
        init_strategy=init_to_value(values=init_values),
    )
    mcmc = MCMC(
        kernel,
        num_samples=self.num_samples,
        warmup_steps=self.warmup_steps,
        disable_progbar=self.disable_progbar,
    )
    mcmc.run()

    samples = mcmc.get_samples(group_by_chain=False)
    if len(samples) == 0:
        raise RuntimeError("Pyro MCMC did not return any posterior samples.")

    first_site = next(iter(samples))
    posterior_count = int(samples[first_site].shape[0])
    if posterior_count <= 0:
        raise RuntimeError("Pyro MCMC returned zero posterior draws.")

    sample_id = self._draw_sample_id(posterior_count)
    sampled_state = working_model.state_dict()
    for site_name, param_name, param_value in param_specs:
        sampled_state[param_name] = samples[site_name][sample_id].to(
            device=param_value.device,
            dtype=param_value.dtype,
        )
    working_model.load_state_dict(sampled_state, strict=True)
    working_model.set_y_scaler(yscaler)
    working_model.eval()
    return working_model

banditry.sampling_oracles.sampler.FeelGoodNLL(fg_lambda=1.0, fg_bound=10.0)

Bases: NLL

Feel-Good Thompson sampling loss factory (Zhang, 2021).

Implements the exploration reweighting of Zhang (2021), "Feel-Good Thompson Sampling for Contextual Bandits and Reinforcement Learning". The returned loss augments the Gaussian NLL with an optimism ("feel-good") bonus: for every distinct context observed so far, a short evolutionary search over the free (non-context) variables finds the current weight sample's maximum predicted value; that value is capped at fg_bound, weighted by the number of observations sharing the context, and fg_lambda times the total is subtracted from the NLL. Posterior draws are thereby tilted towards weights that believe a high value is attainable, adding optimism to Thompson sampling.

Parameters:

Name	Type	Description	Default
fg_lambda	float	Weight of the feel-good term. 0 recovers the plain Gaussian NLL; larger values add more optimism.	1.0
fg_bound	float	Upper cap applied to each context's best predicted value before weighting, preventing the bonus from dominating the likelihood. Must be non-negative.	10.0

Raises:

Type	Description
ValueError	If fg_bound is negative.

Source code in src/banditry/sampling_oracles/sampler.py

def __init__(self, fg_lambda: float = 1.0, fg_bound: float = 10.0):
    if fg_bound < 0:
        raise ValueError(f"fg_bound must be non-negative, got {fg_bound}")
    self.fg_lambda = fg_lambda
    self.fg_bound = fg_bound

__call__(fix_input, space, X)

Bind the observed contexts and return the feel-good loss callable.

Groups X by the context columns named in fix_input to enumerate the distinct contexts and their observation counts. The returned callable computes a scalar total loss: the Gaussian NLL summed over the dataset (rescaled to num_data) minus the feel-good bonus. It expects the sampler to pass model (and, for functional evaluation, model_params) as keyword arguments; without a model, or when fg_lambda == 0, it falls back to the per-element Gaussian NLL.

Parameters:

Name	Type	Description	Default
fix_input	dict	Context-variable mapping; only its keys are used, to group X.	required
space		DesignSpace used by the inner acquisition optimisation.	required
X	DataFrame	Raw observations collected so far; grouped by the context columns.	required

Returns:

Type	Description
Callable[..., FloatTensor]	A callable (pred, target, obs_std, **kwargs) -> FloatTensor implementing
Callable[..., FloatTensor]	the feel-good-adjusted negative log-likelihood.

Source code in src/banditry/sampling_oracles/sampler.py

def __call__(self, fix_input: dict, space, X: pd.DataFrame) -> Callable[..., FloatTensor]:
    """Bind the observed contexts and return the feel-good loss callable.

    Groups ``X`` by the context columns named in ``fix_input`` to enumerate the distinct
    contexts and their observation counts. The returned callable computes a scalar total
    loss: the Gaussian NLL summed over the dataset (rescaled to ``num_data``) minus the
    feel-good bonus. It expects the sampler to pass ``model`` (and, for functional
    evaluation, ``model_params``) as keyword arguments; without a model, or when
    ``fg_lambda == 0``, it falls back to the per-element Gaussian NLL.

    Args:
        fix_input: Context-variable mapping; only its keys are used, to group ``X``.
        space: ``DesignSpace`` used by the inner acquisition optimisation.
        X: Raw observations collected so far; grouped by the context columns.

    Returns:
        A callable ``(pred, target, obs_std, **kwargs) -> FloatTensor`` implementing
        the feel-good-adjusted negative log-likelihood.
    """
    from banditry.optimisation_oracles.gen_alg import EvolutionOpt
    from banditry.optimisation_subroutines.contextual_problem import ContextualProblem
    from banditry.optimisation_subroutines.objectives import ThompsonObjective

    context_names = list(fix_input.keys())

    if not context_names:
        contexts = [{}]
        counts = [len(X)]
    else:
        grouped = X.groupby(context_names).size().reset_index(name="count")
        contexts = grouped[context_names].to_dict("records")
        counts = grouped["count"].tolist()

    fg_lambda = self.fg_lambda
    fg_bound = self.fg_bound

    def nll_fn(pred: FloatTensor, target: FloatTensor, obs_std: FloatTensor, **kwargs) -> FloatTensor:
        base_nll = _gaussian_nll(pred, target, obs_std)

        model = kwargs.get("model")
        model_params = kwargs.get("model_params")

        if model is None or fg_lambda == 0:
            return base_nll

        eval_model = _ModelWrapper(model, model_params)
        num_data = kwargs.get("num_data", pred.shape[0])

        fg_values = []
        for ctx, count in zip(contexts, counts, strict=True):
            objective = ThompsonObjective(eval_model)
            problem = ContextualProblem(objective, space, ctx)
            opt = EvolutionOpt(space, pop=5, max_iters=2, verbose=False)
            rec = opt.optimise(problem)

            Xc_opt, Xe_opt = space.transform(rec)
            f_val = -eval_model(Xc_opt, Xe_opt)
            f_clipped = torch.clamp(f_val, max=fg_bound)
            fg_values.append(count * f_clipped.sum())

        fg_term = torch.stack(fg_values).sum()

        per_dp = base_nll.reshape(base_nll.shape[0], -1).sum(dim=1)
        total_base = per_dp.mean() * num_data

        return total_base - fg_lambda * fg_term

    return nll_fn