SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Nishanth Anand *
McGill University and Mila
nishanth.anand@mail.mcgill.ca
McGill University, Mila, and Deepmind
dprecup@cs.mcgill.ca
Temporal difference (TD) learning is often used to update the estimate of the value function which is used by RL agents to extract useful policies. In this paper, we focus on value function estimation in continual reinforcement learning. We propose to decompose the value function into two components which update at different timescales: a permanent value function, which holds general knowledge that persists over time, and a transient value function, which allows quick adaptation to new situations. We establish theoretical results showing that our approach is well suited for continual learning and draw connections to the complementary learning systems (CLS) theory from neuroscience. Empirically, this approach improves performance significantly on both prediction and control problems.
Deep reinforcement learning (RL) has achieved remarkable successes in complex tasks, e.g Go [40, 41] but in a narrow setting, where the task is well-defined and stable over time. In contrast, humans continually learn throughout their life and adapt to changes in the environment and in their goals. This ability to continually learn and adapt will be crucial for general AI agents who interact with people in order to help them accomplish tasks. One possible explanation for this ability in the natural world is the existence of complementary learning systems (CLS): one which slowly acquires structured knowledge and the second which rapidly learns specifics adapting to the current situation [26]. In contrast, in RL, learning is typically focused on how to optimize return, and the main estimated quantity is one value function (which then drives learning a policy). As a result, when changes occur in the environment, RL systems are faced with the stability-plasticity dilemma [9]: whether to forget past predictions in order to learn new estimates or to preserve old estimates, which may be useful again later, and compromise the accuracy on a new task.
The stability-plasticity dilemma in RL is agnostic to the type of function approximation used to estimate predictions, and it is not simply a side effect of the function approximator having small capacity. In fact, the problem exists even when the RL agent uses a table to represent the value function. To see this, consider policy evaluation on a sequence of tasks whose true value function is different due to different rewards and transition probabilities. A convergent temporal-difference (TD) algorithm [43] aggregates information from all value functions in the task distribution implicitly, thereby losing precision in the estimate of the value function for the current situation. A tracking TD learning algorithm [46] learns predictions for the current task, overwriting estimates of past tasks and re-learning from scratch for each new task, which can require a lot of data. Using a separate function approximator for each task, which is a common approach [23], requires detecting the identity of the current task to update accordingly, which can be challenging and limits the number of tasks that can be considered.
* corresponding author.
As an alternative, we propose a CLS-inspired approach to the stability-plasticity dilemma which relies on maintaining two value function estimates: a permanent one, whose goal is to accumulate "baseline" knowledge from the entire distribution of information to which the agent is exposed over time, and a transient component, whose goal is to learn very quickly using information that is relevant to the current circumstances. This idea has been successfully explored in the context of RL and tree search in the game of Go [39], but the particular approach taken uses domain knowledge to construct features and update rules. Our approach is general, simple, online and model-free, and can be used in continual RL. Our method is also orthogonal to various advancements in the continual RL field, and therefore, it can be combined with other improvements. Our contributions are:
- A conceptual framework for defining permanent and transient value functions, which can be used to estimate value functions;
- RL algorithms for prediction and control using permanent and transient value functions;
- Theoretical results which help clarify the role of these two systems in the semi-continua (multi-task) RL setting; and
- Empirical case studies of the proposed approaches in simple gridworlds, Minigrid [11] JellyBeanWorld (JBW) [31], and MinAtar environments [51].
Let  $S$  be the set of possible states and  $\mathcal{A}$  the set of actions. At each timestep  $t$ , the agent takes action  $A_{t} \in \mathcal{A}$  in state  $S_{t}$  according to its (stochastic) policy  $\pi : S \to Dist(\mathcal{A})$ . As a consequence, the agent receives reward  $R_{t+1}$  and transitions to a new state  $S_{t+1}$ . Let  $G_{t} = \sum_{k=t}^{\infty} \gamma^{k-t} R_{k+1}$  be the discounted return obtained by following  $\pi$  from step  $t$  onward, with  $0 < \gamma < 1$  the discount factor. Let  $v_{\pi}(s) = \mathbb{E}_{\pi}[G_{t}|S_{t} = s]$ . This quantity can be estimated with a function approximator parameterized by  $\mathbf{w}$ , for example using TD learning [43]:
$$
\mathbf {w} _ {t + 1} \leftarrow \mathbf {w} _ {t} + \alpha_ {t} \delta_ {t} \nabla_ {\mathbf {w}} v _ {\mathbf {w}} \left(S _ {t}\right), \tag {1}
$$
where  $\alpha_{t}$  is the learning rate at time  $t$ ,  $\delta_t = R_{t + 1} + \gamma v_{\mathbf{w}}(S_{t + 1}) - v_{\mathbf{w}}(S_t)$  is the TD error (see Sutton and Barto [45] for details).
In the control problem, the agent's goal is to find a policy,  $\pi^{*}$ , that maximizes expected returns. This can be achieved by estimating the optimal action-value function,  $q^{*} = \max_{\pi} q_{\pi}$ , where  $q_{\pi} = \mathbb{E}_{\pi}[G_t | S_t = s, A_t = a]$ , e.g. using Q-learning [50].
"Continual RL" is closely related to transfer learning [48], meta learning [37, 15], multi-task learning [10], and curriculum learning [35]. A thorough review distinguishing these flavours can be found in Khetarpal et al. [24]. We use Continual RL to refer to a setup in which the agent's environment (i.e. the reward or transition dynamics) changes over time. We use the term semi-continual RL for the case when the agent can observe task boundaries (this is similar to the multi-task setting, but does not make assumptions about the task distribution or task IDs). In continual RL, a unique optimal value function may not exist, so the agent has to continually track the value function in order to learn useful policies [46, 2].
Our work focuses on the idea of two learning systems: one learning the "gist" of the task distribution, slowly, the other which is used for tracking can quickly adapt to current circumstances. We instantiate this idea in the context of learning value functions, both for prediction and control, but the idea can be applied broadly (for e.g. policies, GVFs). The concept of decomposing the value function into permanent and transient components was first introduced in model-based RL [39], for the single task of learning to play Go, and using samples drawn from the model to train the transient value function. In contrast, we focus on general continual RL, with samples coming from the environment.
In continual RL, most of the literature is focused on supplementing neural network approaches with better tools, such as designing new optimizers to update weights [44, 19, 12, 4, 14], building new architectures [21, 32], using experience replay to prevent forgetting [34, 36, 22, 8], promoting plasticity explicitly [28, 29], or using regularization techniques from continual supervised learning [25, 20, 30, 27]. Our method is agnostic to the nature of the function approximator used, and therefore, complementary to these advancements (and could in fact be combined with any of these methods).
Our approach has parallels with fast and slow methods [13, 7] which use neural networks as the main learning system and experience replay as a secondary system. Another similar flavor is to use the task
distribution to learn a good starting point, as in distillation-based approaches [49, 18] or MaxQInit [1]. But the latter are multi-task RL methods and their use in continual RL is limited.
3 Defining Permanent and Transient Value Functions
Inspired by CLS theory, we decompose the value function estimated by the agent into two, independently parameterized components: one which slowly acquires general knowledge, the other which quickly learns local nuances but then forgets. We call these components permanent value function,  $V_{\theta}^{(P)}$  and transient value function,  $V_{\mathbf{w}}^{(T)}$ . The overall value function is computed additively:
$$
V ^ {(P T)} (s) = V _ {\theta} ^ {(P)} (s) + V _ {\mathbf {w}} ^ {(T)} (s), \tag {2}
$$
A similar decomposition can be done to the action-value function:
$$
Q ^ {(P T)} (s, a) = Q _ {\theta} ^ {(P)} (s, a) + Q _ {\mathbf {w}} ^ {(T)} (s, a). \tag {3}
$$
Other approaches to combining the two components are possible but additive decomposition is the simplest.
3.1 Permanent Value Function
Similar to the role of the neocortex in the brain, the permanent value function should capture general structure in the value functions from the tasks that the agent has seen in its lifetime. It should also provide good baseline predictions for any task that the agent faces, in order to allow it to learn good estimates with fewer updates. To achieve this in a scenario where the agent goes through a sequence of tasks and knows when the task changes, we might store all the states visited during the current task in a buffer and using them when the task finishes, to update the permanent value function as follows:
$$
\theta_ {k + 1} \leftarrow \theta_ {k} + \bar {\alpha} _ {k} \left(V ^ {(P T)} \left(S _ {k}\right) - V _ {\theta} ^ {(P)} \left(S _ {k}\right)\right) \nabla_ {\theta} V _ {\theta} ^ {(P)} \left(S _ {k}\right), \tag {4}
$$
where  $\overline{\alpha}$  is the learning rate and  $k$  is the iteration index of the buffer. The permanent estimate takes a small step towards the new estimate of the value function in the above update rule: bootstrapping at a slower timescale. We analyze this idea further in Sec. 4.
3.2 Transient Value Function
The transient value function should compute corrections on top of the estimates provided by the permanent value function, in order to approximate more closely the true value function of the current task. We implement this intuition by updating the parameters of transient value function as:
$$
\mathbf {w} _ {t + 1} \leftarrow \mathbf {w} _ {t} + \alpha_ {t} \delta_ {t} \nabla_ {\mathbf {w}} V _ {\mathbf {w}} ^ {(T)} (S _ {t}), \tag {5}
$$
where  $\alpha$  is the learning rate,  $\delta_t = R_{t+1} + \gamma V^{(PT)}(S_{t+1}) - V^{(PT)}(S_t)$  and  $\nabla_{\mathbf{w}}V_{\mathbf{w}}^{(T)}(S_t)$  is the gradient of transient value function at time  $t$ . Eq. (5) is a semi-gradient update rule, in the same sense used to describe TD-learning, as the gradient is with respect to the current estimate only and not the target. Since the permanent and the transient value functions are independently parameterized, the semi-gradient ends up being just  $\nabla_{\mathbf{w}}V_{\mathbf{w}}^{(T)}(S_t)$ . This results in learning only a part of the value function that is not captured by the permanent value function. The transient value function is reset or decayed after transferring the new information into the permanent value function (e.g. at the moment when the task changes), similar to the functioning of the hippocampus [26, 33].
3.3 Learning Algorithm
We can easily turn the updates above into algorithms; a blueprint for the prediction case is shown in Algorithm 1, and the control version is in Appendix 3. The initialization and resets are done appropriately based on the function approximation used.
In the above algorithm, we assumed that the agent can observe task boundaries, though not any additional information (like a task id). We made this assumption to facilitate the theoretical analysis presented in the next section. A fully continual version of the algorithms is presented and evaluated empirically in Sec. 6.
Note that this algorithm is a strict generalization of TD learning: instead of setting  $\mathbf{w}$  to 0 when the task changes, if we were to match the difference between the updated permanent value function and the transient value function instead, we would obtain exactly the TD learning predictions. Another way to get TD is by initializing  $\theta$  to match the initialization of the TD learning algorithm and updating  $\mathbf{w}$  only without resetting. We explore this connection formally in the next section.
In this section, we provide some theoretical results aiming to understand the update rules of permanent and transient value function separately, and then we study
the semi-continual RL setting (because the fully continual RL setting is poorly understood in general from a theoretical point of view, and even the tools and basic setup are somewhat contentious). All our results are for the prediction problem with tabular value function representation; We expect that prediction results with linear function approximation can also be obtained. We present theorems in the main paper and include proofs in appendices B. The appendix also contains analytical equations describing the evolution of the parameters based on these algorithms.
Algorithm 1 PT-TD learning (Prediction)
1: Initialize: Buffer  $\mathcal{B},\theta ,\mathbf{w}$
2: for  $t:0\to \infty$  do
3: Store  $S_{t}$  in  $\mathcal{B}$
4: Observe  $R_{t+1}, S_{t+1}$
5: Update w using Eq. (5)
6: if Task Ends then
7: Update  $\theta$  using  $\mathcal{B}$  and Eq. (4)
8: Reset transient value function, w
9: Reset  $\mathcal{B}$
10: end if
11: end for
4.1 Transient Value Function Results
We provide an understanding of what the transient value function does in the first set of results (Theorems 1-4). We carry out the analysis by fixing  $V^{(P)}$  and only updating  $V^{(T)}$  using samples from one particular task. We use  $V_{t}^{(TD)}$  and  $V_{t}^{(PT)}$  to denote the value function estimates at time  $t$  learned using TD learning and Alg. 1 respectively.
The first result establishes a relationship between the estimates learnt using TD and the transient value function update rule in Eq.(5).
Theorem 1. If  $V_0^{(TD)} = V^{(P)}$ ,  $V_0^{(T)} = 0$  and the two algorithms train on the same samples using the same learning rates, then  $\forall t, V_t^{(PT)} = V_t^{(TD)}$ .
The theorem shows that the estimates learnt by our approach match those learnt by TD-learning in a special case, proving that our algorithm is a strict generalization of TD-learning.
The next two theorems together shows that the transient value function reduces prediction error. First we show the contraction property of the transient value function Bellman operator.
Theorem 2. The expected target in the transient value function update rule is a contraction-mapping Bellman operator:  $\mathcal{T}^{(T)}V^{(T)} = r_{\pi} + \gamma \mathcal{P}_{\pi}V^{(P)} - V^{(P)} + \gamma \mathcal{P}_{\pi}V^{(T)}$ , where  $r_{\pi}$  is expected one-step reward for each state written in vector form and  $\mathcal{P}_{\pi}$  is the transition matrix for policy  $\pi$ .
We get the Bellman operator by writing the expected target in matrix form. We now characterize the fixed point of the transient value function operator in the next theorem.
Theorem 3. The unique fixed point of the transient operator is  $(I - \gamma \mathcal{P}_{\pi})^{-1}r_{\pi} - V^{(P)}$
The theorem confirms the intuition that the transient value function learns only the part of the value function that is not captured in the permanent value function, while the agent still learns the overall value function:  $V^{(PT)} = V^{(P)} + V^{(T)} = V^{(P)} + (I - \gamma \mathcal{P}_{\pi})^{-1}r_{\pi} - V^{(P)} = v_{\pi}$ .
The next theorem relates the fixed point of the transient value function to the value function of a different Markov Reward Process (MRP)
Theorem 4. Let  $\mathcal{M}$  be an MDP in which the agent is performing transient value function updates for  $\pi$ . Define a Markov Reward Process (MRP),  $\widetilde{\mathcal{M}}$ , with state space  $\widetilde{\mathcal{S}}: \{(S_t, A_t, S_{t+1}), \forall t\}$ , transition dynamics  $\widetilde{\mathcal{P}}(x'|x, a') = \mathcal{P}(s''|s', a')$ ,  $\forall x, x' \in \widetilde{\mathcal{S}}, \forall s', s'' \in \mathcal{S}, \forall a' \in \mathcal{A}$  and reward function  $\widetilde{\mathcal{R}}(x, a') = \mathcal{R}(s, a) + \gamma V^{(P)}(s') - V^{(P)}(s)$ . Then, the fixed point of Eq.(5) is the value function of  $\pi$  in  $\widetilde{\mathcal{M}}$ .
The proof follows by comparing the value function of the modified MRP with the fixed point of transient value function. The theorem identifies the structure of the fixed point of the transient value function in the form of the true value function of a different MRP. This result can be used to design an objective for the permanent value function updates; For instance, to reduce the variance of transient value function updates.
4.2 Permanent Value Function Results
We focus on permanent value function in the next two theorems. First, we characterize the fixed point of the permanent value function. Then, we connect the fixed point of the permanent value function to the jumpstart objective [48, 1], which is defined as the initial performance of the agent in the new task before collecting any data. The latter is used to measure the effectiveness of transfer in a zero-shot learning setup, and therefore, is crucial to measure generalization under non-stationarity. We make an assumption on the task distribution for the following proofs:
Assumption 1. Let  $\mathrm{MDP}\mathcal{M}_{\tau} = (\mathcal{S},\mathcal{A},\mathcal{R}_{\tau},\mathcal{P}_{\tau},\gamma)$  denote task  $\tau$ . Then,  $v_{\tau}$  is the value function of task  $\tau$  and let  $\mathbb{E}_{\tau}$  denote the expectation with respect to the task distribution. We assume there are  $N$  tasks and task  $\tau$  is i.i.d. sampled according to  $p_{\tau}$ .
Theorem 5. Following Theorem 2, under assumption 1 and Robbins-Monro step-size conditions, the sequence of updates computed by Eq.(4) contracts to a unique fixed point  $\mathbb{E}_{\tau}[v_{\tau}]$ .
We use the approach described in Example 4.3 from Bertsekas and Tsitsiklis [5] to prove the result. Details are given in Appendix B.2.1
Theorem 6. The fixed point of the permanent value function optimizes the jumpstart objective,  $J = \min_{u\in \mathbb{R}^{|S|}}\frac{1}{2}\mathbb{E}_{\tau}[\| u - v_{\tau}\| _2^2 ]$ .
4.3 Semi-Continual RL Results
We present two important results in this section which prove the usefulness of our approach in value function retention, for past tasks, and fast adaptability on new tasks. For the proofs, we assume that at time  $t$ , both  $V_{t}^{(TD)}$  and  $V_{t}^{(PT)}$  have converged to the true value function of task  $i$  and the agent gets samples from task  $\tau$  from timestep  $t + 1$  onward.
The first result proves that the TD estimates forget the past after seeing enough samples from a new task, while our method retains this knowledge through  $V^{(P)}$ .
Theorem 7. Under assumption 1, there exists  $k_0$  such that  $\forall k \geq k_0$ ,  $\mathbb{E}_{\tau}\left[\left\| V_{t + k}^{(TD)} - v_i\right\| _2^2\right] > \mathbb{E}_{\tau}\left[\left\| V^{(P)} - v_i\right\| _2^2\right]$ ,  $\forall i$ .
The proof uses the sample complexity bound from Bhandari et al. [6] (Sec. 6.3, theorem 1). Details are in Appendix B.3.1.
Next, we prove that our method has a tighter upper bound on errors compared with the TD learning algorithm in expectation.
Theorem 8. Under assumption 1, there exists  $k_0$  such that  $\forall k \leq k_0$ ,  $\mathbb{E}_{\tau} \left[ \left\| v_{\tau} - V_{t + k}^{(PT)} \right\|_2^2 \right]$  has a tighter upper bound compared with  $\mathbb{E}_{\tau} \left[ \left\| v_{\tau} - V_{t + k}^{(TD)} \right\|_2^2 \right]$ ,  $\forall i$ .
The proof uses the sample complexity bound from Bhandari et al. [6] (Sec. 6.3, theorem 1) to get the desired result as detailed in Appendix B.3.2. As  $k \to \infty$ , the estimates of both the algorithms converge to the true value function and their upper bounds collapse to 0. This result doesn't guarantee low error at each timestep, but it provides an insight on why our algorithm has low error.
To further understand the behaviour of our algorithm on a new task, we derive analytical expressions for mean squared error similar to Singh and Dayan [42] for both our algorithm and TD-learning algorithm (see Appendix B.3.3 and B.3.4). Then, we use them to calculate the mean squared error empirically on a toy problem. We observe low error for our algorithm at every timestep for various task distributions (details in Appendix C.2).
5 Experiments: Semi-Continual Reinforcement Learning
We conducted experiments in both the prediction and the control settings on a range of problems. Although the theoretical results handle the prediction setting, we performed control experiments to demonstrate the broad applicability of our approach. In this section, we assume that the agent is faced with a sequence of tasks in which the transition dynamics are the same but the reward function changes. The agent knows when the task has changed. In the sections below, we focus on the effectiveness of our approach. Additional experiments providing other insights (effect of network capacity, sensitivity to hyperparameters, and effect of task distribution) are included in the Appendix (Sec. C.2, C.3, and C.4).
Baselines: Because our approach is built on top of TD-learning and Q-learning, we primarily use the usual versions of TD learning and Q-learning as baselines, which can be viewed as emphasizing stability. We also include a version of these algorithms in which weights are reset when the task switches, which creates significant plasticity. We include more baselines for the continual RL experiments, which are described in Sec. 6. The results for all algorithms are averaged across 30 seeds and shaded regions in the figures indicate  $90\%$  confidence intervals². We use the same procedure to optimize hyperparameters for all algorithms, detailed in Appendix (Sec. C.4).
5.1 Prediction experiments
We ran five experiments on 3 types of navigation tasks, with various function approximators. The policy being evaluated was uniformly random in all experiments. We used root mean squared value error (RMSVE) as a performance measure on the online task and the expected mean squared error (MSE) on the other tasks.
Discrete grid: is a 5x5 environment depicted in Fig. 1a, which has four goal states (one in each corner). The agent starts in the center and chooses an action (up, down, left or right). Each action moves the agent one step along the corresponding direction. Rewards are 0 for all transitions except those leading into a goal state. The goal rewards change as detailed in Table C.1, so the value function also changes as shown in Fig. C.1. For the linear setting, each state is represented using 10 features: five for the row and five for the column. The permanent and transient value functions use the same features. For the deep RL experiment, instead of receiving state information, the agent receives 48x48 RGB images. The agent's location is indicated in red, and goal states are colored in green
(a) Discrete grid.
(b) Minigrid.
Figure 1: Environments (PE).
as shown in Fig. 1a. Goal rewards are multiplied by 10 to diminish the noise due to the random initialization of the value function. For the tabular and linear setting, we run 500 episodes and change the goal rewards after every 50 episodes. For the deep RL experiments, we run 1000 episodes and goal rewards are switched every 100 episodes. The discount factor is 0.9 in all cases.
Figure 2: Prediction results on various problems. Solid lines represent RMSVE on the current task and dotted lines represent MSE on other tasks. Black dotted vertical lines indicate task boundaries.