Metadata-Version: 2.4
Name: jaxdp
Version: 0.4.0
Summary: Exact dynamic programming for finite Markov decision processes in JAX
Author-email: Tolga Ok <T.Ok@tudelft.nl>
License-Expression: MIT
Project-URL: Repository, https://github.com/TolgaOk/jaxdp
Keywords: dynamic programming,jax,markov decision process,reinforcement learning
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENCE.txt
Requires-Dist: chex>=0.1.90
Requires-Dist: jax>=0.8.0
Provides-Extra: dev
Requires-Dist: build>=1.2.2; extra == "dev"
Requires-Dist: flax>=0.10.7; extra == "dev"
Requires-Dist: pytest>=8.4; extra == "dev"
Requires-Dist: rich>=14.1.0; extra == "dev"
Requires-Dist: ruff>=0.12.5; extra == "dev"
Requires-Dist: twine>=6.1.0; extra == "dev"
Requires-Dist: ty>=0.0.70; extra == "dev"
Requires-Dist: tyro>=1.0.6; extra == "dev"
Dynamic: license-file

# jaxdp

[![Python 3.11+](https://img.shields.io/badge/python-3.11%2B-blue)](https://www.python.org)
[![JAX 0.8+](https://img.shields.io/badge/JAX-0.8%2B-green)](https://github.com/jax-ml/jax)
[![version](https://img.shields.io/badge/version-0.4.0-orange)](https://github.com/TolgaOk/jaxdp)

Exact dynamic programming (DP) for finite Markov decision processes (MDPs) in JAX.

`jaxdp` provides `jax.jit` and `jax.vmap` compatible implementations of **MDP**s, DP **operators** and **mappings**, and **planning algorithms**.

## Installation

```bash
uv add jaxdp
uv add "jaxdp[dev]"   # for development
```

## Quick start

Make an MDP and apply value iteration steps.

```python
import jaxdp

mdp = jaxdp.make("frozen-lake-deterministic")
algo = jaxdp.ValueIteration(gamma=0.99)
state = algo.init(mdp)

for _ in range(100):
    state = algo.update(mdp, state)

pi = jaxdp.greedy_map.v(mdp, state.v_val, algo.gamma)
v_pi = jaxdp.policy_eval.v(mdp, pi, algo.gamma)
```

You can use `jax.vmap` to compute target values for different discount factors.

```python
import jax
import jax.numpy as jnp
import jaxdp


mdp = jaxdp.make("garnet")
v_val = jnp.linspace(0.0, 1.0, mdp.state_size)


@jax.jit
@jax.vmap
def target_value(gamma: jax.Array) -> jax.Array:
    q_val = jaxdp.reward.sa(mdp) + gamma * jaxdp.trans_op.sa(mdp, v_val)
    return jnp.max(q_val, axis=0)


gammas = jnp.array([0.9, 0.99, 0.995, 0.999])
v_vals = target_value(gammas)
# >>> v_vals.shape
# (4, ...)
```

See the [component reference](https://github.com/TolgaOk/jaxdp/blob/master/jaxdp/README.md) for the public API and MDPs.

## Citation

If you use `jaxdp` in your research, please cite:

```bibtex
@software{tolgaok_jaxdp_2026,
  author  = {Tolga Ok},
  title   = {{jaxdp}: Exact dynamic programming for finite Markov decision processes in JAX},
  year    = {2026},
  version = {0.4.0},
  url     = {https://github.com/TolgaOk/jaxdp},
}
```
