Combining Reinforcement Learning and Tensor Networks, with an Application to Dynamical Large Deviations

We present a framework to integrate tensor network (TN) methods with reinforcement learning (RL) for solving dynamical optimization tasks. We consider the RL actor-critic method, a model-free approach for solving RL problems, and introduce TNs as the approximators for its policy and value functions. Our “actor-critic with tensor networks” (ACTeN) method is especially well suited to problems with large and factorizable state and action spaces. As an illustration of the applicability of ACTeN we solve the exponentially hard task of sampling rare trajectories in two paradigmatic stochastic models, the East model of glasses and the asymmetric simple exclusion process, the latter being particularly challenging to other methods due to the absence of detailed balance. With substantial potential for further integration with the vast array of existing RL methods, the approach introduced here is promising both for applications in physics and to multi-agent RL problems more generally.

Figure 1

Actor-critic with tensor networks. (a) Sketch of a Markov decision process. (b) In actor-critic RL, the state is passed to an actor, which chooses the action, and to a critic, which values the state given the reward. This value is used to improve the actor’s policy. In ACTeN, the function approximators for actor and critic are tensor networks. (c) Top: typical trajectory of the ASEP at half-filling and $L=50$ sites with one particle highlighted (blue), shown for 3000 steps. Bottom: trajectory with a current large deviation, sampled from the ACTeN solution for biasing (counting) field $\lambda =-3$. See the text for details.

Figure 2

Dynamical large deviations in the East model using ACTeN. Scaled-cumulant generating function for the dynamical activity of the East model as a function of biasing field $\lambda$ from ACTeN (symbols), for $L=50$ and PBC. Our RL results coincide with those obtained from the current state-of-the-art method using DMRG, cf. Ref. [64] (which is possible since the East model obeys detailed balance). Inset: Kinetic constraint of the East model; a spin, $s_i$, can flip, $s_i\to 1-s_i$, only if the spin to the left is up, $s_{i-1}=1$.

Figure 3

Dynamical large deviations in the ASEP using ACTeN. (a) In the ASEP particles can only move to an unoccupied neighboring site, with probability $p$ to the left and $q=1-p$ to the right. (b) SCGF for the time-integrated particle current as a function of biasing field. We show results from ACTeN for $p=0.1$ (squares) and $p=1/2$ (diamonds). The lack of detailed balance for PBC and $p\ne 1/2$ prevents straightforward application of DMRG, but for small sizes (here $L=14$) we can compare to exact diagonalization (blue curve for $p=0.1$, green for $p=1/2$). (c) SCGF for $p=0.1$ from ACTeN for size $L=50$ which is beyond the scope of ED. Compared to $L=14$ (blue curve from ED), we see that ACTeN captures the flattening of the SCGF for larger sizes indicative of a LD phase transition, cf. Ref. [32]. The inset shows the smooth convergence of our ACTeN numerics with $L$ for two values of $\lambda$. (d) Since ACTeN provides direct access to the optimal dynamics, observables such as the time-integrated current can be evaluated directly (black squares for $L=50$). We show for comparison the numerical differentiation of the ACTeN SCGF (red circles) and of the ED SCGF at $L=14$ (blue line).

Figure 4

Training procedure and learning curves. (a) For each bias [we show $\lambda =-1$ (top row), $\lambda =1$ (middle row), $\lambda =2$ (bottom row)] TN-based policies and value functions are produced via actor-critic optimization. These are initiated at random for $L=4$ with $\chi =16$ and trained for ${10}^6$ steps. Every 5000 training steps the average reward of the policy is evaluated over ${10}^4$ steps (black squares) and the weights of the policy (which we call a “snapshot” for that time) are stored. The evaluated values can be compared to the training estimate of $r(\pi )$ (red circles), which tends to overestimate $r(\pi )$ initially. The policy snapshot with the highest evaluated $r$ (blue dashed line) is used to initiate the policy for higher values of $L$. This is repeated every $\mathrm{\Delta }L=2$ up to $L=50$, with $L=14$ shown here. (b) For each bias, several policies (here six) are independently trained via the same procedure from different random initial conditions. This produces a distribution of evaluated average rewards, here represented by the median (black squares) and interquartile range (red-shaded region). The policy with the maximum average reward at each $L$ is selected as the optimal dynamics (blue triangles). (c) Same as (a) for $L=50$. The learning curves appear noisier than in (a) but note that the vertical scale is much smaller. The learning rate is kept fixed throughout. (d) The distribution of $r$ across parallel agents for $L=50$ is again much tighter than for $L=14$.