Control-Aware Representations for Model-based Reinforcement Learning

Control of non-linear dynamical systems is a key problem in control theory. Many methods have been developed with different levels of success in different classes of such problems. The majority of these methods assume that a model of the system is known and the underlying state of the system is low-dimensional and observable. These requirements limit the usage of these techniques in controlling dynamical systems from high-dimensional raw sensory data (e.g., image and audio), where the system dynamics is unknown, a scenario often seen in modern reinforcement learning (RL).

Recent years have witnessed the rapid development of a large arsenal of model-free RL algorithms, such as DQN (Mnih et al., 2013), TRPO (Schulman et al., 2015a), PPO (Schulman et al., 2017), and SAC (Haarnoja et al., 2018), with impressive success in solving high-dimensional control problems. However, most of this success has been limited to simulated environments (e.g., computer games), mainly due to the fact that these algorithms often require a large number of samples from the environment. This restricts their applicability in real-world physical systems, for which data collection is often a difficult process. On the other hand, model-based RL algorithms, such as PILCO (Deisenroth and Rasmussen, 2011), MBPO (Janner et al., 2019), and Visual Foresight (Ebert et al., 2018), despite their success, still have many issues in handling the difficulties of learning a model (dynamics) in a high-dimensional (pixel) space.

To address this issue, a class of algorithms have been developed that are based on learning a low-dimensional latent (embedding) space and a latent model (dynamics), and then using this model to control the system in the latent space. This class has been referred to as learning controllable embedding (LCE) and includes recently developed algorithms, such as E2C (Watter et al., 2015), RCE (Banijamali et al., 2018), SOLAR (Zhang et al., 2019), PCC (Levine et al., 2020), Dreamer (Hafner et al., 2020), and PC3 (Shu et al., 2020). The following two properties are extremely important in designing LCE models and algorithms. First, to learn a representation that is the most suitable for the control process at hand. This suggests incorporating the control algorithm in the process of learning representation. This view of learning control-aware representations is aligned with the value-aware and policy-aware model learning, VAML (Farahmand, 2018) and PAML (Abachi et al., 2020), frameworks that have been recently proposed in model-based RL. Second, to interleave the representation learning and control processes, and to update them both, using a unifying objective function. This allows to have an end-to-end framework for representation learning and control.

LCE methods, such as SOLAR and Dreamer, have taken steps towards the second objective by performing representation learning and control in an online fashion. This is in contrast to offline methods like E2C, RCE, PCC, and PC3, that learn a representation once and then use it in the entire control process. On the other hand, methods like PCC and PC3 address the first objective by adding a term to their representation learning loss function that accounts for the curvature of the latent dynamics. This term regularizes the representation towards smoother latent dynamics, which are suitable for the locally-linear controllers, e.g., iLQR (Li and Todorov, 2004), used by these methods.

In this paper, we take a few steps towards the above two objectives. We first formulate a LCE model to learn representations that are suitable to be used by a policy iteration (PI) style algorithm in the latent space. We call this model control-aware representation learning (CARL). We derive a loss function for CARL that exhibits a close connection to the prediction, consistency, and curvature (PCC) principle for representation learning, proposed in Levine et al. (2020). We derive three implementations of CARL: offline, online, and value-guided. Similar to offline LCE methods, such as E2C, RCE, PCC, and PC3, in offline CARL, we first learn a representation and then use it in the entire control process. However, in offline CARL, we replace the locally-linear control algorithm (e.g., iLQR) used by these LCE methods with a PI-style (actor-critic) RL algorithm. Our choice of RL algorithm is the model-based implementation of soft actor-critic (SAC) (Haarnoja et al., 2018). Our experiments show significant performance improvement by replacing iLQR with SAC. Online CARL is an iterative algorithm in which at each iteration, we first learn a latent representation by minimizing the CARL loss, and then perform several policy updates using SAC in this latent space. Our experiments with online CARL show further performance gain over its offline version. Finally, in value-guided CARL (V-CARL), we optimize a weighted version of the CARL loss function, in which the weights depend on the TD-error of the current policy. This would help to further incorporate the control algorithm in the representation learning process. We evaluate the proposed algorithms by extensive experiments on benchmark tasks and compare them with several LCE baselines.

We are interested in learning control policies for non-linear dynamical systems, where the states $s\in\mathcal{S}\subseteq\mathbb{R}^{n_{s}}$ are not fully observed and we only have access to their high-dimensional observations $x\in\mathcal{X}\subseteq\mathbb{R}^{n_{x}},\;n_{x}\gg n_{s}$. This scenario captures many practical applications in which we interact with a system only through high-dimensional sensory signals, such as image and audio. We assume that the observations $x$ have been selected such that we can model the system in the observation space using a Markov decision process (MDP)¹¹1A method to ensure observations are Markovian is to buffer them for several time steps (Mnih et al., 2013). $\mathcal{M}_{\mathcal{X}}=\langle\mathcal{X},\mathcal{A},r,P,\gamma\rangle$, where $\mathcal{X}$ and $\mathcal{A}$ are observation and action spaces; $r:\mathcal{X}\times\mathcal{A}\rightarrow\mathbb{R}$ is the reward function with maximum value $R_{\max}$, defined by the designer of the system to achieve the control objective;²²2For example, in a goal tracking problem in which the agent (robot) aims at finding the shortest path to reach the observation goal $x_{g}$ (the observation corresponding to the goal state $s_{g}$), we may define the reward for each observation $x$ as the negative of its distance to $x_{g}$, i.e., $-\|x-x_{g}\|^{2}$. $P:\mathcal{X}\times\mathcal{A}\rightarrow\mathbb{P}(\mathcal{X})$ is the unknown transition kernel; and $\gamma\in(0,1)$ is the discount factor. Our goal is to find a mapping from observations to control signals, $\mu:\mathcal{X}\rightarrow\mathbb{P}(\mathcal{A})$, with maximum expected return, i.e., $J(\mu)=\mathbb{E}[\sum_{t=0}^{\infty}\gamma^{t}r(x_{t},a_{t})\mid P,\mu]$.

Since the observations $x$ are high-dimensional and the observation dynamics $P$ is unknown, solving the control problem in the observation space may not be efficient. As discussed in Section 1, the class of learning controllable embedding (LCE) algorithms addresses this issue by learning a low-dimensional latent (embedding) space $\mathcal{Z}\subseteq\mathbb{R}^{n_{z}},\;n_{z}\ll n_{x}$ and a latent state dynamics, and controlling the system there. The main idea behind LCE is to learn an encoder $E:\mathcal{X}\rightarrow\mathbb{P}(\mathcal{Z})$, a latent space dynamics $F:\mathcal{Z}\times\mathcal{A}\rightarrow\mathbb{P}(\mathcal{Z})$, and a decoder $D:\mathcal{Z}\rightarrow\mathbb{P}(\mathcal{X})$,³³3Some recent LCE models, such as PC3 (Shu et al., 2020), are advocating latent models without a decoder. Although we are aware of the merits of such approach, we use a decoder in the models proposed in this paper. such that a good or optimal controller (policy) in $\mathcal{Z}$ minimizes the expected loss in the observation space $\mathcal{X}$. This means that if we model the control problem in $\mathcal{Z}$ as a MDP $\mathcal{M}_{\mathcal{Z}}=\langle\mathcal{Z},\mathcal{A},\bar{r},F,\gamma\rangle$ and solve it using a model-based RL algorithm to obtain a policy $\pi:\mathcal{Z}\rightarrow\mathbb{P}(\mathcal{A})$, the image of $\pi$ in the observation space, i.e., $(\pi\circ E)(a|x)=\int_{z}dE(z|x)\pi(a|z)$, should have a high return. Thus, the loss function to learn $\mathcal{Z}$ and $(E,F,D)$ from observations $\{(x_{t},a_{t},r_{t},x_{t+1})\}$ should be designed to comply with this goal.

This is why in this paper, we propose a LCE framework that tries to incorporate the control algorithm used in the latent space in the representation learning process. We call this model, control-aware representation learning (CARL). In CARL, we set the class of control (RL) algorithms used in the latent space to approximate policy iteration (PI), and more specifically to soft actor-critic (SAC) (Haarnoja et al., 2018). Before describing CARL in details in the following sections, we present a number of useful definitions and notations here.

For any policy $\mu$ in $\mathcal{X}$, we define its value function $U_{\mu}$ and Bellman operator $T_{\mu}$ as

for all $x\!\in\!\mathcal{X}$ and $U:\mathcal{X}\!\rightarrow\!\mathbb{R}$, where $r_{\mu}(x)\!=\!\int_{a}d\mu(a|x)r(x,a)$ and $P_{\mu}(x^{\prime}|x)\!=\!\int_{a}d\mu(a|x)P(x^{\prime}|x,a)$ are the reward function and dynamics induced by $\mu$.

Similarly, for any policy $\pi$ in $\mathcal{Z}$, we define its induced reward function and dynamics as $\bar{r}_{\pi}(z)=\int_{a}d\pi(a|z)\bar{r}(z,a)$ and $F_{\pi}(z^{\prime}|z)=\int_{a}d\pi(a|z)F(z^{\prime}|z,a)$. We also define its value function $V_{\pi}$ and Bellman operator $T_{\pi}$, for all $z\in\mathcal{Z}$ and $V:\mathcal{Z}\rightarrow\mathbb{R}$, as

For any policy $\pi$ and value function $V$ in the latent space $\mathcal{Z}$, we denote by $\pi\circ E$ and $V\circ E$, their image in the observation space $\mathcal{X}$, given encoder $E$, and define them as

In this section, we formulate our LCE model, which we refer to as control-aware representation learning (CARL). As described in Section 2, CARL is a model for learning a low-dimensional latent space $\mathcal{Z}$ and the latent dynamics, from data generated in the observation space $\mathcal{X}$, such that this representation is suitable to be used by a policy iteration (PI) algorithm in $\mathcal{Z}$. In order to derive the loss function used by CARL to learn $\mathcal{Z}$ and its dynamics, i.e., $(E,F,D,\bar{r})$, we first describe how the representation learning can be interleaved with PI in $\mathcal{Z}$. Algorithm 1 contains the pseudo-code of the resulting algorithm, which we refer to as latent space learning policy iteration (LSLPI).

Each iteration $i$ of LSLPI starts with a policy $\mu^{(i)}$ in the observation space $\mathcal{X}$, which is the mapping of the improved policy in $\mathcal{Z}$ in iteration $i-1$, i.e., $\pi_{+}^{(i-1)}$, back in $\mathcal{X}$ through the encoder $E^{(i-1)}$ (Lines 6 and 7). We then compute $\pi^{(i)}$, the current policy in $\mathcal{Z}$, as the image of $\mu^{(i)}$ in $\mathcal{Z}$ through the encoder $E^{(i)}$ (Line 4). Note that $E^{(i)}$ is the encoder learned at the end of iteration $i-1$ (Line 8). We then use the latent space dynamics $F^{(i)}$ learned at the end of iteration $i-1$ (Line 8), and first compute the value function of $\pi^{(i)}$ in the policy evaluation or critic step, i.e., $V^{(i)}=V_{\pi^{(i)}}$ (Line 5), and then use $V^{(i)}$ to compute the improved policy $\pi_{+}^{(i)}$, as the greedy policy w.r.t. $V^{(i)}$, i.e., $\pi^{(i+1)}=\mathcal{G}[V^{(i)}]$, in the policy improvement or actor step (Line 6). Using the samples in the buffer $\mathcal{D}$, together with the current policies in $\mathcal{Z}$, i.e., $\pi^{(i)}$ and $\pi_{+}^{(i)}$, we learn the new representation $(E^{(i+1)},F^{(i+1)},D^{(i+1)},\bar{r}^{(i+1)})$ (Line 8). Finally, we generate samples $\mathcal{D}^{(i+1)}$ by following $\mu^{(i+1)}$, the image of the improved policy $\pi_{+}^{(i)}$ back in $\mathcal{X}$ using the old encoder $E^{i}$ (Line 7), and add it to the buffer $\mathcal{D}$ (Line 9), and the algorithm iterates. It is important to note that both critic and actor operate in the low-dimensional latent space $\mathcal{Z}$.

LSLPI is a PI algorithm in $\mathcal{Z}$. However, what is desired is that it also acts as a PI algorithm in $\mathcal{X}$, i.e., it results in (monotonic) policy improvement in $\mathcal{X}$, i.e., $U_{\mu^{(i+1)}}\geq U_{\mu^{(i)}}$. Therefore, we define the representation learning loss function in CARL, such that it ensures that LSLPI also results in policy improvement in $\mathcal{X}$. The following theorem, whose proof is reported in Appendix B, shows the relationship between the value functions of two consecutive polices generated by LSLPI in $\mathcal{X}$.

It is easy to see that LSLPI guarantees (policy) improvement in $\mathcal{X}$, if the terms in the parentheses on the RHS of (4) are zero. We now describe these terms. The last term on the RHS of (4) is the KL between $\pi^{(i)}\circ E$ and $\mu^{(i)}=\pi^{(i)}\circ E^{(i)}$. This term can be seen as a regularizer to keep the new encoder $E$ close to the old one $E^{(i)}$. The four terms in (5) are: (I) The encoding-decoding error to ensure $x\approx(D\circ E)(x)$; (II) The error that measures the mismatch between the reward of taking action according to policy $\pi\circ E$ at $x\in\mathcal{X}$, and the reward of taking action according to policy $\pi$ at the image of $x$ in $\mathcal{Z}$ under $E$; (III) The error in predicting the next observation through paths in $\mathcal{X}$ and $\mathcal{Z}$. This is the error between $x^{\prime}$ and $\hat{x}^{\prime}$ shown in Fig. 1(a); and (IV) The error in predicting the next latent state through paths in $\mathcal{X}$ and $\mathcal{Z}$. This is the error between $z^{\prime}$ and $\tilde{z}^{\prime}$ shown in Fig. 1(b).

Representation Learning in CARL:   Theorem 1 provides us with a recipe (loss function) to learn the latent space $\mathcal{Z}$ and $(E,F,D,\bar{r})$. In CARL, we propose to learn a representation for which the terms in the parentheses on the RHS of (4) are small. As mentioned earlier, the second term, $L_{\text{reg}}(E,E_{-},\pi,x)$, can be considered as a regularizer to keep the new encoder $E$ close to the old one $E_{-}$. Term (II) that measures the mismatch between rewards can be kept small, or even zero, if the designer of the system selects the rewards in a compatible way. Although CARL allows us to learn a reward function in the latent space, similar to several other LCE works (Watter et al., 2015; Banijamali et al., 2018; Levine et al., 2020; Shu et al., 2020), in this paper, we assume that a compatible reward function is given. Terms (III) and (IV) are the equivalent of the prediction and consistency terms in PCC (Levine et al., 2020) for a particular latent space policy $\pi$. Since PCC has been designed for an offline setting (i.e., one-shot representation learning and control), the prediction and consistency terms are independent of a particular policy and are defined for state-action pairs. While CARL is designed for an online setting (i.e., interleaving representation learning and control), and thus, its loss function at each iteration depends on the current latent space policies $\pi$ and $\pi_{+}$. As we will see in Section 4.1, our offline implementation of CARL uses a loss function very similar to that of PCC. Note that (IV) is slightly different than the consistency term in PCC. However, if we upper-bound it using Jensen inequality as $(\mathrm{IV})\leq L_{c}(E,F,\pi,x):=\int_{x^{\prime}\in\mathcal{X}}dP_{\pi\circ E}(x^{\prime}|x)\cdot D_{\text{KL}}\big{(}E(\cdot|x^{\prime})\;||\;(F_{\pi}\circ E)(\cdot|x)\big{)}$, we obtain the loss term $L_{\text{c}}(E,F,\pi,x)$, which is very similar to the consistency term in PCC. Similar to PCC, we also add a curvature loss to the loss function of CARL to encourage having a smoother latent space dynamics $F_{\pi}$. Putting all these terms together, we obtain the CARL loss function as

where $(\lambda_{\text{ed}},\lambda_{\text{p}},\lambda_{\text{c}},\lambda_{\text{cur}},\lambda_{\text{reg}})$ are hyper-parameters of the algorithm, $(L_{\text{ed}},L_{\text{p}})$ are the encoding-decoding and prediction losses defined in (5), $L_{\text{c}}$ is the consistency loss defined above, $L_{\text{cur}}$ is the curvature loss, and $L_{\text{reg}}$ is the regularizer that ensures the new encoder remains close to the old one. We set $\lambda_{\text{reg}}$ to a small value not to allow $L_{\text{reg}}$ to play a major role in our implementations.

The CARL loss function in (6) introduces an optimization problem that takes a policy $\pi$ as input and learns a representation suitable for its evaluation and improvement. To optimize this loss in practice, similar to the PCC model (Levine et al., 2020), we define $\widehat{P}=D\circ F_{\pi}\circ E$ as a latent variable model that factorizes as $\widehat{P}(x_{t+1},z_{t},\hat{z}_{t+1}|x_{t},\pi)=\widehat{P}(z_{t}|x_{t})\widehat{P}(\hat{z}_{t+1}|z_{t},\pi)\widehat{P}(x_{t+1}|\hat{z}_{t+1})$, and use a variational approximation to the interactable negative log-likelihood of the loss terms in (6). The variational bounds for these terms can be obtained similar to Eqs. 6 and 7 in Levine et al. (2020). Below we describe three instantiations of the CARL model in practice. Implementation details can be found in Algorithm 2 in Appendix E. While CARL is compatible with most PI-style (actor-critic) RL algorithms, following a recent work, MBRL (Janner et al., 2019), we choose SAC as the RL algorithm in CARL. Since most actor-critic algorithms are based on first-order gradient updates, as discussed in Section 3, we regularize the curvature of the latent dynamics $F$ (see Eqs. 8 and 9 in Levine et al. 2020) in CARL to improve its empirical stability and performance in policy learning.

In this section, we experiment with the following continuous control domains (see Appendix F for more descriptions): (i) Planar System, (ii) Inverted Pendulum, (iii) Cartpole, (iv) 3-link manipulator, and compare the performance of our CARL algorithms with two LCE baselines: PCC (Levine et al., 2020) and SOLAR (Zhang et al., 2019). These tasks have underlying start and goal states that are “not” observable to the algorithms, instead the algorithms only have access to the start and goal image observations. To evaluate the performance of the control algorithms, similar to Levine et al. (2020), we report the $\%$-time spent in the goal. The initial policy that is used for data generation is uniformly random. To measure performance reproducibility for each experiment, we 1) train $25$ models and 2) perform $10$ control tasks for each model. For SOLAR, due to its high computation cost we only train and evaluate $10$ different models. Besides the average results, we also report the results from the best LCE models, averaged over the $10$ control tasks.

In this paper, we argued for incorporating control in the representation learning process and for the interaction between control and representation learning in learning controllable embedding (LCE) algorithms. We proposed a LCE model called control-aware representation learning (CARL) that learns representations suitable for policy iteration (PI) style control algorithms. We proposed three implementations of CARL that combine representation learning with model-based soft actor-critic (SAC), as the controller, in offline and online fashions. In the third implementation, called value-guided CARL, we further included the control process in representation learning by optimizing a weighted version of the CARL loss function, in which the weights depend on the TD-error of the current policy. We evaluated the proposed algorithms on benchmark tasks and compared them with several LCE baselines. The experiments show the importance of SAC as the controller and of the online implementation. Future directions include 1) investigating other PI-style algorithms in place of SAC, 2) developing LCE models suitable for value iteration style algorithms, and 3) identifying other forms of bias for learning an effective embedding and latent dynamics.

Controlling non-linear dynamical systems from high-dimensional observation is challenging. Direct application of model-free and model-based reinforcement learning algorithms to this problem may not be efficient, due to requiring a large number of samples from the real system (model-free) and the challenges of learning a model in a high-dimensional space (model-based). A popular approach to address this problem is learning controllable embedding (LCE), i.e., learning a low-dimensional latent space and a latent space model, and performing the optimal control in this learned latent space. This work is a step towards end-to-end representation learning and control in this setting. We propose methods that interleave representation learning and control, which allows us to learn control-aware representations, i.e., representations that are suitable for the control problem at hand.