jaxdp

Dynamic programming provides a framework for studying reinforcement learning (RL). Dynamic programming forms the foundation of many modern RL algorithms and helps us understand the theory behind the operators used in value estimation and policy improvement.

jaxdp provides JAX-compatible operators, mappings, and MDPs for studying dynamic programming. We can combine these components to implement algorithms and study their behavior across different MDPs.

Components

In jaxdp, most operators and mappings take an MDP dataclass as an explicit argument. An MDP (jaxdp.MDP) contains JAX arrays for the transition probabilities $P$, rewards $R$, initial distribution $\mu$, and termination indicators. jaxdp.make provides predefined recipes for environments such as Frozen Lake, Garnet, and Four Rooms. These are selected by name, following the same convention as Gymnasium’s make.

In [1]: import jaxdp

In [2]: mdp = jaxdp.make("garnet")
In [3]: mdp.transition.shape # (action, next_state, state)
Out[3]: (4, 10, 10)

Operators use the MDP through this shared interface. Let $\mathcal{T}_V^*$ and $\mathcal{T}_Q^*$ be the Bellman optimality operators on state-value and action-value functions, respectively. jaxdp.bellman_opt_op provides .v and .q methods for computing $\mathcal{T}_V^*v$ and $\mathcal{T}_Q^*q$. The suffix distinguishes state-value and action-value representations. This convention extends to mappings such as reward, where reward.s(mdp, policy) returns the state reward $r^\pi: S \to \mathbb{R}$ and reward.sa(mdp) returns the state-action reward $r: S \times A \to \mathbb{R}$.

In [4]: jaxdp.bellman_opt_op.v(mdp, v, gamma).shape
Out[4]: (10,)
In [5]: jaxdp.bellman_opt_op.q(mdp, q, gamma).shape
Out[5]: (4, 10)

In [6]: jaxdp.reward.s(mdp, policy).shape
Out[6]: (10,)
In [7]: jaxdp.reward.sa(mdp).shape
Out[7]: (4, 10)

With the Bellman optimality operator, we can implement value iteration through the recurrence $v_{k+1} = \mathcal{T}_V^*v_k$. jaxdp.ValueIteration implements this recurrence and, together with several other planning algorithms in the library, follows the interface shown below.

In [8]: algo = jaxdp.ValueIteration(gamma=gamma)
In [9]: alg_state = algo.init(mdp)
In [10]: alg_state = algo.update(mdp, alg_state)

In jaxdp, the variables that an algorithm carries between iterations are grouped in a dataclass called State. algo.init(mdp) creates the initial state, and each call to algo.update(mdp, alg_state) returns a new State for the next iteration. For value iteration, the state contains the current estimate v_val, initialized to $v_0 = 0$. More elaborate updates retain additional variables through the same interface.

For example, jaxdp.MomentumValueIteration includes a momentum term proportional to $v_k - v_{k-1}$. Computing this difference requires the previous estimate alongside the current one, so State retains both.

@chex.dataclass(frozen=True)
class State:
    v_val: jax.Array       # Current value estimate
    prev_v_val: jax.Array  # Previous value estimate

For the full API and available environments, see the component reference.

Leveraging JAX

By passing the algorithm state explicitly between updates, we can express a complete run with jax.lax.fori_loop. Since the update rules in jaxdp operate on JAX arrays, we can also apply automatic vectorization to the complete run through vmap, a transformation that batches a function over selected input axes. This lets us evaluate the same algorithm across multiple seeds and discount factors, with jit providing just-in-time compilation of the batched computation.

@jax.jit
@partial(jax.vmap, in_axes=(0, None))
@partial(jax.vmap, in_axes=(None, 0))
def run(key, gamma):
    mdp = make(key)
    algo = jaxdp.ValueIteration(gamma=gamma)
    state = jax.lax.fori_loop(
        0, 1_000, lambda _, s: algo.update(mdp, s), algo.init(mdp)
    )
    return state.v_val

Once we vmap and jit the run function, we can evaluate it across multiple seeds and discount factors.

In [11]: gammas = jnp.array([0.9, 0.95, 0.99, 0.999])
In [12]: keys = jax.random.split(jax.random.key(0), 100)

In [13]: vals = run(keys, gammas)
In [14]: vals.shape
Out[14]: (100, 4, ...)

In the above example, the inner vmap evaluates all four discount factors for a fixed seed, and the outer vmap repeats this sweep across 100 seeds. Each combination runs 1,000 updates, so vals contains the final value functions arranged along (seed, discount, state) axes.

To compare execution times across algorithms, we batch $M = 1{,}000$ discount factors for each of four fixed MDPs. With $N = 1{,}000$ updates per discount factor, each measurement covers $MN = 10^6$ algorithm updates. The environments are compiled separately because their state and action dimensions differ.

We run each batch five times on an Apple M5 CPU and measure the time in seconds to complete one million updates across the discount factors. We plot the median runtime as a bar and add a capped interval for the interquartile range to show variation between runs. To compare throughput across environments, we convert the median runtimes into updates per second and use a shared color scale.

Across the four environments, jaxdp averages less than one second per million updates with each of the four algorithms.

Comparing convergence

jaxdp separates the update rules from the MDPs and evaluation operators, letting us reuse the same experiment loop across algorithms. We can add new algorithms by implementing the init and update methods.

The functional approach adopted in jaxdp lets us use init and update to compare algorithm convergence within a common experimental setup. For instance, we run four algorithms on four MDPs, namely Frozen Lake, Four Rooms, Long Forest, and Garnet, for 500 steps at $\gamma = 0.99$. The figure below shows how the Bellman residual and value error change over these steps.

Curve colors indicate the value of the selected metric after 500 updates (the final error).

Anderson acceleration can reduce the number of updates needed for convergence, although each update requires more computation. Whether this improvement carries over to the corresponding estimator remains a statistical question.

Conclusion

jaxdp provides a setting for investigating these questions at the operator level. By working directly with the operators and using optimal value functions and policies as references, we can examine how RL algorithms approach these solutions and understand what drives their convergence. This analysis can guide the design of new estimators, whose statistical properties we can then investigate theoretically and empirically. For instance, algorithms derived for optimization frequently inspired their RL counterparts. For more rigorous coverage of the connection between optimization and RL algorithms, see Ok et al. (2025).

Try it

jaxdp requires Python 3.11 or newer and JAX 0.8 or newer.

uv add jaxdp

To install via pip, use pip install jaxdp.

The component reference lists the public API and available MDP recipes. Contributors can install the development dependencies with uv add "jaxdp[dev]". Citation metadata is available in CITATION.cff.