A GPU-native solver for small-state MDPs — exact value iteration on a grid, looking for feedback

RL is ridiculously cool when the state space is huge or the dynamics unknown. But for a large class of problems where you do have a model, the state is relatively small (<8 dims or so?), and you want an exact policy across the entire support, backward induction on a grid is still a fast and viable option.

I couldn’t find a good implementation of this using a GPU backend, so I built bellgrid. It’s a PyTorch-based DP solver for mixed continuous/discrete state and actions.

The bellman update is embarrassingly parallel, and you can see in the readme the speedup you get from GPU – ~45x for a realistic lifecycle problem (80s -> 1.8s). Additionally I have a bunch of analytical and numerical reference problems to test correctness.

I’d love feedback if the API seems pleasant to write against and any example problems you think would be interesting to cover.

Repo: https://github.com/tbb300/bellgrid · MIT · pip install bellgrid