An Adaptive State Aggregation Algorithm for Markov Decision Processes

When a pre-spcified aggregation is given, tsitsiklis1996 ; van2006 give performance bounds on the aggregated system and propose variants of value iteration. There are also a variety of ways to perform state aggregation based on different criteria. ferns2012 ; dean1997 analyze partitioning the state space by grouping states whose transition model and reward function are close. mccallum1997 proposes aggregate states that have the same optimal action and similar $Q$-values for these actions. jong2005 develops aggregation techniques such that states are aggregated if they have the same optimal action. We refer the readers to li2006 ; abel2016 ; abel2020 for a more comprehensive survey on this topic.

On the dynamical learning side, our adaptive state-aggregated value iteration algorithm is also related to the aggregation-disaggregation method used to accelerate the convergence of value iterations (bertsekas1988 ; schweitzer1985 ; mendelssohn1982 ; bertsekas2018 ) in policy evaluation, i.e., to evaluate the value function of a policy. Among those works, that of bertsekas1988 is closest to our approach. Assuming the underlying Markov processes to be ergodic, the authors propose to group states based on Bellman residuals in between runs of value iteration. They also allow states to be aggregated or disaggregated at every abstraction step. However, the value iteration step in our setting is different from the value iteration in bertsekas1988 for the reason that value iteration in policy evaluation will not involve the Bellman operator. Aside from state aggregation, a variety of other methods have been studied to accelerate value iteration. herzberg1994 proposes iterative algorithms based on a one-step look-ahead approach. shlakhter2010 combines the so called “projective operator” with value iteration and achieve better efficiency. anderson1965 ; fang2009 ; zhang2020 analyze the Anderson mixing approach to speed up the convergence of fixed-point problems.

Dynamic learning of the aggregate states has also been studied more generally in MDP and reinforcement learning (RL) settings. hostetler2014 proposes a class of $Q$-value–based state aggregations and applies them to Monte Carlo tree search. slivkins2011 uses data-driven discretization to adaptively discretize state and action space in a contextual bandit setting. ortner2013 develops an algorithm for learning state aggregation in an online setting by leveraging confidence intervals. sinclair2019 designs a $Q$-learning algorithm based on data-driven adaptive discretization of the state-action space. For more state abstraction techniques see baras2000 ; dean1997 ; jiang2014 ; abel2019 .

We consider an infinite-horizon Markov Decision Process $M=(\mathcal{S},\mathcal{A},P,r,\gamma,\rho)$, consisting of: a finite state space $\mathcal{S}$; a finite action space $\mathcal{A}$; the probability transition model $P$, where $P(s^{\prime}|s,a)$ denotes the probability of transitioning to state $s^{\prime}$ conditioned on the state-action pair $(s,a)$; the immediate cost function $r:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$, which denotes the immediate cost—or reward—obtained from a particular state-action pair; the discount factor $\gamma\in[0,1)$; and the initial distribution over states in $\mathcal{S}$, which we denote by $\rho$.

A policy $\pi:\mathcal{S}\to\mathcal{A}$ specifies the agent’s action based on the current state, either deterministically or stochastically. A policy induces a distribution over the trajectories $\tau=(s_{t},a_{t},r_{t})_{t=0}^{\infty}$, where $s_{0}\sim\rho$, and $a_{t}\sim\pi(\cdot|s_{t})$ and $s_{t+1}\sim P(\cdot|s_{t},a_{t})$ for $t\geq 0$.

A value function $\bm{V}:\mathcal{S}\to\mathbb{R}^{|\mathcal{S}|}$ assigns a value to each state; as $|\mathcal{S}|<\infty$, $\bm{V}$ can also be represented by a finite-length vector $(V(1),...,V(|\mathcal{S}|))^{\top}\in\mathbb{R}^{|\mathcal{S}|}$. (In this paper, we view all vector as vector functions mapping from the index to the corresponding entry.) Moreover, each policy $\pi$ is associated with a value function $\bm{V}^{\pi}:\mathcal{S}\to\mathbb{R}$, which is defined to be the discounted sum of future rewards starting at $s$ and with policy $\pi$:

As noted above, we represent both the value function corresponding to the policy $\pi$ and the value vector $(V^{\pi}(1),...,V^{\pi}(|\mathcal{S}|))^{\top}$ as $\bm{V}^{\pi}$.

For each state $s\in\mathcal{S}$ and action $a\in\mathcal{A}$ belonging to state $s$, let $\bm{P}_{s,a}\in\mathbb{R}^{|\mathcal{S}|}$ denote the vector of transition probabilities resulting from taking the action $a$ in state $s$. We call a policy $\pi$ greedy with respect to a given value $\bm{V}\in\mathbb{R}^{|\mathcal{S}|}$ if

We define $\bm{T}:\mathbb{R}^{|\mathcal{S}|}\to\mathbb{R}^{|\mathcal{S}|}$ to be the dynamic programming operator such that $(\bm{T}\bm{V})(s)=T_{s}(\bm{V})$ where

The optimal value function $\bm{V}^{*}$ is the unique solution of the equation $\bm{V}^{*}=T\bm{V}^{*}$. A common approach to find $\bm{V}^{*}$ is value iteration, which, given an initial guess $\bm{V}_{0}$, generates a sequence of value functions $\{\bm{V}_{t}\}_{t=0}^{\infty}$ such that $\bm{V}_{t+1}=\bm{T}\bm{V}_{t}$. The sequence $\{\bm{V}_{t}\}_{t=0}^{\infty}$ converges to $V^{*}$ as $t$ goes to $+\infty$ Bel .

The state space of MDPs can be very large. State aggregation divides the state space $\mathcal{S}$ into $K$ subsets and views each collection of states as a mega-state. Then, the value function generated by the mega-states can be used to approximate the optimal value $\bm{V}^{*}$.

To represent a state aggregation, we define the matrix $\bm{\Phi}\in\mathbb{R}^{|\mathcal{S}|\times K}$. We set $\phi_{i,j}=1$ if state $i$ is in the $j$-th mega-state, and let $\phi_{i,j}=0$ otherwise; i.e., column $j$ of $\bm{\Phi}$ indicates whether each state belongs to mega-state $j$. The state-reduction matrix $\bm{\Phi}$ also induces a partition $\{S_{i}\}_{i=1}^{K}$ on $\mathcal{S}$, i.e., $\mathcal{S}=\bigcup_{i=1}^{K}S_{i}$ and $S_{i}\cap S_{j}=\emptyset$ for $i\neq j$. Denote by $\bm{W}\in\mathbb{R}^{K}$ the cost-to-go value function for the aggregated state, and note that the current value of $\bm{W}$ induces a value function $\tilde{\bm{V}}(\bm{W})\in\mathbb{R}^{|\mathcal{S}|}$ on the original state space, where

Given a pre-specified aggregation, we seek the value function $\bm{W}$ for the aggregated states such that

where $\bm{e}=(e_{1},...,e_{K})^{\top}$ and $e_{j}=\max_{s_{1},s_{2}\in S_{j}}|V^{*}(s_{1})-V^{*}(s_{2})|$ for $j=1,...,K$. Intuitively, Eq. $\eqref{eqn_prefix_error}$ justifies our approach of aggregating states that have similar optimal cost-to-go values. We then state the algorithm that will converge to the correct cost-to-go values while satisfying Eq. $\eqref{eqn_prefix_error}$.

Algorithm 1 takes a similar form in Stochastic Approximation robbins1951 ; wasan2004 , and will converge almost surely to a unique cost-to-go value. Here $\alpha_{t}$ is the step size of the learning algorithm; by taking, e.g., $\alpha_{t}=1$, we recover the formula of value iteration. The following convergence result is proved in tsitsiklis1996 .

Proposition 1 states that if we are able to partition the state space such that the maximum difference of the optimal value function within each mega-state is small, the value function produced by Algorithm 1 can approximate the optimal value up to $\frac{\|\bm{e}\|_{\infty}}{1-\gamma}$, and the policy associated with the approximated value function will also be close to the optimal policy.

In order to generate an efficient approximation, Proposition 1 requires a pre-specified aggregation scheme such that $\max_{s_{1},s_{2}\in S_{i}}|V^{*}(s_{1})-V^{*}(s_{2})|$ is small for every $i$ to guarantee the appropriate level of convergence for Algorithm 1. Without knowing $\bm{V}^{*}$, is it still possible to control the approximation error? In this section we answer in the affirmative by introducing an adaptive state aggregation scheme that learns the correct state aggregations online as it learns the true cost-to-go values.

Given the current cost-to-go value vector $V\in\mathbb{R}^{|\mathcal{S}|}$, let $b_{1}=\min_{s\in\mathcal{S}}V(s)$, let $b_{2}=\max_{s\in\mathcal{S}}V(s)$. Group the cost-to-go values among disjoint subintervals of the form $\lceil(b_{2}-b_{1})/\varepsilon\rceil$. Next, let $\Delta=(b_{2}-b_{1})/\varepsilon$, and let $S_{j}$ to be the $j$-th mega-state, which contains all the states whose current estimated cost-to-go value falls in the interval $[b_{1}+(j-1)\varepsilon,b_{1}+j\varepsilon)$. Grouping the states in this way reduces the state size from $|\mathcal{S}|$ states to at most $\lceil(b_{2}-b_{1})/\varepsilon\rceil$ mega-states. See Algorithm 2 for further details.

Without the knowledge of $\bm{V}^{*}$ in advance, one must periodically perform value iteration on $\mathcal{S}$ to learn the correct aggregation to help with adapting the aggregation scheme. As a result, our algorithm alternates between two phases: in the global update phase the algorithm performs value iteration on $\mathcal{S}$; in the aggregated update phase, the algorithm starts to group together states with similar cost-to-go values based on the result of the last global update, and then performs aggregated updates as in Algorithm 1.

We denote by $\{\mathcal{A}_{i}\}_{i=1}^{\infty}$ the intervals of iterations in which the algorithm performs state-aggregated updates, and we denote by $\{\mathcal{B}_{i}\}_{i=1}^{\infty}$ the intervals of iterations in which the algorithm performs global update. As a consequence, $b<a$ for any $a\in\mathcal{A}_{i}$ and $b\in\mathcal{B}_{i}$; likewise, $a<b$ for any $a\in\mathcal{A}_{i}$ and $b\in\mathcal{B}_{i+1}$.

We then present our adaptive algorithm. For a pre-speficied number of iterations $n$, the time horizon $[1,n)$ is divided into intervals of the form $\mathcal{B}_{1},\mathcal{A}_{1},\mathcal{B}_{2},\mathcal{A}_{2},\ldots$. Every time the algorithm exits an interval of global updates $\mathcal{B}_{i}$, it runs Algorithm 2 based on the current cost-to-go value and the parameter $\varepsilon$, using the output of Algorithm 2 for the current state aggregation and cost-to-go values for $\mathcal{A}_{i}$. Similarly, every time the algorithm exits an interval of state-aggregated updates $\mathcal{A}_{i}$, it sets $\tilde{\bm{V}}(\bm{W})$, where $\bm{W}$ is the current cost-to-go value for the aggregated space, as the initial cost-to-go value for the subsequent interval of global iterations.

From Proposition 1, if we fix the state-aggregation parameter $\varepsilon$, even with perfect information, state aggregated value iteration will generate an approximation of the cost-to-go values with $\ell^{\infty}$ error bounded by $\varepsilon/(1-\gamma)$. This bound is sharp, as shown in tsitsiklis1996 . Such error is negligible in the early phase of the algorithm, but the error would accumulate in the later phase of the algorithm and prevent the algorithm from converging to the optimal value. As a result, it is not desirable for $\limsup|\mathcal{A}_{t}|\to\infty$.¹¹1 By setting $\varepsilon$ adaptively, one might achieve better complexity and error bound by setting $\lim\sup|\mathcal{A}_{t}|\to\infty$; however, adaptively choosing $\varepsilon$ lies beyond the scope of this work.

We state asymptotic convergence results for Algorithm 3; proofs can be found in the supplementary materials. For the remainder of the paper, by a slight abuse of notation, we denote by $\bm{V}_{t}$ the current cost-to-go value at iteration $t$. More specifically, if the current algorithm is in phase $\mathcal{B}_{i}$, $\bm{V}_{t}$ is the updated cost-to-go value for global value iteration. If the algorithm is in phase $\mathcal{A}_{i}$, $\bm{V}_{t}$ represents $\tilde{\bm{V}}(\bm{W}_{t})$.

Notice that the result of Theorem 1 is consistent with Proposition 1: due to state aggregation we suffer from the same $\frac{1}{1-\gamma}$ for the error bound. However, our algorithm maintains the same order of error bound without knowing the optimal value function $\bm{V}^{*}$.

We also identify the existence of “stable field” for our algorithm, and we prove that under some specific choice of the learning rate $\alpha_{t}$, with probability one the value function will stay within the stable field.

Such results provide guidance for choice of parameters. Indeed, the experimental results are consistent with the theory presented above.

We consider two problems in testing our algorithm. The first, which we term the “standard maze problem” consists of a $d_{1}\times d_{2}\times\cdots d_{n}$ grid of positions. Each position is connected to one or more adjacent positions. Moving from position to position incurs a constant cost, except for moving to the terminal state of the maze (the position $(1,\ldots,1)$) which incurs a constant reward.³³3 We rescale action costs in both the standard and terrain mazes to ensure that the maximum cost-to-go is exactly 100. There is a unique path from each position to the terminal state. A two-dimensional $20\times 20$ standard maze, in which the player can move, depending on their position, up, down left, or right is illustrated in Figure 1(a).

The second problem is the “terrain maze problem.” As in a standard maze, each state in the terrain maze represents a position in a $d_{1}\times\cdots\times d_{n}$ grid. As before, we imagine that the player can move from state to state only by travelling a single unit in any direction. (The player is only constrained from moving out of the grid entirely.) The player receives a reward for reaching the final square of the maze, which again we place at position $(1,\ldots,1)$. However, in contrast to the standard maze game, the player’s movements incur different costs at different positions. In particular, the maze is determined by a “height function” $H:\{1,\ldots,d_{1}\}\times\cdots\times\{1,\ldots,d_{n}\}\to\mathbb{R}$. The cost of movement is set to be the difference in heights between the player’s destination position and their current position, normalized appropriately.

In both problems, we also allow for stochasticity controlled by a parameter $p$ which gives the probability that a player moves in their intended direction. For $p=1$, the MDP is deterministic; otherwise, with probability $1-p$, the player moves in a different direction chosen uniformly at random.

In both problems, the actions available at any state correspond to a very sparse transition probability vectors, since players are constrained to move along cardinal directions at a rate of a single unit. However, in the standard maze game, the cost-to-go at any position is extremely sensitive to the costs-to-go at states that are very distant. In a $10\times 10$ standard maze, the initial tile (i.e., position (1, 1)) is often between 25 and 30 units away from the destination tile (i.e., position (10, 10)). In contrast, in the terrain maze game, the cost-to-go is much less sensitive to far away positions, because local immediate costs, dictated by the slopes one must climb or go down to move locally, are more significant.

We conclude the experiments section by showing the performance of our method on a real-world use case in continuous control. These problems often involve solving complex tasks with high-dimensional sensory input. The idea typically involves teaching an autonomous agent, usually a robot, to successfully complete some task or achieve some goal state. These problems are often very tough as they reside in the continuous state space (and many times action space) domain.

We already showcased the significant reduction in algorithm update costs during learning on grid-world problems in comparison with value iteration. Our goal for this section is to showcase real world examples of how our method may be practically applied in the field of continuous control. This is important not only to emphasize that our idea has a practical use case, but also to further showcase its ability to scale into extremely large (and continuous) dimensional problems.