SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xuanle Zhao $^{1,2}$ , Duzhen Zhang $^{1,2}$ , Liyuan Han $^{1,2}$ , Tielin Zhang $^{1,2*}$ ; Bo Xu $^{1,2,3}$
$^{1}$ Institute of Automation, Chinese Academy of Sciences, Beijing, China
$^{2}$ School of Artificial Intelligence, University of Chinese Academy of Sciences, Beijing, China
<sup>3</sup>Center for Excellence in Brain Science and Intelligence Technology,
Chinese Academy of Sciences, Shanghai, China
{zhaoxuanle2022, zhangduzhen2019, hanliyuan2019, tielin.zhang, xubo}@ia.ac.cn
Neural ordinary differential equations (ODEs) are widely recognized as the standard for modeling physical mechanisms, which help to perform approximate inference in unknown physical or biological environments. In partially observable (PO) environments, how to infer unseen information from raw observations puzzled the agents. By using a recurrent policy with a compact context, context-based reinforcement learning provides a flexible way to extract unobservable information from historical transitions. To help the agent extract more dynamics-related information, we present a novel ODE-based recurrent model combined with model-free reinforcement learning (RL) framework to solve partially observable Markov decision processes (POMDPs). We experimentally demonstrate the efficacy of our methods across various PO continuous control and meta-RL tasks. Furthermore, our experiments illustrate that our method is robust against irregular observations, owing to the ability of ODEs to model irregularly-sampled time series.
Conventional reinforcement learning (RL) is typically cast as a problem of solving fully observable Markov decision process (MDP) tasks which are trained and tested on the same task. However, most practical applications require the agents to handle some degree of partially observable and irregular observations. Humans are always good at solving these kinds of tasks by extracting crucial information from past observations and actions, while conventional RL agents do not have the ability to extract information relevant to the tasks.
In recent works, several categories of methods have been proposed to solve PO problems. The most straightforward one is to include all historical observations as input [Lee et al., 2020]. However, this kind of method is impractical in application owing to the dimension increment of the observation inputs. Another category is based on model-free methods which generally use recurrent neural networks (RNNs) as memory units and function approximators to extract historical task-related information to solve partial observable problems [Igl et al., 2018, Kapturowski et al., 2019, Han et al., 2019, Jaderberg et al., 2019]. Recurrent neural networks are commonly employed to encode historical state transitions into contextual representations, leveraging their capacity to process sequential information. These representations are then used in conjunction with the current state to find the optimal policy. This baseline is simple and efficient, as it obtains near-optimal results and only requires changing the policy and value network into a recurrent version. The third category considers using a model to estimate a belief state from historical transitions [Lee et al., 2020]. Agent receives
belief states as observations that contain useful information to solve the tasks, this category is much similar to the second one. However, almost all notable breakthroughs have been showcased in discrete-time problems such as games and discrete control problems [Mnih et al., 2015, Silver et al., 2016], whereas most physical and biological systems in the real world are inherently continuous in time and follow differential equation dynamics.
Differential equations are commonly considered the benchmark for modeling physical mechanisms [Scholkopf et al., 2021]. They permit people to forecast the future behaviors of physical systems and the statistical interdependencies among variables. Moreover, they provide physical insights, explain the system's functioning, and justify reasons for their causal structure [Peters et al., 2022, Zhang et al., 2023b]. A discretized control system will converge to an ODE as the time increment approaches zero. However, there is no consensus on how to choose the right discretization step sizes and account for irregular observation times. In current paradigms, step sizes are fixed in advance and vary in environments. For example, in classic control tasks (such as CartPole and Pendulum), observations are sampled by regularly discretized timesteps [Brockman et al., 2016].
In this study, similar to context-based RL, we developed an ODE-based recurrent model (GRU-ODE) that encodes historical observations, actions, and rewards into a latent variable. Previous works have shown that ODE-based recurrent models could simulate the physical dynamic system in an auto-regressive fashion, such as Latent-ODE [Rubanova et al., 2019] and ODE-LSTM [Lechner and Hasani, 2020]. The GRU-ODE is based on recurrent networks, which modify the GRU topology by separating its memory cell into time-continuous states. At each time step, the model first iterates using the conventional GRU equations and then computing the differentiation based on the ODEs [Chen et al., 2018]. We then propose a method to solve PO tasks by training the GRU-ODE and the actor-critic algorithm. We show owing to the differential process, our model could extract more accurate unobservable information from partially observable environments. Our approach outperforms baselines in several regular PO and Meta-RL control tasks and irregular PO tasks.
2 Preliminary and Related work
Typical POMDPs are used to describe decision or control problems in which the underlying states of the environments cannot be directly observed. A POMDP is a tuple  $(S, \mathcal{A}, \mathcal{O}, T, O, R, \gamma)$ .  $\mathcal{S}$  is a set of states,  $\mathcal{A}$  is a set of actions, and  $T: S \times \mathcal{A}$  is the state-transition probability function. Let  $\mathcal{O}$  be the observation set and let  $O: S \times \mathcal{A} \times \mathcal{O} \to [0,1]$  be the observation probability. The reward function  $R: S \times \mathcal{A}$  decides the reward during state transition and  $\gamma$  is the discount factor. The objective is to learn a policy function to maximize expected discounted rewards  $\mathbb{E}\left[\sum_{t=0}^{T-1} \gamma^t r_{t+1}\right]$  with a finite horizon  $T$ .
Common PO tasks include scenarios of partially occluded states [Heess et al., 2015], randomly dropped frames [Hausknecht and Stone, 2015], egocentric images [Zhu et al., 2017], and randomly noised observations [Meng et al., 2021]. These kinds of problems are hard to solve because the dimension of historical observations grows linearly with timesteps. Prior works consider learning a belief state to encode underlying environments by extracting information from historical state transitions [Zintgraf et al., 2019]. Some other methods use memory-based policy, which takes the entire history states as inputs [Lee et al., 2020]. Recent works have widely used recurrent models to equip relevant algorithms [Ni et al., 2022, Han et al., 2019]. These recurrent methods could be further divided into model-free and model-based methods according to whether their training objectives contain transition functions.
In the meta-RL setting, some indicator of the task is unobserved and methods are quite similar to that of PO tasks. Prominent approaches utilize recurrent networks for fast adaption [Duan et al., 2016]. Recent studies demonstrate that the recurrent model-free algorithm serves as a robust baseline for meta-RL [Ni et al., 2022]. This is accomplished by equipping the policy network with a recurrent model, which aggregates past state and action transitions into an auxiliary context input [Zintgraf et al., 2019]. Recent works also consider training non-recurrent models to encode the context [Mu et al., 2022, Guo et al., 2022]. These approaches can be categorized into context-based meta-RL methods. Another category uses policy gradients [Finn et al., 2017, Xu et al., 2018a] or hyperparameters [Xu et al., 2018b] to learn from aggregated experience are defined as gradient-based methods.
Neural ODEs have been widely used to tackle irregular time series [Chen et al., 2018]. Standard RNN treats observations as a sequence of tokens and does not account for variable timestep between observations. When facing irregular-sized data, another interpolation or generative adversarial network is used to perform interpolation and imputation. However, these methods are not robust in unseen environments. Chen et al. [2018] use residual neural networks with time-invariant dynamics to solve the ODE initial-value problem (IVP). Recent works build continuous time models by adding ODEs to the recurrent cells update process [De Brouwer et al., 2019, Kidger et al., 2020]. CT-LSTM [Mei and Eisner, 2017] combines the LSTM architecture with the continuous-time neural Hawkes process. ODE-LSTM [Lechner and Hasani, 2020] transforms the internal dynamical flow of LSTM to a continuous-time model and demonstrates its efficacy in learning kinematics simulation. Liquid time-constant networks (LTC) [Hasani et al., 2021, Lechner et al., 2020] represent dynamic systems with liquid time constants coupled to their hidden states and compute outputs by numerical differential equation solvers.
Recent studies in scientific machine learning have attempted to utilize dynamic encoding or physical mechanisms within their respective fields, such as DeepONet [Lu et al., 2019] and PINNs [Cuomo et al., 2022, Krishnapriyan et al., 2021]. Incorporating physical insights has been shown to enhance the efficiency of neural networks and the ability of representation learning [Zhang et al., 2022, Yang et al., 2021, Han et al., 2023, Zhang et al., 2023a]. Some previous works combine ODEs with model-based RL. Du et al. [2020] use Latent-ODE to learn the transition model to solve semi-Markov decision processes (SMDPs), which also shows variables encoded by Latent-ODE capture unobservable state representations in PO environment. Yildiz et al. [2021] infer the unknown state evolution differentials with Bayesian neural ODEs. Some other works use ODEs to learn the dynamic or transition functions and apply them to continuous and irregular time interval tasks [Salehi et al., 2022, Ainsworth et al., 2021].
This section first describes constructing ODE-based recurrent models to encode historical information into context variables for model-free RL. Then, we describe the overall training objective and how to use our model to optimize the policy network and value function.
In the typical POMDP setting, the reward and transition functions share some common structures across the MDPs. Some embeddings must represent the detail of the task information, which could not be accessed in advance. Context-based RL methods attempt to solve partial observable problems by allowing the policy network to receive an auxiliary context variable which includes historical information. The context variable at time step  $t$  is inferred from the agents' experience up to the current states,
$$
\tau_ {t} = \left(o _ {0}, a _ {0}, r _ {1}, o _ {1}, a _ {1}, \dots , o _ {t - 1}, a _ {t - 1}, r _ {t}, o _ {t}\right). \tag {1}
$$
There should be a superscript  $i$  to denote various episodes, we drop it for ease of notation. In POMDP environments, observations are denoted as  $o_t$  instead of  $s_t$ , as the underlying states are only partially observable and must be inferred based on the available observations. Based on the above  $\tau_{:t}$ , it is sufficient to encode historical information into a context variable  $z_t$ , rather than modeling transition and reward dynamics which consist of millions of parameters. Generally, the memory-based policy is defined as  $\pi(a_t|\tau_{:t})$ , conditioning on the whole history.
Following ODE-RNN [Rubanova et al., 2019], we extend hidden state transitions in RNNs to continuous-time dynamics defined by Neural ODEs. As RNNs with exponentially-decayed hidden state follow an implicit ODE of the form  $\frac{dh(t)}{dt} = -\tau h$  with  $h(t_0) = h_0$ , where  $\tau$  is a parameter, the hidden states could be modeled by Neural ODEs. Therefore, we propose our ODE-based recurrent model.
GRU-ODE. In our models, transformations of hidden states and transitions between latent states are computed using recurrent cells and neural ODEs, respectively. The resulting model is
$$
\tilde {h} _ {t} = \operatorname {G R U C e l l} \left(h _ {t - 1}, x _ {t}\right),
$$
$$
h _ {t} = \text {O D E S o l v e} \left(f _ {\theta}, \tilde {h} _ {t}, d t\right), \tag {2}
$$
$$
z _ {t} = \mathcal {N} (\mu_ {h _ {t}}, \sigma_ {h _ {t}}) = o (h _ {t}),
$$
where the context variable  $z_{t}$  is defined as a Gaussian distribution parameterized by mean  $\mu_{h_t}$  and standard deviation  $\sigma_{h_t}$ .  $x_{t}$  is the concatenation of  $o_{t}$ ,  $a_{t - 1}$  and  $r_t$ . The gate of the GRUCell is controlled by the output of ODESolve  $h_t$ , instead of its own output  $\tilde{h}_t$ . The ODESolve are numerical differential solvers implemented by neural networks
$$
\text {O D E S o l v e} \left(f _ {\theta}, h (t), d t\right) = h (t) + \int_ {t} ^ {t + d t} f _ {\theta} (h (t)) d t, \tag {3}
$$
where  $f_{\theta}$  computes the differential of  $h(t)$  which are implemented by neural networks. The time increment  $dt$  in the numerical ODE solver is fixed for conventional control tasks and varied in irregular observation tasks. The input of our model  $x_{t}$  is the concatenation of observations  $o_{t}$ , actions  $a_{t-1}$  and rewards  $r_{t}$ . We found that the reward signals could help the agent learn unobservable information, leading to improved agent performance in PO environments. As the reward function typically consists of  $r(s, a) = d(s, s^{*}) + c||a||_{2}$ , where the first term measures the distance between the current state and the objective state and the latter term contains the penalty of the action magnitude. Therefore, the unobservable state information could be inferred based on tuples of observations, actions, and rewards. To verify the rationality of our proposed model, we simplify the classic control task. In PO dynamic systems, the underlying physics of the task is always supposed to be second-order, the observable states are expressed in terms of either position  $\theta$  or angular velocity  $\omega$ . As the state and action space could be denoted as  $s(t) = [\theta(t), \omega(t)]^T$  and  $a(t)$  respectively, the dynamic system with explicit control is defined as follows:
$$
\frac {d \theta (t)}{d t} = \omega (t), \quad \frac {d \omega (t)}{d t} = \mathcal {A} (\theta (t), \omega (t), a (t)), \tag {4}
$$
where  $\mathcal{A}(\cdot)$  is referred to as acceleration field. Simplifying the update equation, this procedure could be modeled as differential equations with explicit actions. The current state could be modeled using ODESolve with initial state as  $s(t) = s(0) + \int_0^t f_\theta (s(\tau),a(\tau))d\tau$ , which could be solved by neural ODEs [Chen et al., 2018]. Previous neural ODE models were always used for time-series predictions, which do not contain explicit actions. Chiappa et al. [2017] incorporate actions as auxiliary inputs of recurrent cells to build recurrent simulators. We combine them to build our recurrent models. In addition to states, actions and rewards can also be solved as initial value problems.
Proposition 1. The action  $a(t)$  is differential with time like the state  $s(t)$ , as the action is controlled by a deterministic policy function  $\pi(\cdot)$ . The derivative of  $a(t)$  with respect to time  $t$  could be expressed as
$$
\frac {d a (t)}{d t} = \frac {d \pi (s (t))}{d t} = \frac {d \pi (s (t))}{d s (t)} \frac {d s (t)}{d t}. \tag {5}
$$
The policy function  $\pi(\cdot)$  is implemented with neural networks, so  $d\pi(s(t)) / ds(t)$  is the differential between the output and input of the policy network and could be computed by the chain rule. Following the above formulation, the reward function is also differentiable as it is always the linear combination of state  $s(t)$  and action  $a(t)$ . Therefore, the update process of GRU-ODE input  $x(t)$ , the concatenation of state, action and reward, could be modeled as ODEs and solved by ODESolve. As the latent context variable  $z_t$  reflects the general dynamic across different episodes, we reckon the update process of  $z_t$  could be modeled in the same way with explicit input  $x(t)$ . Our model explicitly changes the discrete update into a continuous case, which represents latent variables more precisely. Existing implementations of state transitions perform inaccurate numerical integration routines, which results in rough estimation. For example, in the CartPole task crude numerical solver lead approximations bifurcate near the pole hangs up positions [Yildiz et al., 2021]. The accuracy of the approximation is dependent on the stepsize of the solver. Consequently, continuous methods tend to provide a closer approximation to the true value.
Figure 1: Network architecture. Trajectories of observations, actions, and rewards are sampled and stored in the buffer. During the training procedure,  $\tau_{t-1:t}$  is processed online using a GRUCell to produce the embedding  $\tilde{h}_t$ . Subsequently,  $\tilde{h}_t$  is passed through an ODE solver and MLP to encode the context variable  $z_t$ . We use separate encoders for the actor and critic networks to improve the performance.
Proposition 2. The global and local truncation errors of the Euler numerical solver are  $O(h)$  and  $O(h^2)$  respectively. As the stepsize  $h$  approaches zero, the truncation error tends to  $\lim_{h\to 0}|O(h)| = 0$ . Assuming the increment  $dt$  in the ODEs equals the observation interval, the discretized system will converge to the actual ODE as the time increment approaches zero.
Similar to relevant context-based methods, we use a recurrent ODE-based model to encode past trajectories into a variable. Compared with ODE-based models with encoder-decoder structures, models transformed from typical recurrent models could predict online at each time step. As opposed to standard RNNs, ODE-based RNNs learn the dynamics between observations, rather than having them pre-defined, which enables handling irregular data and modeling continuous sequences without relying on pre-assumptions about the dynamics of the observation sequences [Rubanova et al., 2019]. Moreover, we found that owing to their underlying physical principles, ODEs are suitable for inferring unknown physical information, which is significant in partially observable environments.
Now we have the GRU-ODE model to represent the context variable, we move on to the question of identifying learning procedure and optimal policy. As mentioned above, with partial observations, latent context  $z_{t}$  is used to provide historically unknown information. Prior works have shown using separate recurrent encoders for the actor and critic could achieve high rewards [Meng et al., 2021]. Ni et al. [2022] show the gradient norms of recurrent context encoders vary in the actor and critic resulting in the gradient of critic loss dominating the actors', which will impact the training process of the actor's context encoder. Therefore, we use separate GRU-ODE models for context representation and implement them into Actor-Critic algorithms to improve the final performance.
Model implementation. The general training objective maximizes the discounted reward in an episode. We implement our method based on TD3 Fujimoto et al. [2018] and SAC Haarnoja et al. [2018], so the policy network and value function should be parameterized differently. Deep neural networks are used to represent individual components. There are :
1. The recurrent context encoder GRU-ODE model  $g_{\phi}(z_t|\tau ;t)$ , parameterised by  $\phi$ . We use superscripts  $a$  or  $c$  to denote the encoder and context variables used for the policy network or the value function.
2. The conditioned policy network (actor)  $\pi_{\theta}(a_t|o_t,g_\phi^a (z_t^a |\tau ;t))$  parameterized by  $\theta$  with context variable  $z_{t}^{a}$
3. The conditioned value function (critic)  $q_{\psi}(Q_t|o_t, a_t, g_{\phi}^c(z_t^c|\tau:t))$  parameterized by  $\phi$  with context variable  $z_t^c$ .
The network architecture is shown in Fig. 1. The context variable  $z_{t}$  is regarded as a posterior and represented by the distribution's parameters, the details are discussed later. Compared to modeling the transition or reward function directly using the network, learning distribution embeddings need fewer parameters which makes inference easier.
Training objectives. The conventional training objective of the policy network  $\pi_{\theta}(\cdot)$  is to maximize the expected return
$$
\mathcal {J} (\theta) = \mathbb {E} _ {\pi_ {\theta}} \left[ \sum_ {t = 0} ^ {T} \gamma^ {t} R \left(r _ {t + 1} \mid s _ {t}, a _ {t}\right) \right]. \tag {6}
$$
After experimental analysis, it was discovered that utilizing  $\mathcal{J}(\theta)$  directly in conditioned policy network and value function with recurrent encoders results in significant variation of the embedding variable  $z_{t}$  within an episode. This variation can lead to a degradation of the context encoder's representation ability in PO environments. Therefore, we adjust the original training objective with an embedding constraint term to regularize  $z_{t}$ . Now, the overall training objective is to maximize
$$
\mathcal {L} (\phi , \theta , \psi) = \mathbb {E} _ {\rho} \left[ \mathcal {J} (\psi , \phi , \theta) - \lambda \mathcal {K} (\phi) \right], \tag {7}
$$
where  $\rho$  denotes the trajectory distribution induced by the sample policy. Expectations are approximated by Monte Carlo samples. Conditioned delayed policy update algorithms are implemented in our method, thus  $\mathcal{J}(\psi ,\phi ,\theta)$  are used to denote actor and critic networks with different parameters. To regularize the context variable  $z_{t}$ , we add the Kullback-Leibler (KL) divergence  $\kappa (\phi)$  between posterior and prior distributions
$$
\mathcal {K} (\phi) = \sum_ {t = 0} ^ {T} D _ {K L} \left(g _ {\phi} \left(z _ {t} \mid \tau ; t\right) \mid \mid p \left(z _ {t}\right)\right). \tag {8}
$$
We set the prior to our previous posterior,  $g_{\phi}(z_{t - 1}|\tau ;t - 1)$ , with initial prior  $g_{\phi}(z_0) = \mathcal{N}(0,I)$ .  $D_{KL}$  denotes the KL divergence, which could be optimized by the reparameterization trick. Since both terms in  $\mathcal{K}(\phi)$  are parameterized by Gaussian distributions, the KL-divergence can be analytically calculated as
$$
D _ {K L} \left(g _ {\phi} \left(z _ {t} \mid \tau_ {: t}\right) \| p \left(z _ {t}\right)\right) = \log \frac {\sigma_ {\phi , t}}{\sigma_ {\phi , t - 1}} + \frac {\left(\mu_ {\phi , t} - \mu_ {\phi , t - 1}\right) ^ {2} + \sigma_ {\phi , t} ^ {2}}{2 \sigma_ {\phi , t - 1} ^ {2}} - \frac {1}{2} \tag {9}
$$
In this work, past trajectories are encoded by GRU-ODE. Eq. 7 is optimized end-to-end, and  $\lambda$  is the balance weight that supervises the encoder learning objective against the RL objective  $\mathcal{I}(\psi ,\phi ,\theta)$  which is objective to find an optimal policy and accurate value function. To improve the efficiency of the training process, we sample fixed-length truncated sequences as mini-batches for training our model, rather than using the whole episodes.
Previous work mentioned that training the encoder and the policy network separately using data from different buffers can prevent the gradients of opposing losses from interfering with each other. Zintgraf et al. [2019] follow this setting to pre-train the encoder in advance before the policy network. We found this setting is unnecessary in practice, the data buffer is shared in our method.