Utils¶

This section contains the API documentation for utility functions and helpers.

General Utilities¶

The utils module contains general-purpose utility functions.

skagent.utils.apply_fun_to_vals(fun, vals)¶

Applies a function to the arguments defined in vals. This is equivalent to fun(**vals), except that vals may contain keys that are not named arguments of fun.

Parameters:

fun (callable)
vals (dict)

skagent.utils.compute_gradients_for_tensors(tensors_dict, wrt, create_graph=False)¶

Compute gradients for a dictionary of tensors with respect to variables.

This function computes gradients using PyTorch’s autograd. It is used by the gradient methods on BellmanPeriod (grad_reward_function, grad_transition_function, and grad_pre_state_function).

For batched inputs, this computes the diagonal of the Jacobian: for each batch element i, we compute ∂target[i]/∂var[i]. The implementation uses grad(target.sum(), var), which yields the correct diagonal whenever target[i] depends only on var[i] (sample-independence). This assumption holds when each batch element is computed independently, which is the case for element-wise reward and transition functions in this library.

For scalar inputs, .sum() is a no-op, so the same code path handles both cases without branching.

Parameters:

tensors_dict (dict[str, Tensor]) – Dictionary mapping symbol names to tensors to compute gradients for.
wrt (dict[str, Tensor]) – Dictionary of variables to compute gradients with respect to. Keys are variable names, values are tensors with requires_grad=True.
create_graph (bool) – If True, the graph of the derivative is constructed, allowing higher-order derivatives and end-to-end training. Default: False.

Returns:

Nested dictionary of gradients for each tensor symbol and variable: {tensor_sym: {var_name: gradient}}. A variable with no computational graph path to the target receives a zero tensor, which is the value of the partial derivative under structural independence.

Return type:

dict[str, dict[str, Tensor]]

Raises:

ValueError – If a wrt tensor does not have requires_grad=True.

skagent.utils.compute_parameter_difference(params1, params2)¶: Compute the L2 norm of the difference between two parameter vectors.

skagent.utils.create_vectorized_function_wrapper_with_mapping(lambda_func, param_to_column)¶

Create a vectorized wrapper that automatically maps lambda parameters to correct tensor columns based on a parameter-to-column mapping.

Parameters:

lambda_func – Original lambda function
param_to_column – A mapping from parameter names to columns of a tensor

skagent.utils.extract_parameters(network)¶: Extract all parameters from a PyTorch network into a flat tensor.

skagent.utils.fischer_burmeister(a, h, eps=1e-12)¶

Compute the Fischer-Burmeister function for smooth complementarity.

The Fischer-Burmeister function replaces the complementarity conditions \(a \geq 0,\; h \geq 0,\; ah = 0\) with the equivalent smooth equation:

\[\text{FB}(a, h) = a + h - \sqrt{a^2 + h^2} = 0\]

This is differentiable everywhere, unlike \(\min(a, h) = 0\).

Following Maliar et al. (2021, JME) equation (25).

Parameters:

a (Tensor) – First argument (e.g., slack variable \(1 - c/w\)).
h (Tensor) – Second argument (e.g., unit-free Lagrange multiplier).
eps (float) – Regularization constant added inside the square root to keep the gradient finite at the origin. At the default eps=1e-12, FB(0, 0) = -sqrt(eps) ≈ -1e-6 rather than exactly zero. This is below typical convergence tolerances but should be accounted for in tests or with very tight tolerances.

Returns:

Fischer-Burmeister residual. Approximately zero when the complementarity conditions are satisfied.

Return type:

Tensor

skagent.utils.param_names(fun)¶

The names of fun’s parameters, in order.

Equivalent to tuple(inspect.signature(fun).parameters), memoized on the function object. inspect.signature costs on the order of ten microseconds, and the block machinery asks the same functions for their argument names once per variable per pass.

Parameters:: fun (callable)
Return type:: tuple of str

skagent.utils.plot_block_diagram(block, title=None, calibration=None, figsize=(9, 7), discount=None)¶

Render block’s model diagram into a new matplotlib figure.

display() shows the diagram in a notebook and returns the image; this puts the image on a figure instead, which is what a script or a documentation gallery needs. It also avoids display’s side effect of emitting the SVG through IPython, which in a plain script prints a line of noise.

Parameters:

block (Block) – The block to draw.
title (str, optional) – Figure title. Omitted if not given.
calibration (dict, optional) – Passed to visualize(), where it says which symbols are static rather than dynamic. Defaults to empty.
figsize (tuple of float) – Figure size in inches, for diagrams that are wider than they are tall or the reverse.
discount (str, optional) – The calibration symbol serving as the discount factor, passed to visualize().

skagent.utils.reconcile(vec_a, vec_b)¶: Returns a new vector with the values of vec_b but with the object type and shape of vec_a.

skagent.utils.takes_arguments(fun)¶

Whether fun accepts at least one argument.

For a plain Python function the answer is read off the code object, which is constant-time and needs no memo. Callers asking this question are often handed a function built fresh for the call – a decision rule wrapping a candidate action, say – for which a memo keyed on the function object would miss every time.

Parameters:: fun (callable)
Return type:: bool

Example Usage¶

Calling a Function from a Dictionary of Values¶

apply_fun_to_vals calls a function using only the entries of a dictionary that match the function’s named arguments, ignoring the rest:

from skagent.utils import apply_fun_to_vals

def transition(a, b):
    return a + b

# The extra key "c" is ignored
apply_fun_to_vals(transition, {"a": 1.0, "b": 2.0, "c": 99.0})  # 3.0

Smooth Complementarity Conditions¶

fischer_burmeister replaces the complementarity conditions \(a \geq 0,\; h \geq 0,\; ah = 0\) with a single smooth equation, which is useful for occasionally binding constraints in loss functions:

import torch

from skagent.utils import fischer_burmeister

a = torch.tensor([0.0, 1.0, 3.0])
h = torch.tensor([2.0, 0.0, 4.0])
fischer_burmeister(a, h)  # tensor([0., 0., 2.])

Note

Grid construction tools (Grid, make_grid, cartesian_product) live in skagent.grid and are documented on the Algorithms page.