Batched solves#

Batched iLQR solves: many problems that share a model, e.g. one input trajectory per recorded trial.

A batch is described as in jax.vmap: in_axes is a pytree prefix of the arguments whose leaves are 0 (the leaf has a leading batch axis) or None (the leaf is shared by the whole batch). Three strategies map the solver over the batch:

  • “vmap”: one vectorised solve. The solver’s while loops stop once every batch element has converged, so the batch costs about as many iterations as its slowest element. Usually the fastest choice on GPU.

  • “map”: lax.map, one element (or one chunk of batch_size elements, vectorised) at a time, so slow elements only hold up their own chunk and memory stays bounded. Usually the fastest choice on CPU.

  • “shard”: shard_map over a 1-D mesh of devices (by default jax.devices()), with “vmap” or “map” (inner) on each device, so each device stops on its own shard. The batch is padded to a multiple of the number of devices. After jax.distributed.initialize(), jax.devices() spans every host, so the same call runs on a multi-host cluster.

Slow elements (stragglers) cost the most under “vmap” and “shard”. When the same batch is solved repeatedly, e.g. once per training epoch, return_info=True gives the iterations each element needed. Sorting the batch by them and solving with “map” in chunks of batch_size groups similar elements together. Starting each solve from the previous solution (Us_init) usually cuts the iterations.

Gradients of batched_dilqr() are the implicit gradients of diffilqrax.diff_ilqr.dilqr(), attached after the batched solve, so jax.grad of a loss summed over the batch works with every strategy without differentiating the solver loops.

diffilqrax.batch.batch_map(fn, in_axes=0, strategy='vmap', batch_size=None, devices=None, inner='vmap')[source]#

Map fn over the leading axis of the arguments selected by in_axes.

Parameters:

  • fn (Callable) – Function of one batch element.

  • in_axes (Any) – Pytree prefix of the arguments with leaves 0 (batched) or None (shared), as in jax.vmap, by default 0 (every argument batched).

  • strategy (str) – “vmap” (default), “map” or “shard”, see the module docstring.

  • batch_size (Optional[int]) – Chunk size of lax.map for strategy “map” (and inner “map”); by default one element at a time.

  • devices (Optional[Sequence[Device]]) – Devices of strategy “shard”, by default jax.devices().

  • inner (str) – Mapping on each device for strategy “shard”: “vmap” (default) or “map”.

Returns:

Function of the batched arguments whose outputs have a leading batch axis.

Return type:

Callable

diffilqrax.batch.batched_ilqr_solver(model, params, Us_init, in_axes=(0, 0), parallel=False, strategy='vmap', batch_size=None, devices=None, inner='vmap', **solver_kwargs)[source]#

Solve a batch of iLQR problems that share model.

Parameters:

  • model (System) – The system model, shared by the batch.

  • params (iLQRParams) – The iLQR parameters, with a leading batch axis on the leaves selected by in_axes[0].

  • Us_init (Array) – Initial inputs, [BxTxm] if batched or [Txm] if shared (in_axes[1] = None).

  • in_axes (Any) – Batch axes of (params, Us_init) as in jax.vmap (0 or None leaves), by default (0, 0). For example (iLQRParams(x0=None, theta=theta_axes), None) batches only the theta leaves marked 0 in theta_axes.

  • parallel (bool) – Use diffilqrax.parallel_ilqr.pilqr_solver() instead of diffilqrax.ilqr.ilqr_solver(), by default False.

  • strategy (str) – See batch_map().

  • batch_size (Optional[int]) – See batch_map().

  • devices (Optional[Sequence[Device]]) – See batch_map().

  • inner (str) – See batch_map().

  • **solver_kwargs – Keyword arguments of the solver, e.g. max_iter, kkt_tol, linesearch. With return_info=True the outputs end with the per-problem convergence flags and iteration counts.

Returns:

The solver outputs (see diffilqrax.ilqr.ilqr_solver()), each with a leading batch axis.

Return type:

Tuple

diffilqrax.batch.batched_dilqr(model, params, Us_init, in_axes=(0, 0), parallel=False, strategy='vmap', batch_size=None, devices=None, inner='vmap', exact_hessian=True, adjoint='parallel', **solver_kwargs)[source]#

Differentiable solutions of a batch of iLQR problems that share model.

The batch is solved by batched_ilqr_solver() without gradients. The implicit gradients of diffilqrax.diff_ilqr.dilqr() are then attached to each solution, one LQR solve per problem, mapped with vmap (strategy “vmap”, or “shard” with inner “vmap”) or lax.map (otherwise, chunked by batch_size). Reverse mode therefore never enters the solver loops or shard_map.

Parameters:
Returns:

The optimal state-control trajectories [Bx(T+1)x(n+m)], with implicit gradients wrt params.

Return type:

Array