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.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.

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.