Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Going faster to see further: GPU-accelerated value iteration and simulation for perishable inventory control using JAX

Joseph Farrington, Kezhi Li, Wai Keong Wong, Martin Utley

TL;DR

This work addresses the computational intractability of optimal replenishment policies for perishable inventory with complex age-structured state spaces. It introduces a JAX-based, GPU-accelerated value-iteration framework (with vmap/pmap) and gymnax-based simulators to solve large Markov decision processes, extending tractability to problems with millions of states and features like lead-time, substitution, and uncertain useful life. The authors demonstrate feasibility across three scenarios, achieving near-optimal performance with heuristic policies (optimality gaps typically below 2.5%) and achieving substantial wall-time reductions compared to CPU baselines, including up to up to 16M-state problems converging in hours on consumer GPUs and minutes on multi-GPU clusters. The work provides open-source code and Colab resources, highlighting practical impact for operational research and inventory management by enabling rigorous benchmarking and faster experimentation on accessible hardware. The broader significance lies in enabling large-scale, realistic policy evaluation and benchmark development, potentially accelerating advances in perishable logistics and related stochastic optimization problems.

Abstract

Value iteration can find the optimal replenishment policy for a perishable inventory problem, but is computationally demanding due to the large state spaces that are required to represent the age profile of stock. The parallel processing capabilities of modern GPUs can reduce the wall time required to run value iteration by updating many states simultaneously. The adoption of GPU-accelerated approaches has been limited in operational research relative to other fields like machine learning, in which new software frameworks have made GPU programming widely accessible. We used the Python library JAX to implement value iteration and simulators of the underlying Markov decision processes in a high-level API, and relied on this library's function transformations and compiler to efficiently utilize GPU hardware. Our method can extend use of value iteration to settings that were previously considered infeasible or impractical. We demonstrate this on example scenarios from three recent studies which include problems with over 16 million states and additional problem features, such as substitution between products, that increase computational complexity. We compare the performance of the optimal replenishment policies to heuristic policies, fitted using simulation optimization in JAX which allowed the parallel evaluation of multiple candidate policy parameters on thousands of simulated years. The heuristic policies gave a maximum optimality gap of 2.49%. Our general approach may be applicable to a wide range of problems in operational research that would benefit from large-scale parallel computation on consumer-grade GPU hardware.

Going faster to see further: GPU-accelerated value iteration and simulation for perishable inventory control using JAX

TL;DR

This work addresses the computational intractability of optimal replenishment policies for perishable inventory with complex age-structured state spaces. It introduces a JAX-based, GPU-accelerated value-iteration framework (with vmap/pmap) and gymnax-based simulators to solve large Markov decision processes, extending tractability to problems with millions of states and features like lead-time, substitution, and uncertain useful life. The authors demonstrate feasibility across three scenarios, achieving near-optimal performance with heuristic policies (optimality gaps typically below 2.5%) and achieving substantial wall-time reductions compared to CPU baselines, including up to up to 16M-state problems converging in hours on consumer GPUs and minutes on multi-GPU clusters. The work provides open-source code and Colab resources, highlighting practical impact for operational research and inventory management by enabling rigorous benchmarking and faster experimentation on accessible hardware. The broader significance lies in enabling large-scale, realistic policy evaluation and benchmark development, potentially accelerating advances in perishable logistics and related stochastic optimization problems.

Abstract

Value iteration can find the optimal replenishment policy for a perishable inventory problem, but is computationally demanding due to the large state spaces that are required to represent the age profile of stock. The parallel processing capabilities of modern GPUs can reduce the wall time required to run value iteration by updating many states simultaneously. The adoption of GPU-accelerated approaches has been limited in operational research relative to other fields like machine learning, in which new software frameworks have made GPU programming widely accessible. We used the Python library JAX to implement value iteration and simulators of the underlying Markov decision processes in a high-level API, and relied on this library's function transformations and compiler to efficiently utilize GPU hardware. Our method can extend use of value iteration to settings that were previously considered infeasible or impractical. We demonstrate this on example scenarios from three recent studies which include problems with over 16 million states and additional problem features, such as substitution between products, that increase computational complexity. We compare the performance of the optimal replenishment policies to heuristic policies, fitted using simulation optimization in JAX which allowed the parallel evaluation of multiple candidate policy parameters on thousands of simulated years. The heuristic policies gave a maximum optimality gap of 2.49%. Our general approach may be applicable to a wide range of problems in operational research that would benefit from large-scale parallel computation on consumer-grade GPU hardware.
Paper Structure (29 sections, 33 equations, 2 figures, 17 tables, 1 algorithm)

This paper contains 29 sections, 33 equations, 2 figures, 17 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related work
  3. Methods
  4. Scenarios
  5. Value iteration
  6. Simulation of the Markov decision processes
  7. Reproducibility
  8. Hardware
  9. Scenario A: lead time may be greater than one period
  10. Problem description
  11. Results
  12. Scenario B: substitution between two perishable products
  13. Problem description
  14. Results
  15. Scenario C: periodic demand and uncertain useful life on arrival
  16. ...and 14 more sections

Figures (2)

  • Figure 1: Wall times required to run value iteration for experiment 1 with $m=3$ for Scenario B using different GPUs. The Nvidia GeForce RTX 3060 is a consumer-grade GPU. The Nvidia A100 40GB is a data-centre grade GPU. JAX enables our method to be run on multiple identical GPUs without any code changes.
  • Figure 2: The optimality gap between the heuristic policy that achieved the highest mean return and the value iteration policy plotted against the wall time required to run value iteration for the experiments from Scenarios A, B and C.