Utilities#

Includes utility functions for the project. Generic functions to generate data, seeds, etc.

diffilqrax.utils.keygen(key, nkeys)[source]#

Generate randomness that JAX can use by splitting the JAX keys.

Parameters:

  • key (jax.random.PRNGKey) – The random key for JAX.

  • nkeys (int) – Number of keys to generate.

Returns:

A tuple containing the new key for further generators and the key generator.

Return type:

tuple

diffilqrax.utils.initialise_stable_dynamics(key, n_dim, T, radii=0.6)[source]#

Generate a state matrix with stable dynamics (all eigenvalues inside the unit circle) discrete dynamics.

Parameters:

  • key (Tuple[int, int]) – Random key.

  • n_dim (int) – State dimensions.

  • T (int) – Time steps.

  • radii (float) – Spectral radius, by default 0.6.

Returns:

Matrix A with stable dynamics.

Return type:

Array

diffilqrax.utils.initialise_stable_time_varying_dynamics(key, n_dim, T, radii=0.6)[source]#

Generate a state matrix with stable dynamics (all eigenvalues inside the unit circle).

Parameters:

  • key (Tuple[int, int]) – Random key.

  • n_dim (int) – State dimensions.

  • T (int) – Time steps.

  • radii (float) – Spectral radius, by default 0.6.

Returns:

Matrix A with stable dynamics.

Return type:

Array

diffilqrax.utils.linearise(fun)[source]#

Find Jacobian with respect to x and u inputs.

Parameters:

fun (Callable) – Function with arguments (t, x, u, params).

Returns:

Jacobian tuple evaluated at args 1 and 2.

Return type:

Callable

diffilqrax.utils.quadratise(fun)[source]#

Find Hessian with respect to x and u inputs.

Parameters:

fun (Callable) – Function with arguments (t, x, u, params).

Returns:

Hessian tuple cross evaluated with args 1 and 2.

Return type:

Callable

diffilqrax.utils.time_map(fun)[source]#

Vectorise function in time. Assumes 0th-axis is time for x and u args of fun, the last arg (theta) of Callable function assumed to be time-invariant.

Parameters:

fun (Callable) – Function that takes args (t, x[Txn], u[Txm], theta).

Returns:

Vectorised function along args 1 and 2 0th-axis.

Return type:

Callable

diffilqrax.utils.smooth_l1(x, eps=1e-06)[source]#

Smooth approximation of the L1 norm, sum(sqrt(x^2 + eps) - sqrt(eps)).

Twice differentiable, so it can be used as an iLQR sparsity cost. Approaches sum(|x|) as eps -> 0; anneal eps by re-solving with smaller values, warm-starting each solve.

Parameters:

  • x (Array) – Input array.

  • eps (float) – Smoothing parameter, by default 1e-6.

Returns:

Smoothed L1 norm of x.

Return type:

Array