Original: C ONTROL -A WARE R EPRESENTATIONS FOR M ODELBASED R EINFORCEMENT L EARNING
Revised: C ONTROL -A WARE R EPRESENTATIONS FOR M ODELBASED R EINFORCEMENT L EARNING

Original: Anonymous authors Paper under double-blind review
Revised: Anonymous authors Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: A major challenge in modern reinforcement learning (RL) is efﬁcient control of dynamical systems from high-dimensional sensory observations.
Revised: A major challenge in modern reinforcement learning (RL) is efﬁcient control of dynamical systems from high-dimensional sensory observations.

Original: Learning controllable embedding (LCE) is a promising approach that addresses this challenge by embedding the observations into a lower-dimensional latent space, estimating the latent dynamics, and utilizing it to perform control in the latent space.
Revised: Learning controllable embedding (LCE) is a promising approach that addresses this challenge by embedding the observations into a lower-dimensional latent space, estimating the latent dynamics, and utilizing it to perform control in the latent space.

Original: Two important questions in this area are how to learn a representation that is amenable to the control problem at hand, and how to achieve an end-to-end framework for representation learning and control.
Revised: Two important questions in this area are how to learn a representation that is amenable to the control problem at hand, and how to achieve an end-to-end framework for representation learning and control.

Original: In this paper, we take a few steps towards addressing these questions.
Revised: In this paper, we take a few steps towards addressing these questions.

Original: We ﬁrst formulate a LCE model to learn representations that are suitable to be used by a policy iteration style algorithm in the latent space.
Revised: We ﬁrst formulate a LCE model to learn representations that are suitable to be used by a policy iteration style algorithm in the latent space.

Original: We call this model control-aware representation learning (CARL).
Revised: We call this model control-aware representation learning (CARL).

Original: We derive a loss function and three implementations for CARL.
Revised: We derive a loss function and three implementations for CARL.

Original: In the ofﬂine implementation, we replace the locally-linear control algorithm (e.g., iLQR) used by the existing LCE methods with a RL algorithm, namely model-based soft actor-critic, and show that it results in signiﬁcant improvement.
Revised: In the ofﬂine implementation, we replace the locally-linear control algorithm (e.g., iLQR) used by the existing LCE methods with a RL algorithm, namely model-based soft actor-critic, and show that it results in signiﬁcant improvement.

Original: In online CARL, we interleave representation learning and control, and demonstrate further gain in performance.
Revised: In online CARL, we interleave representation learning and control, and demonstrate further gain in performance.

Original: Finally, we propose value-guided CARL, a variation in which we optimize a weighted version of the CARL loss function, where the weights depend on the TD-error of the current policy.
Revised: Finally, we propose value-guided CARL, a variation in which we optimize a weighted version of the CARL loss function, where the weights depend on the TD-error of the current policy.

Original: We evaluate the proposed algorithms by extensive experiments on benchmark tasks and compare them with several LCE baselines.
Revised: We evaluate the proposed algorithms by extensive experiments on benchmark tasks and compare them with several LCE baselines.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Control of non-linear dynamical systems is a key problem in control theory.
Revised: Control of non-linear dynamical systems is a key problem in control theory.

Original: Many methods have been developed with different levels of success in different classes of such problems.
Revised: Many methods have been developed with different levels of success in different classes of such problems.

Original: The majority of these methods assume that a model of the system is known and its underlying state is low-dimensional and observable.
Revised: The majority of these methods assume that a model of the system is known and its underlying state is low-dimensional and observable.

Original: These requirements limit the usage of these techniques in controlling dynamical systems from high-dimensional raw sensory data (e.g., image), where the system dynamics is unknown, a scenario often seen in modern reinforcement learning (RL).
Revised: These requirements limit the usage of these techniques in controlling dynamical systems from high-dimensional raw sensory data (e.g., image), where the system dynamics is unknown, a scenario often seen in modern reinforcement learning (RL).

Original: Recent years have witnessed a 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.
Revised: Recent years have witnessed a rapid development of a large arsenal of model-free RL algorithms, such as DQN (Mnih et al., 2013), TRPO (Schulman et al., 2015), PPO (Schulman et al., 2017), and SAC (Haarnoja et al., 2018), with impressive success in solving high-dimensional control problems.

Original: 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.
Revised: 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.

Original: This restricts their applicability in real-world physical systems, for which data collection is often a difﬁcult process.
Revised: This restricts their applicability in real-world physical systems, for which data collection is often a difﬁcult process.

Original: On the other hand, model-based RL algorithms, such as PILCO (Deisenroth & Rasmussen, 2011), MBPO (Janner et al., 2019), and Visual Foresight (Ebert et al., 2018), despite their success, still face difﬁculties in learning a model (dynamics) in a high-dimensional (pixel) space.
Revised: On the other hand, model-based RL algorithms, such as PILCO (Deisenroth & Rasmussen, 2011), MBPO (Janner et al., 2019), and Visual Foresight (Ebert et al., 2018), despite their success, still face difﬁculties in learning a model (dynamics) in a high-dimensional (pixel) space.

Original: To address the problems faced by model-free and model-based RL algorithms in solving highdimensional control problems, a class of algorithms have been developed, whose main idea is to ﬁrst learn a low-dimensional latent (embedding) space and a latent model (dynamics), and then use this model to control the system in the latent space.
Revised: To address the problems faced by model-free and model-based RL algorithms in solving highdimensional control problems, a class of algorithms have been developed, whose main idea is to ﬁrst learn a low-dimensional latent (embedding) space and a latent model (dynamics), and then use this model to control the system in the latent space.

Original: This class has been referred to as learning controllable embedding (LCE) and includes algorithms, such as E2C (Watter et al., 2015), RCE (Banijamali et al., 2018), SOLAR (Zhang et al., 2019), PCC (Levine et al., 2020),
Revised: This class has been referred to as learning controllable embedding (LCE) and includes algorithms, such as E2C (Watter et al., 2015), RCE (Banijamali et al., 2018), SOLAR (Zhang et al., 2019), PCC (Levine et al., 2020),

Original: Dreamer (Hafner et al., 2020), and
Revised: Dreamer (Hafner et al., 2020a;b),

Original: PC3 (Shu et al.,  2020).
Revised: PC3 (Shu et al., 2020), and SLAC (Lee et al., 2020).

Original: The following two properties are extremely important in designing LCE models and algorithms.
Revised: The following two properties are extremely important in designing LCE models and algorithms.

Original: First , to learn a representation that is the most suitable for the control problem at hand.
Revised: First , to learn a representation that is the most suitable for the control problem at hand.

Original: This suggests incorporating the control algorithm in the process of learning representation.
Revised: This suggests incorporating the control algorithm in the process of learning representation.

Original: Second , to interleave the representation learning and control, and to update them both, using a unifying objective function.
Revised: Second , to interleave the representation learning and control, and to update them both, using a unifying objective function.

Original: This allows to have an end-to-end framework for representation learning and control.
Revised: This allows to have an end-to-end framework for representation learning and control.

Original: LCE methods, such as SOLAR and Dreamer,  have taken steps towards the second objective by performing representation learning and control in an online fashion.
Revised: LCE methods, such as SOLAR, Dreamer, and SLAC, have taken steps towards the second objective by performing representation learning and control in an online fashion.

Original: This is in contrast to ofﬂine methods like E2C, RCE, PCC, and PC3 that learn a representation once and then use it in the entire control process.
Revised: This is in contrast to ofﬂine methods like E2C, RCE, PCC, and PC3 that learn a representation once and then use it in the entire control process.

Original: On the other hand, methods like PCC and PC3 address the ﬁrst objective by adding a term to their representation learning loss function that accounts for the curvature of the latent dynamics.
Revised: On the other hand, methods like PCC and PC3 address the ﬁrst objective by adding a term to their representation learning loss function that accounts for the curvature of the latent dynamics.

Original: This term regularizes the representation towards smoother latent dynamics, which are suitable for the locally-linear controllers, e.g., iLQR (Li & Todorov, 2004), used by these methods.
Revised: This term regularizes the representation towards smoother latent dynamics, which are suitable for the locally-linear controllers, e.g., iLQR (Li & Todorov, 2004), used by these methods.

Original: In this paper, we take a few steps towards the above two objectives.
Revised: In this paper, we take a few steps towards the above two objectives.

Original: We ﬁrst formulate a LCE model to learn representations that are suitable to be used by a policy iteration (PI) style algorithm in the latent space.
Revised: We ﬁrst formulate a LCE model to learn representations that are suitable to be used by a policy iteration (PI) style algorithm in the latent space.

Original: We call this model control-aware representation learning (CARL) and derive a loss function for it that exhibits a close connection to the prediction, consistency, and curvature (PCC) principle for representation learning (Levine et al., 2020).
Revised: We call this model control-aware representation learning (CARL) and derive a loss function for it that exhibits a close connection to the prediction, consistency, and curvature (PCC) principle for representation learning (Levine et al., 2020).

Original: We derive three implementations of CARL: ofﬂine , online , and value-guided .
Revised: We derive three implementations of CARL: ofﬂine , online , and value-guided .

Original: Similar to ofﬂine LCE methods, such as E2C, RCE, PCC, and PC3, in ofﬂine CARL, we ﬁrst learn a representation and then use it in the entire control process.
Revised: Similar to ofﬂine LCE methods, such as E2C, RCE, PCC, and PC3, in ofﬂine CARL, we ﬁrst learn a representation and then use it in the entire control process.

Original: However, in ofﬂine CARL, we replace the locally-linear control algorithm (e.g., iLQR) used by these LCE methods with a PI-style (actor-critic) RL algorithm.
Revised: However, in ofﬂine CARL, we replace the locally-linear control algorithm (e.g., iLQR) used by these LCE methods with a PI-style (actor-critic) RL algorithm.

Original: Our choice of RL algorithm is the model-based implementation of soft actor-critic (SAC) (Haarnoja et al., 2018).
Revised: Our choice of RL algorithm is the model-based implementation of soft actor-critic (SAC) (Haarnoja et al., 2018).

Original: Our experiments show signiﬁcant performance improvement by replacing iLQR with SAC.
Revised: Our experiments show signiﬁcant performance improvement by replacing iLQR with SAC.

Original: Online CARL is an iterative algorithm in which at each iteration, we ﬁrst learn a latent representation by minimizing the CARL loss, and then perform several policy updates using SAC in this latent space.
Revised: Online CARL is an iterative algorithm in which at each iteration, we ﬁrst learn a latent representation by minimizing the CARL loss, and then perform several policy updates using SAC in this latent space.

Original: Our experiments with online CARL show further performance gain over its ofﬂine version.
Revised: Our experiments with online CARL show further performance gain over its ofﬂine version.

Original: Finally, in value-guided CARL (V-CARL), we optimize a weighted version of the CARL loss function, in which the weights depend on the
Revised: Finally, in value-guided CARL (V-CARL), we optimize a weighted version of the CARL loss function, in which the weights depend on the

Original: TD-error of the current policy.
Revised: TD-error of the current policy.

Original: This would help to further incorporate the control algorithm in the representation learning process.
Revised: This would help to further incorporate the control algorithm in the representation learning process.

Original: We evaluate the proposed algorithms by extensive experiments on benchmark tasks and compare them with several LCE baselines: PCC, SOLAR, and Dreamer.
Revised: We evaluate the proposed algorithms by extensive experiments on benchmark tasks and compare them with several LCE baselines: PCC, SOLAR, and Dreamer.

Original:  P ROBLEM F ORMULATION
Revised:  P ROBLEM F ORMULATION

Original: We are interested in learning control policies for non-linear dynamical systems, where the states s 2 S ✓ R n s are not fully observed and we only have access to their high-dimensional observations x 2 X ✓ R n x , n x  n s .
Revised: We are interested in learning control policies for non-linear dynamical systems, where the states s ∈ S ⊆ R n s are not fully observed and we only have access to their high-dimensional observations x ∈ X ⊆ R n x , n x (cid:29) n s .

Original: This scenario captures many practical applications in which we interact with a system only through high-dimensional sensory signals, such as image and audio.
Revised: This scenario captures many practical applications in which we interact with a system only through high-dimensional sensory signals, such as image and audio.

Original: 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) 1 M X 
Revised: 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) 1 M X =

Original: = h X , A , r, P,  i , where X and A are observation and action spaces; r : X ⇥ A !
Revised: (cid:104)X , A , r, P, γ (cid:105) , where X and A are observation and action spaces; r : X ×

Original:  R is the reward function with maximum value R max , deﬁned by the designer of the system to achieve the control objective; 2 P : X ⇥ A ! P ( X ) is the unknown transition kernel; and  2 (0 , 1) is the discount factor.
Revised: A → R is the reward function with maximum value R max , deﬁned by the designer of the system to achieve the control objective; 2 P : X ×A → P ( X ) is the unknown transition kernel; and γ ∈ (0 , 1) is the discount factor.

Original: Our goal is to ﬁnd a mapping from observations to control signals, µ : X ! P ( A ) , with maximum expected return, i.e., J ( µ ) =
Revised: Our goal is to ﬁnd a mapping from observations to control signals, µ : X → P ( A ) , with maximum expected return, i.e., J ( µ ) =

Original: E [ P 1 t [-=
 -] t r ( x t , a t ) | P, µ ] .
Revised: E [ (cid:80) ∞ t [-=
 -]=0 γ t r ( x t , a t ) | P, µ ] .

Original: 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 efﬁcient.
Revised: 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 efﬁcient.

Original: As discussed in Section 1, the class of learning controllable embedding (LCE) algorithms addresses this by learning a low-dimensional latent (embedding) space Z ✓ R n z , n z ⌧ n x , together with a latent dynamics, and controlling the system there.
Revised: As discussed in Section 1, the class of learning controllable embedding (LCE) algorithms addresses this by learning a low-dimensional latent (embedding) space Z ⊆ R n z , n z (cid:28) n x , together with a latent dynamics, and controlling the system there.

Original: The main idea behind LCE is to learn an encoder E : X ! P ( Z ) , a latent space dynamics F : 
Revised: The main idea behind LCE is to learn an encoder E : X → P ( Z ) , a latent space dynamics F : Z ×

Original: Z ⇥ A ! P ( Z ) , and a decoder D : Z ! P ( X ) , 3 such that a good or optimal controller (policy) in Z performs well in the observation space X .
Revised: A → P ( Z ) , and a decoder D : Z → P ( X ) , 3 such that a good or optimal controller (policy) in Z performs well in the observation space X .

Original: This means that if we model the control problem in Z as a MDP M Z = h Z , A , ¯ r, F,  i and solve it using a model-based RL algorithm to obtain a policy ⇡ : Z ! P ( A ) , the image of ⇡ back in the observation space, i.e., ( ⇡  E )( a | x ) =
Revised: This means that if we model the control problem in Z as a MDP M Z = (cid:104)Z , A , ¯ r, F, γ (cid:105) and solve it using a model-based RL algorithm to obtain a policy π : Z → P ( A ) , the image of π back in the observation space, i.e., ( π ◦ E )( a | x ) =

Original: R z dE ( z | x ) ⇡ ( a | z ) , should have high expected return.
Revised: (cid:82) z dE ( z | x ) π ( a | z ) , should have high expected return.

Original: Thus, the loss function to learn
Revised: Thus, the loss function to learn

Original: Z and ( E, F, D ) from observations { ( x t , a t , r t , x t +1 ) } should be designed to comply with this goal.
Revised: Z and ( E, F, D ) from observations { ( x t , a t , r t , x t +1 ) } should be designed to comply with this goal.

Original: 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.
Revised: 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.

Original: We call this model, control-aware rep-
Revised: We call this model, control-aware rep-

Original: 1: Inputs: E (0) , F (0) , D (0) ; 2: Initialization: µ (0) = random policy; D samples generated from µ (0) ; 3: for i = 0 , 1 , . .
Revised: 1: Inputs

Original:  do 4:
Revised: D do 4:

Original: Compute ⇡ ( i ) such that µ ( i ) =
Revised: Compute π latent ◦ ; # µ ≈ π

Original: ⇡ ( i )  E( i ) ; # set ⇡ ( i ) to be the image of µ ( i )in the latent space 5:
Revised: i ◦ E i ) 5:

Original: Compute the value function of ⇡ ( i ) and set V (i ) =
Revised: Compute the value function of π i ) and set V ( ) =

Original: V ⇡ ( i ) ; # policy evaluation (critic) 6:
Revised: V π ( i ) ; # policy evaluation (critic) 6:

Original: Compute the greedy policy w.r.t. V( i ) and set ⇡ ( i )+ =
Revised: Compute the greedy policy w r.t. V i ) and set π ( i )+ =

Original: G [ V ( i ) ] ; # policy improvement (actor) 7:
Revised: G [ V ( i ) ] ; # policy improvement (actor) 7:

Original: Set µ ( i +1) =
Revised: Set µ ( i +1) =

Original: ⇡ ( i )+  E ( i ) ; # project the improved policy ⇡ ( i )+ back into the observation space 8:
Revised: π ( i )+ ◦ E ( i ) ; # project the improved policy π i )+ back into the observation space 8:

Original: Learn ( E ( i+1) , F ( i+1) , D ( i +1) , ¯ r
Revised: Learn ( E ( i , F ( i , D ( i +

Original:  ( i+1) )
Revised: , ¯ r ( i

Original: from D , ⇡ ( i) , and ⇡ ( i )+ ; # representation learning 9:
Revised: from D π ( i and π ( i )+ ; # representation learning 9:

Original: Generate samples D ( i +1) =
Revised: Generate samples D ( i +1) =

Original: { ( x t , a t , r t , x t +1 ) } nt =1 from µ ( i +1) ;
Revised: { ( x t , a t , r t , x t +1 ) } nt =1 from µ ( i +1) ;

Original: D  D
Revised: D ← D ∪ D

Original: [ D ( i +1) ; 10: end for resentation learning (CARL).
Revised: ( i +1) ; 10: end for resentation learning (CARL).

Original: In CARL, we set the class of control (RL) algorithms used in the latent space to approximate policy iteration (PI), and more speciﬁcally to soft actor-critic (SAC) (Haarnoja et al., 2018).
Revised: In CARL, we set the class of control (RL) algorithms used in the latent space to approximate policy iteration (PI), and more speciﬁcally to soft actor-critic (SAC) (Haarnoja et al., 2018).

Original: Before describing CARL in details in the following sections, we present a number of useful deﬁnitions and notations here.
Revised: Before describing CARL in details in the following sections, we present a number of useful deﬁnitions and notations here.

Original: For any policy µ in X , we deﬁne its value function U µ and Bellman operator T µ as
Revised: For any policy µ in X , we deﬁne its value function U µ and Bellman operator T µ as

Original: [Equation]
Revised: [Equation]

Revised: r µ (

Original: [Equation]
Revised: [Equation]

Revised: P

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: for all x 2 X and U : X ! R , where r µ ( x )
Revised: for all x ∈X and U : X → R , where r µ ( x )=

Original: =
Revised: (cid:82) a dµ ( a | x ) r ( x, a ) and P µ ( x (cid:48)

Original: R a dµ ( a | x )r ( x, a ) and Pµ ( x 0 | x )=
Revised: | x )= (cid:82) a dµ ( a | x ) P ( x (cid:48)

Original: R a dµ ( a | x ) P ( x 0 | x, a ) are the reward function and dynamics induced by µ .
Revised: | x, a ) are the reward function and dynamics induced by µ .

Original: Similarly, for any policy ⇡ in Z , we de- ﬁne its induced reward function and dynamics as ¯ r ⇡ ( z ) =
Revised: Similarly, for any policy π in Z , we deﬁne its induced reward function and dynamics as ¯ r π ( z ) =

Original: R a d ⇡ ( a | z )¯ r ( z, a ) and F ⇡ ( z 0 | z ) = R a d ⇡ ( a | z ) F ( z 0 | z, a ) .
Revised: (cid:82) a dπ ( a | z )¯ r ( z, a ) and F π ( z (cid:48) | z ) = (cid:82) a dπ ( a | z ) F ( z (cid:48) | z, a ) .

Original: We also deﬁne its value function V ⇡ and Bellman operator T ⇡ as 
Revised: We also deﬁne its value function V π and Bellman operator T π as V π ( z

Original: [Equation]
Revised: [Equation]

Revised: =

Original:  For any policy ⇡ and value function V in the latent space Z , we denote by ⇡  E and V  E , their image in the observation space X , given encoder E , and deﬁne them as
Revised: E [ [Equation] For any policy π and value function V in the latent space Z , we denote by π ◦ E and V ◦ E , their image in the observation space X , given encoder E , and deﬁne them as

Original: [Equation]
Revised: [Equation]

Original:  CARL M ODEL : A C ONTROL P ERSPECTIVE
Revised:  CARL M ODEL : A C ONTROL P ERSPECTIVE

Original: In this section, we formulate our LCE model, which we refer to as control-aware representation learning (CARL).
Revised: In this section, we formulate our LCE model, which we refer to as control-aware representation learning (CARL).

Original: As described in Section 2, CARL is a model for learning a low-dimensional latent space Z and the latent dynamics, from data generated in the observation space X , such that this representation is suitable to be used by a policy iteration (PI) style algorithm in Z .
Revised: As described in Section 2, CARL is a model for learning a low-dimensional latent space Z and the latent dynamics, from data generated in the observation space X , such that this representation is suitable to be used by a policy iteration (PI) style algorithm in Z .

Original: In order to derive the loss function used by CARL to learn Z and its dynamics, i.e., ( E, F, D, ¯ r ) , we ﬁrst describe how the representation learning can be interleaved with PI in Z .
Revised: In order to derive the loss function used by CARL to learn Z and its dynamics, i.e., ( E, F, D, ¯ r ) , we ﬁrst describe how the representation learning can be interleaved with PI in Z .

Original: Algorithm 1 contains the pseudo-code of the resulting algorithm, which we refer to as latent space learning policy iteration (LSLPI).
Revised: Algorithm 1 contains the pseudo-code of the resulting algorithm, which we refer to as latent space learning policy iteration (LSLPI).

Original: Each iteration i of LSLPI starts with a policy µ ( i ) in the observation space X , which is the mapping of the improved policy in Z in iteration i  1 , i.e., ⇡ ( i  1) + , back in X through the encoder E
Revised: Each iteration i of LSLPI starts with a policy µ ( i ) in the observation space X , which is the mapping of the improved policy in Z in iteration i − 1 , i.e., π ( i − 1) + , back in X through the encoder E

Original: ( i  1) (Lines 6 and 7).
Revised: ( i − 1) (Lines 6 and 7).

Original: We then compute ⇡ ( i ) , the current policy in Z , as the image of µ ( i ) in Z through the encoder E ( i ) (Line 4).
Revised: We then compute π ( i ) , the current policy in Z , as the image of µ ( i ) in Z through the encoder E ( i ) (Line 4).

Original: Note that E ( i ) is the encoder learned at the end of iteration i  1 (Line 8).
Revised: Note that E ( i ) is the encoder learned at the end of iteration i − 1 (Line 8).

Original: We then use the latent space dynamics F ( i ) learned at the end of iteration i  1 (Line 8), and ﬁrst compute the value function of ⇡ ( i ) in the policy evaluation or critic step, i.e., V ( i )
Revised: We then use the latent space dynamics F ( i ) learned at the end of iteration i − 1 (Line 8), and ﬁrst compute the value function of π ( i ) in the policy evaluation or critic step, i.e., V ( i )

Original: =
Revised: =

Original: V ⇡ ( i ) (Line 5), and then use V ( i ) to compute the improved policy ⇡ ( i )+ , as the greedy policy w.r.t.
Revised: V π ( i ) (Line 5), and then use V ( i ) to compute the improved policy π ( i )+ , as the greedy policy w.r.t.

Original: V ( i ) , i.e., ⇡ ( i +1) =
Revised: V ( i ) , i.e., π ( i +1) =

Original: G [ V ( i ) ] , in the policy improvement or actor step (Line 6).
Revised: G [ V ( i ) ] , in the policy improvement or actor step (Line 6).

Original: Using the samples in the buffer D , together with the current policies in Z , i.e., ⇡ ( i ) and ⇡ ( i )+ , we learn the new representation ( E ( i +1) , F ( i +1) , D ( i +1) , ¯ r
Revised: Using the samples in the buffer D , together with the current policies in Z , i.e., π ( i ) and π ( i )+ , we learn the new representation ( E ( i +1) , F ( i +1) , D ( i +1) , ¯ r

Original: ( i +1) )
Revised: ( i +1) )

Original: (Line 8).
Revised: (Line 8).

Original: Finally, we generate samples D ( i +1) by following µ ( i +1) , the image of the improved policy ⇡ ( i )+ back in X using the old encoder E ( i ) (Line 7), and add it to the buffer D (Line 9), and the algorithm iterates.
Revised: Finally, we generate samples D ( i +1) by following µ ( i +1) , the image of the improved policy π ( i )+ back in X using the old encoder E ( i ) (Line 7), and add it to the buffer D (Line 9), and the algorithm iterates.

Original: It is important to note that both critic and actor operate in the low-dimensional latent space Z .
Revised: It is important to note that both critic and actor operate in the low-dimensional latent space Z .

Original: LSLPI is a PI algorithm in Z .
Revised: LSLPI is a PI algorithm in Z .

Original: However, what is desired is that it also acts as a PI algorithm in X , i.e., it results in (monotonic) policy improvement in X , i.e., U µ ( i +1)  U µ ( i ) .
Revised: However, what is desired is that it also acts as a PI algorithm in X , i.e., it results in (monotonic) policy improvement in X , i.e., U µ ( i +1) ≥ U µ ( i ) .

Original: Therefore, we deﬁne the representation learning loss function for CARL, such that it ensures LSLPI also results in policy improvement in X .
Revised: Therefore, we deﬁne the representation learning loss function for CARL, such that it ensures LSLPI also results in policy improvement in X .

Original: The following theorem, whose proof is reported in Appendix A, shows the relationship between the value functions of two consecutive polices generated by LSLPI in X .
Revised: The following theorem, whose proof is reported in Appendix A, shows the relationship between the value functions of two consecutive polices generated by LSLPI in X .

Original: Theorem 1.
Revised: Theorem 1.

Original: Let µ , µ + , ⇡ , ⇡ + , and ( E, F, D, ¯ r ) be the policies µ ( i ) , µ ( i +1) , ⇡ ( i ) , ⇡ ( i )+ , and the learned latent representation ( E ( i +1) , F ( i +1) , D ( i +1) , ¯ r
Revised: Let µ , µ + , π , π + , and ( E, F, D, ¯ r ) be the policies µ ( i ) , µ ( i +1) , π ( i ) , π ( i )+ , and the learned latent representation ( E ( i +1) , F ( i +1) , D ( i +1) , ¯ r

Original: ( i +1) )
Revised: ( i +1) )

Original: at iteration i of the LSLPI algorithm (Algorithm 1).
Revised: at iteration i of the LSLPI algorithm (Algorithm 1).

Original: Then, the following holds for the value functions of µ and µ + :
Revised: Then, the following holds for the value functions of µ and µ + :

Original: [Equation]
Revised: [Equation]

Original: for all x 2 X , where  the error term  for a policy ⇡ is given by 
Revised: for all x ∈ X , where d [Equation] x (cid:96) = x (cid:48) | x 0 ; π ◦ E ) is the γ -stationary distribution induced by policy π ◦ E , and the error term ∆ for a policy π is given by ∆( E,

Original: [Equation]
Revised: [Equation]

Revised: , D,

Original: [Equation]
Revised: [Equation]

Original:  | {z } (
Revised: r, [Equation] , x ) =

Original: [Equation]
Revised: [Equation]

Original: ⌘ . It is easy to see that LSLPI guarantees (policy) improvement in X , if the terms in the parentheses on the RHS of (4) are zero.
Revised: It is easy to see that LSLPI guarantees (policy) improvement in X , if the terms in the parentheses on the RHS of (4) are zero.

Original: We now describe these terms.
Revised: We now describe these terms.

Original: The last term on the RHS of (4) is the
Revised: The last term on the RHS of (4) is the

Original: KL between ⇡ ( i )  E and µ ( i ) =
Revised: KL between π ( i ) ◦ E and µ ( i ) =

Original: ⇡ ( i )  E ( i ) .
Revised: π ( i ) ◦ E ( i ) .

Original: This term can be seen as a regularizer to keep the new encoder E close to the old one E ( i ) .
Revised: This term can be seen as a regularizer to keep the new encoder E close to the old one E ( i ) .

Original: The four terms in (5) are: (I)
Revised: The four terms in (5) are: (I)

Original: The encoding-decoding error to ensure x ⇡ ( D  E )( x ) ; (II) The error that measures the mismatch between the reward of taking action according to policy ⇡  E at x 2 X , and the reward of taking action according to policy ⇡ at the image of x in Z under E ; (III)
Revised: The encoding-decoding error to ensure x ≈ ( D ◦ E )( x ) ; (II) The error that measures the mismatch between the reward of taking action according to policy π ◦ E at x ∈ X , and the reward of taking action according to policy π at the image of x in Z under E ; (III)

Original: The error in predicting the next observation through paths in X and Z .
Revised: The error in predicting the next observation through paths in X and Z .

Original: This is the error between x 0 and ˆ x 0 shown in Fig. 1(a); and (IV)
Revised: This is the error between x (cid:48) and ˆ x (cid:48) shown in Fig. 1(a); and (IV)

Original: The error in predicting the next latent state through paths in X and Z .
Revised: The error in predicting the next latent state through paths in X and Z .

Original: This is the error between z 0 and ˜ z 0 shown in Fig. 1(b).
Revised: This is the error between z (cid:48) and ˜ z (cid:48) shown in Fig. 1(b).

Original: Representation Learning in CARL Theorem 1 provides us with a recipe (loss function) to learn the latent space Z and ( E, F, D, ¯ r ) .
Revised: Representation Learning in CARL Theorem 1 provides us with a recipe (loss function) to learn the latent space Z and ( E, F, D, ¯ r ) .

Original: In CARL, we propose to learn a representation for which the terms in the parentheses on the RHS of (4) are small.
Revised: In CARL, we propose to learn a representation for which the terms in the parentheses on the RHS of (4) are small.

Original: As mentioned earlier, the second term,
Revised: As mentioned earlier, the second term,

Original: L reg ( E, µ, ⇡ , x ) , can be considered as a regularizer to keep the new encoder E close to the old one E  , when the policy µ is given by ⇡  E  .
Revised: L reg ( E, µ, π, x ) , can be considered as a regularizer to keep the new encoder E close to the old one E − , when the policy µ is given by π ◦ E − .

Original: Term (I) minimizes the reconstruction error between encoder and decoder, which is standard for training auto-encoders (Kingma & Welling, 2013).
Revised: Term (I) minimizes the reconstruction error between encoder and decoder, which is standard for training auto-encoders (Kingma & Welling, 2013).

Original: 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.
Revised: 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 4 .

Original: 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.
Revised: 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 latent reward function is given.

Original: 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 ⇡ .
Revised: 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 π .

Original: Since PCC has been designed for an ofﬂine setting (i.e., one-shot representation learning and control), its prediction and consistency terms are independent of a particular policy and are deﬁned for state-action pairs.
Revised: Since PCC has been designed for an ofﬂine setting (i.e., one-shot representation learning and control), its prediction and consistency terms are independent of a particular policy and are deﬁned for state-action pairs.

Original: 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 ⇡ and ⇡ + .
Revised: 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 π and π + .

Original: As we will see in Section 4, in our ofﬂine implementation of CARL, these two terms are similar to prediction and consistency terms in
Revised: As we will see in Section 4, in our ofﬂine implementation of CARL, these two terms are similar to prediction and consistency terms in

Original: PCC.
Revised: PCC.

Original: Note that (IV) is slightly different than the consistency term in PCC.
Revised: Note that (IV) is slightly different than the consistency term in PCC.

Original: However, if we upper-bound it using Jensen inequality: (IV)  L c ( E, F, ⇡ , x ) :=
Revised: However, if we upper-bound it using Jensen inequality: (IV) ≤

Original: [Equation]
Revised: [Equation]

Original: the resulted loss, L c ( E, F, ⇡ , x ) , would be similar to the consistency term in PCC.
Revised: the resulted loss, L c ( E, F, π, x ) , would be similar to the consistency term in PCC.

Original: Similar to PCC, we also add a curvature loss to the loss function of CARL to encourage having a smoother latent space dynamics F ⇡ .
Revised: Similar to PCC, we also add a curvature loss to the loss function of CARL to encourage having a smoother latent space dynamics F π .

Original: Putting all these terms together, we obtain the following loss function for CARL:
Revised: Putting all these terms together, we obtain the following loss function for CARL:

Original: [Equation]
Revised: [Equation]

Original: [Equation] where (  ed ,  p ,  c ,  cur ,  reg ) are hyper-parameters  of the algorithm,
Revised: where ( λ ed , λ p , λ c , λ cur , λ reg ) are hyper-parameters 5 of the algorithm,

Original: ( L ed , L p ) are the encodingdecoding and prediction losses deﬁned in (5), L c is the consistency loss deﬁned above, L cur is the curvature loss, and L reg is the regularizer that ensures the new encoder remains close to the old one.
Revised: ( L ed , L p ) are the encodingdecoding and prediction losses deﬁned in (5), L c is the consistency loss deﬁned above, L cur = E x,u [ E (cid:15) (cid:2) f Z ( z + (cid:15) z , u + (cid:15) u ) − f Z

Original:   D IFFERENT I MPLEMENTATIONS OF CARL
Revised: ( z, u ) − ( ∇ z f Z ( z, u ) · (cid:15) z + ∇ u f Z ( z, u ) · (cid:15) u ) (cid:107) 22 (cid:3) | E ] is the curvature loss that regulates the 2 nd derivative of f Z , the mean of latent dynamics F , in which (cid:15) z , (cid:15) u are standard Gaussian noise, and L reg is the regularizer that ensures the new encoder remains close to the old one.  D IFFERENT I MPLEMENTATIONS OF CARL

Original: The CARL loss function in (6) introduces an optimization problem that takes a policy ⇡ in Z as input and learns a representation suitable for its evaluation and improvement.
Revised: The CARL loss function in (6) introduces an optimization problem that takes a policy π in Z as input and learns a representation suitable for its evaluation and improvement.

Original: To optimize this loss in practice, similar to the PCC model (Levine et al., 2020), we deﬁne 
Revised: To optimize this loss in practice, similar to the PCC model (Levine et al., 2020), we deﬁne (cid:98) P

Original: [Equation]
Revised: =

Original:  as a latent variable model that is factorized as 
Revised: D ◦ F π ◦ E as a latent variable model that is factorized as (cid:98) P ( x t +1 , z t , ˆ z t +1 | x t , π ) = (cid:98) P

Original: [Equation] 1 | z t , ⇡ ) b P ( x t +1 | ˆ z t +1 ) , and use a variational approximation to the interactable negative log-likelihood of the loss terms in (6).
Revised: z t | x t ) (cid:98) P (ˆ z t +1 | z t , π ) (cid:98) P ( x t +1 | ˆ z t +1 ) , and use a variational approximation to the interactable negative log-likelihood of the loss terms in (6).

Original: The variational bounds for these terms can be obtained similar to Eqs.
Revised: The variational bounds for these terms can be obtained similar to Eqs.

Original: 6 and 7 in Levine et al.
Revised: 6 and 7 in Levine et al.

Original: Below we describe three instantiations of the CARL model in practice.
Revised: Below we describe three instantiations of the CARL model in practice.

Original: Implementation details can be found in Algorithm 2 in Appendix D.
Revised: Implementation details can be found in Algorithm 2 in Appendix D.

Original: Although CARL is compatible with most PI-style (actor-critic)
Revised: Although CARL is compatible with most PI-style (actor-critic)

Original: RL algorithms, we choose soft actor-critic (SAC) (Haarnoja et al., 2018) as its control algorithm.
Revised: RL algorithms, we choose soft actor-critic (SAC) (Haarnoja et al., 2018) as its control algorithm.

Original: Since most actor-critic algorithms are based on ﬁrst-order gradient updates, as discussed in Section 3, we regularize the curvature of the latent dynamics F (see Eqs.
Revised: Since most actor-critic algorithms are based on ﬁrst-order gradient updates, as discussed in Section 3, we regularize the curvature of the latent dynamics F (see Eqs.

Original: 8 and 9 in Levine et al. 2020) in CARL to improve its empirical stability and performance in policy learning.
Revised: 8 and 9 in Levine et al. 2020) in CARL to improve its empirical stability and performance in policy learning.

Original: Ofﬂine CARL We ﬁrst implement CARL in an ofﬂine setting, where we generate a (relatively) large batch of observation samples { ( x t , a t , r t , x t +1 ) }
Revised: Ofﬂine CARL We ﬁrst implement CARL in an ofﬂine setting, where we generate a (relatively) large batch of observation samples { ( x t , a t , r t , x t +1 ) }

Original: Nt =1 using an exploratory (e.g., random) policy.
Revised: Nt =1 using an exploratory (e.g., random) policy.

Original: We then use this batch to optimize the CARL’s loss function (6) via the variational approximation scheme described above, and learn a latent representation Z and ( E, F, D ) .
Revised: We then use this batch to optimize the CARL’s loss function (6) via the variational approximation scheme described above, and learn a latent representation Z and ( E, F, D ) .

Original: Finally, we solve the decision problem in Z using a model-based RL algorithm, which in our case is model-based SAC.
Revised: Finally, we solve the decision problem in Z using a model-based RL algorithm, which in our case is model-based SAC 6 .

Original: The learned policy ˆ ⇡ ⇤ in Z is then used to control the system from observations as a t ⇠ (ˆ ⇡ ⇤  E )( ·| x t ) .
Revised: The learned policy ˆ π ∗ in Z is then used to control the system from observations as a t ∼ (ˆ π ∗ ◦ E )( ·| x t ) .

Original: This is the setting that has been used in several recent LCE works, such as E2C (Watter et al., 2015), RCE (Banijamali et al., 2018), PCC (Levine et al., 2020), and PC3 (Shu et al., 2020).
Revised: This is the setting that has been used in several recent LCE works, such as E2C (Watter et al., 2015), RCE (Banijamali et al., 2018), PCC (Levine et al., 2020), and PC3 (Shu et al., 2020).

Original: Our ofﬂine implementation is different than those in which 1) we replace their locally-linear control algorithm, namely iterative LQR (iLQR) (Li & Todorov, 2004), with model-based SAC, which results in signiﬁcant performance improvement, as shown in Section 5, and 2) we optimize the CARL loss function, that despite close connection, is still different than the one used by PCC.
Revised: Our ofﬂine implementation is different than those in which 1) we replace their locally-linear control algorithm, namely iterative LQR (iLQR) (Li & Todorov, 2004), with model-based SAC, which results in signiﬁcant performance improvement, as shown in Section 5, and 2) we optimize the CARL loss function, that despite close connection, is still different than the one used by PCC.

Original: The CARL loss function presented in Section 3 has been designed for an online setting in which at each iteration, it takes a policy as input and learns a representation that is suitable for evaluating and improving this policy.
Revised: The CARL loss function presented in Section 3 has been designed for an online setting in which at each iteration, it takes a policy as input and learns a representation that is suitable for evaluating and improving this policy.

Original: However, in the ofﬂine setting, the learned representation should be good for any policy generated in the course of running the PI-style control algorithm.
Revised: However, in the ofﬂine setting, the learned representation should be good for any policy generated in the course of running the PI-style control algorithm.

Original: Therefore, we marginalize out the policy from the (online) CARL’s loss function and use the RHS of the following corollary (proof in Appendix B) to construct the CARL loss function used in our ofﬂine experiments.
Revised: Therefore, we marginalize out the policy from the (online) CARL’s loss function and use the RHS of the following corollary (proof in Appendix B) to construct the CARL loss function used in our ofﬂine experiments

Original: Corollary 2.
Revised: Corollary 2.

Original: Let µ and µ + be two consecutive policies in X generated by a PI-style control algorithm in the latent space constructed by ( E,F,D, ¯ r ) .
Revised: Let µ and µ + be two consecutive policies in X generated by a PI-style control algorithm in the latent space constructed by ( E,F,D, ¯ r ) .

Original: Then, the following holds for the value functions of µ and µ + , where  is deﬁned by (5) (in modulo replacing sampled action a ⇠ ⇡  E with action a ):
Revised: Then, the following holds for the value functions of µ and µ + , where ∆ is deﬁned by (5) (in modulo replacing sampled action a ∼ π ◦ E with action a ):

Original: [Equation]
Revised: [Equation]

Original: Online CARL
Revised: Online CARL

Original: In the online implementation of CARL, at each iteration i , the current policy ⇡ ( i ) is the improved policy of the last iteration, ⇡ ( i  1) + .
Revised: In the online implementation of CARL, at each iteration i , the current policy π ( i ) is the improved policy of the last iteration, π ( i − 1) + .

Original: We ﬁrst generate a relatively (to ofﬂine CARL) small batch of samples using the image of the current policy in X , i.e., µ ( i ) = 
Revised: We ﬁrst generate a relatively (to ofﬂine CARL) small batch of samples using the image of the current policy in X , i.e., µ ( i ) = π ( i ) ◦

Original: ⇡ ( i )  E ( i  1) , and then learn a representation ( E ( i ) , F ( i ) , D ( i ) ) suitable for evaluating and improving the image of µ ( i ) in Z under the new encoder E ( i ) .
Revised: E ( i − 1) , and then learn a representation ( E ( i ) , F ( i ) , D ( i ) ) suitable for evaluating and improving the image of µ ( i ) in Z under the new encoder E ( i ) .

Original: This means that with the new representation, the current policy that was the image of µ ( i ) in Z under E ( i  1) , should be replaced by its image ⇡ ( i ) under the new encoder, i.e., ⇡ ( i )  E ( i ) ⇡ µ ( i ) .
Revised: This means that with the new representation, the current policy that was the image of µ ( i ) in Z under E ( i − 1) , should be replaced by its image π ( i ) under the new encoder, i.e., π ( i ) ◦ E ( i ) ≈ µ ( i ) .

Original: In online CARL, we address this by the following policy distillation step in which we minimize the following loss:
Revised: In online CARL, we address this by the following policy distillation step in which we minimize the following loss:

Original: [Equation]
Revised: [Equation]

Original: After the current policy ⇡ ( i ) is set, we perform multiple steps of (model-based) SAC in Z using the current model, ( F ( i ) , ¯ r ( i ) ) , and then send the resulting policy ⇡ ( i )+ to the next iteration.
Revised: After the current policy π ( i ) is set, we perform multiple steps of (model-based) SAC in Z using the current model, ( F ( i ) , ¯ r ( i ) ) , and then send the resulting policy π ( i )+ to the next iteration.

Original:  Value-Guided CARL (V-CARL)
Revised: 3. Value-Guided CARL (V-CARL)

Original: While Theorem 1 shows that minimizing the loss in (6) guarantees performance improvement, this loss does not contain any information about the performance of the current policy µ , and thus, the LCE model trained with this loss may have low accuracy in regions of the latent space that are crucial for learning good RL policies.
Revised: While Theorem 1 shows that minimizing the loss in (6) guarantees performance improvement, this loss does not contain any information about the performance of the current policy µ , and thus, the LCE model trained with this loss may have low accuracy in regions of the latent space that are crucial for learning good RL policies.

Original: In V-CARL, we tackle this issue by modifying the loss function in a way that the resulted LCE model has more accuracy in regions with higher anticipated future returns.
Revised: In V-CARL, we tackle this issue by modifying the loss function in a way that the resulted LCE model has more accuracy in regions with higher anticipated future returns.

Original: To derive the V-CARL’s loss function, we use the variational model-based policy optimization (VMBPO) framework by Chow et al.
Revised: To derive the V-CARL’s loss function, we use the variational model-based policy optimization (VMBPO) framework by Chow et al.

Original: (2020) in which the optimal dynamics for model-based
Revised: (2020) in which the optimal dynamics for model-based

Original: RL can be expressed in closed-form as P ⇤ ( x 0 | x, a )
Revised: RL can be expressed in closed-form as P ∗ ( x (cid:48) | x, a )

Original: =
Revised: =

Original: P ( x 0 | x, a ) · exp  ⌧ ( r ( x, a ) +  ˜ U µ ( x 0 )  ˜ W µ ( x, a ))  , where ˜ U µ ( x ) := 1 ⌧ log E ⇥ exp  ⌧
Revised: P ( x (cid:48) | x, a ) · exp (cid:0) τγ ( r ( x, a ) + γ ˜ U µ ( x (cid:48) ) − ˜ W µ ( x, a )) (cid:1) , where ˜ U µ ( x ) :

Original: P 1  t [-=
 -] t r µ,t  | P µ , x 0 = x ⇤ and ˜ W µ ( x, a ) := r ( x, a ) + ⌧ log E x 0 ⇠ P ( [Equation]
Revised: = 1 τ log E (cid:2) exp (cid:0) τ (cid:80) ∞ t [-=
 -]=0 γ t r µ,t (cid:1) | P µ , x 0 = x (cid:3) and ˜ W µ ( x, a ) :

Original: | x,a ) [exp( ⌧ U µ ( x 0 ))] are the optimistic value and action-value functions 5 of policy µ , and ⌧ > 0 is a temperature parameter.
Revised: = r ( x, a ) + γτ log E x (cid:48) ∼ P ( ·| x,a ) [exp( τU µ ( x (cid:48) ))] are the optimistic value and action-value functions 8 of policy µ , and τ > 0 is a temperature parameter.

Original: Note that in the VMBPO framework, the optimal dynamics
Revised: Note that in the VMBPO framework, the optimal dynamics

Original: P ⇤ is value-aware, because it re-weighs P with an exponential-twisting weight exp( ⌧ w ( x, a, x 0 )) , where w ( x, a, x 0 ) :
Revised: P ∗ is value-aware, because it re-weighs P with an exponential-twisting weight exp( τγ w ( x, a, x (cid:48) )) , where w ( x, a, x (cid:48) ) :

Original: = r ( x, a ) +  ˜ U µ ( x 0 )  ˜ W µ ( x, a ) is the temporal difference (TD) error.
Revised: = r ( x, a ) + γ ˜ U µ ( x (cid:48) ) − ˜ W µ ( x, a ) is the temporal difference (TD) error.

Original: In V-CARL, we use the VMBPO framework to modify the CARL’s prediction loss L p ( E, F, D, ⇡ , x ) .
Revised: In V-CARL, we use the VMBPO framework to modify the CARL’s prediction loss L p ( E, F, D, π, x ) .

Original: Since the regularizer loss L reg ( E, µ, ⇡ , x ) in CARL forces policies ⇡  E and µ to be close to each other, we may replace the transition dynamics P ⇡  E with P µ in L p .
Revised: Since the regularizer loss L reg ( E, µ, π, x ) in CARL forces policies π ◦ E and µ to be close to each other, we may replace the transition dynamics P π ◦ E with P µ in L p .

Original: This makes minimizing L p equivalent to maximizing the log-likelihood 
Revised: This makes minimizing L p equivalent to maximizing the log-likelihood (cid:82) x (cid:48)

Original: [Equation]
Revised: dP µ ( x (cid:48) | x ) · log( D ◦ F π ◦ E )( x (cid:48)

Original: | [Equation]
Revised: | x ) .

Original: Finally, we replace P µ with P ⇤ µ in this log-likelihood and obtain 
Revised: Finally, we replace P µ with P ∗ µ in this log-likelihood and obtain (cid:82) a dµ ( a | x ) (cid:82) x (cid:48) dP

Original: [Equation]
Revised: ( x (cid:48) | x, a ) · exp( τγ · w ( x, a, x (cid:48) ))

Original:  , which is a weighted (by the exponential TD w ( x, a, x 0 ) ) log-likelihood function (w.r.t. P ).
Revised: · log( D ◦ F π ◦ E )( x (cid:48) | x ) , which is a weighted (by the exponential TD w ( x, a, x (cid:48) ) ) log-likelihood function (w.r.t. P ).

Original: Note that this weight depends on the optimistic value functions ˜
Revised: Note that this weight depends on the optimistic value functions ˜

Original: U µ and ˜ W µ .
Revised: U µ and ˜ W µ .

Original: When ⌧ > 0 is small (see Appendix C for more details), these value functions can be approximated by their standard counterparts, i.e., ˜ U µ ( x ) ⇡ U µ ( x ) and ˜ W µ ( x, a ) ⇡ W µ ( x, a ):= r ( x, a )+ 
Revised: When τ > 0 is small (see Appendix C for more details), these value functions can be approximated by their standard counterparts, i.e., ˜ U µ ( x ) ≈ U µ ( x ) and ˜ W µ ( x, a ) ≈ W µ ( x, a ):= r ( x, a )+ (cid:82) x (cid:48)

Original: R x 0 dP ( x 0 | x, a ) U µ ( x 0 ) , which can be further approximated by their latent-space counterparts, i.e., U µ ( x ) 
Revised: dP ( x (cid:48) | x, a ) U µ ( x (cid:48) ) , which can be further approximated by their latent-space counterparts, i.e., U µ ( x ) ≈

Original: [Equation] ) and W µ ( x, a ) ⇡
Revised: ( V π ◦ E )( x ) and W µ ( x, a ) ≈

Original: ( Q ⇡  E )( x, a ) , according to Lemma 5 in Appendix A.1.
Revised: ( Q π ◦ E )( x, a ) , according to Lemma 5 in Appendix A.1.

Original: Since the latent reward function ¯ r is deﬁned such that r ( x, a ) ⇡ (¯ r  E )( x, a ) , we may write the TD-error w ( x, a, x 0 ) in terms of the encoder E and the latent value functions as 
Revised: Since the latent reward function ¯ r is deﬁned such that r ( x, a ) ≈ (¯ r ◦ E )( x, a ) , we may write the TD-error w ( x, a, x (cid:48) ) in terms of the encoder E and the latent value functions as (cid:98) w ( x, a, x (cid:48) ) :

Original: [Equation]
Revised: = (cid:82) z,z (cid:48) dE ( z | x ) · dE ( z (cid:48) | x (cid:48) ) · (¯ r ( z, a ) − Q π ( z, a ) + γV π ( z (cid:48) )) .

Original:  E XPERIMENTAL R ESULTS
Revised:  E XPERIMENTAL R ESULTS

Original: In this section, we experiment with the following continuous control domains: (i) Planar System, (ii)
Revised: In this section, we experiment with the following continuous control domains: (i) Planar System, (ii)

Original: Inverted Pendulum (Swingup), (iii) Cartpole, (iv) Three-link Manipulator (3-Pole), and compare the performance of our CARL algorithms with three LCE baselines: PCC (Levine et al., 2020),
Revised: Inverted Pendulum (Swingup), (iii) Cartpole, (iv) Three-link Manipulator (3-Pole), and compare the performance of our CARL algorithms with three LCE baselines: PCC (Levine et al., 2020),

Original: SOLAR (Zhang et al., 2019),  and two implementations of Dreamer (Hafner et al., 2020) (described below).
Revised: SOLAR (Zhang et al., 2019), SLAC (Lee et al., 2020), and two implementations of Dreamer (Hafner et al., 2020a) (described below).

Original: These tasks have underlying start and goal states that are “not” observable, instead, the algorithms only have access to the start and goal observations.
Revised: These tasks have underlying start and goal states that are “not” observable, instead, the algorithms only have access to the start and goal observations.

Original: We report the detailed setup of the experiments in Appendix E, in particular, the description of the domains in Appendix E.1 and the implementation of the algorithms in Appendix E.3.
Revised: We report the detailed setup of the experiments in Appendix E, in particular, the description of the domains in Appendix E.1 and the implementation of the algorithms in Appendix E.3.

Original: To evaluate the performance of the algorithms, similar to Levine et al.
Revised: To evaluate the performance of the algorithms, similar to Levine et al.

Original: , we report the % - time spent in the goal.
Revised: , we report the % - time spent in the goal.

Original: The initial policy that is used for data generation is uniformly random (see
Revised: The initial policy that is used for data generation is uniformly random (see

Original: Appendix E.2 for more details).
Revised: Appendix E.2 for more details).

Original: To measure performance reproducibility for each experiment, we (i) train 25 models, and (ii) perform 10 control tasks for each model.
Revised: To measure performance reproducibility for each experiment, we (i) train 25 models, and (ii) perform 10 control tasks for each model.

Original: For SOLAR, due to its high computation cost, we only train and evaluate 10 different models.
Revised: For SOLAR, due to its high computation cost, we only train and evaluate 10 different models.

Original: Besides the average results, we also report the results from the best LCE models, averaged over the 10 control tasks.
Revised: Besides the average results, we also report the results from the best LCE models, averaged over the 10 control tasks.

Original: General Results Table 1 shows the means and standard errors of % -time spent in goal, averaged over all models and control tasks, and averaged over all control tasks for the best model.
Revised: General Results Table 1 shows the means and standard errors of % -time spent in goal, averaged over all models and control tasks, and averaged over all control tasks for the best model.

Original: To compare data efﬁciency, we also report the number of samples required to train the latent space and controller in each algorithm.
Revised: To compare data efﬁciency, we also report the number of samples required to train the latent space and controller in each algorithm.

Original: We also show the training curves (performance vs. number of samples) of the algorithms in Fig.
Revised: We also show the training curves (performance vs. number of samples) of the algorithms in Fig.

Original: We report more experiments and ablation studies in Appendix F.
Revised: We report more experiments and ablation studies in Appendix F.

Original: Below summarizes our main observations of the experiments.
Revised: Below summarizes our main observations of the experiments.

Original: First, ofﬂine CARL that uses modelbased SAC as its control algorithm achieves signiﬁcantly better performance than PCC that uses iLQR in all tasks.
Revised: First, ofﬂine CARL that uses modelbased SAC as its control algorithm achieves signiﬁcantly better performance than PCC that uses iLQR in all tasks.

Original: This can be attributed to the advantage that SAC is more robust and effective in non- (locally)-linear environments.
Revised: This can be attributed to the advantage that SAC is more robust and effective in non- (locally)-linear environments.

Original: We report more detailed comparison between PCC and ofﬂine CARL in Appendix F.3, where we explicitly compare their control performance and latent representation maps.
Revised: We report more detailed comparison between PCC and ofﬂine CARL in Appendix F.3, where we explicitly compare their control performance and latent representation maps.

Original: Second, in all tasks, online CARL is more data-efﬁcient than its ofﬂine counterpart, i.e., it achieves similar or better performance with fewer samples.
Revised: Second, in all tasks, online CARL is more data-efﬁcient than its ofﬂine counterpart, i.e., it achieves similar or better performance with fewer samples.

Original: In particular, online CARL is notably superior in Planar, Cartpole, and Swingup, in which it achieves similar performance to ofﬂine CARL with 2 , 2 . 5 , and 4 times less samples, respectively (see Fig. 2).
Revised: In particular, online CARL is notably superior in Planar, Cartpole, and Swingup, in which it achieves similar performance to ofﬂine CARL with 2 , 2 . 5 , and 4 times less samples, respectively (see Fig. 2).

Original: In Appendix F.3, we show how the latent representation of online CARL progressively improves through the iterations of the algorithm (in particular, see Fig. 11).
Revised: In Appendix F.3, we show how the latent representation of online CARL progressively improves through the iterations of the algorithm (in particular, see Fig. 11).

Original: Third, in the simpler tasks (Planar, Swingup, Cartpole), V-CARL performs even better than online CARL.
Revised: Third, in the simpler tasks (Planar, Swingup, Cartpole), V-CARL performs even better than online CARL.

Original: This corroborates our hypothesis that CARL can achieve extra improvement when its LCE model is more accurate in the regions of the latent space with higher temporal difference (regions with higher anticipated future return).
Revised: This corroborates our hypothesis that CARL can achieve extra improvement when its LCE model is more accurate in the regions of the latent space with higher temporal difference (regions with higher anticipated future return).

Original: In 3-pole, the performance of
Revised: In 3-pole, the performance of

Original: V-CARL is worse than online CARL.
Revised: V-CARL is worse than online CARL.

Original: This is likely due to the instability in representation learning resulted from sample variance ampliﬁcation by the exponential-TD weight.
Revised: This is likely due to the instability in representation learning resulted from sample variance ampliﬁcation by the exponential-TD weight.

Original: Fourth, SOLAR requires signiﬁcantly more samples to learn a reasonable latent space for control, and with limited data it fails to converge to a good policy.
Revised: Fourth, SOLAR requires signiﬁcantly more samples to learn a reasonable latent space for control, and with limited data it fails to converge to a good policy.

Original: Even with the ﬁne-tuned latent space from Zhang et al.
Revised: Even with the ﬁne-tuned latent space from Zhang et al.

Original: (2019), its performance is incomparable to those of CARL variants and Dreamer.
Revised: (2019), its performance is incomparable to those of CARL variants and Dreamer.

Original: We report more experiments with SOLAR in Appendix F.5, in which we show that SOLAR can perform better, especially in Planar when we ﬁx the start and goal locations.
Revised: We report more experiments with SOLAR in Appendix F.5, in which we show that SOLAR can perform better, especially in Planar when we ﬁx the start and goal locations.

Original: However, the improved performance is still incomparable with those of CARL and Dreamer.
Revised: However, the improved performance is still incomparable with those of CARL and Dreamer.

Original: Fifth, we include an ablation study in Appendix F.2 to demonstrate how each term of the CARL’s loss function impacts policy learning.
Revised: Fifth, we include an ablation study in Appendix F.2 to demonstrate how each term of the CARL’s loss function impacts policy learning.

Original: It shows the importance of the prediction and consistency terms, without which the resulting algorithms struggle, and the (relatively) minor role of the curvature and encoder-decoder terms in the performance of the algorithms.
Revised: It shows the importance of the prediction and consistency terms, without which the resulting algorithms struggle, and the (relatively) minor role of the curvature and encoder-decoder terms in the performance of the algorithms.

Original: Dreamer As described in Section 2, most LCE algorithms, including E2C, PCC, and CARL variants, assume the observation space X is selected such that the system is Markovian there.
Revised: Dreamer As described in Section 2, most LCE algorithms, including E2C, PCC, and CARL variants, assume the observation space X is selected such that the system is Markovian there.

Original: In contrast, Dreamer does not make this assumption and has been designed for more general class of control problems that can be modeled as POMDPs.
Revised: In contrast, Dreamer does not make this assumption and has been designed for more general class of control problems that can be modeled as POMDPs.

Original: Thus, it is expected that it performs inferior (requires more samples to achieve the same performance) to CARL when the system is Markov in the observation space.
Revised: Thus, it is expected that it performs inferior (requires more samples to achieve the same performance) to CARL when the system is Markov in the observation space.

Original: Moreover, CARL and other LCE methods deﬁne the reward as the negative distance to the goal in the latent space.
Revised: Moreover, CARL and other LCE methods deﬁne the reward as the negative distance to the goal in the latent space.

Original: This cannot be done in Dreamer, where the encoder is an
Revised: This cannot be done in Dreamer, where the encoder is an

Original: RNN that takes an entire observation trajectory as input.
Revised: RNN that takes an entire observation trajectory as input.

Original: To address this, we propose two methods to train the Dreamer’s reward function in the latent space, which we refer to as Dreamer Pixel and
Revised: To address this, we propose two methods to train the Dreamer’s reward function in the latent space, which we refer to as Dreamer Pixel and

Original: Dreamer Oracle.
Revised: Dreamer Oracle.

Original: While Dreamer Pixel uses the negative distance to the goal in the observation space X as the signal to train the reward function, Dreamer Oracle uses the negative distance in the (unobserved) underlying state space S .
Revised: While Dreamer Pixel uses the negative distance to the goal in the observation space X as the signal to train the reward function, Dreamer Oracle uses the negative distance in the (unobserved) underlying state space S .

Original: Thus, it is more fair to compare the CARL algorithms with
Revised: Thus, it is more fair to compare the CARL algorithms with

Original: Dreamer Pixel than Dreamer Oracle that has the advantage of having access to the underlying state space (see Appendix F.6 for more details).
Revised: Dreamer Pixel than Dreamer Oracle that has the advantage of having access to the underlying state space (see Appendix F.6 for more details).

Original: As it was expected, our results show that although both
Revised: As it was expected, our results show that although both

Original: Dreamer’s implementations learn reasonably-performing policies for most tasks (except Planar), they require twice to 100 -times more samples to achieve the same performance as the CARL algorithms.
Revised: Dreamer’s implementations learn reasonably-performing policies for most tasks (except Planar), they require twice to 100 -times more samples to achieve the same performance as the CARL algorithms.

Original: We report longer (more samples) experiments with Dreamer on all tasks in Appendix F.6 (Fig. 12). 
Revised: We report longer (more samples) experiments with Dreamer on all tasks in Appendix F.6 (Fig. 12). [Figure] [Figure] [Table]

Original: [Figure]
Revised: [Figure]

Revised: [Table]

Original: Results with Environment-biased Sampling
Revised: [Figure] [Table] [Figure] [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Original: In the previous experiments, all the online LCE algorithms are warm-started with data collected by a uniformly random policy over the entire
Revised: [Figure] [Table] [Figure] [Table]

Revised: [Figure]

Original: [Table]
Revised: [Table]

Original:  environment.
Revised: [Figure] [Table] Results with Environment-biased Sampling In the previous experiments, all the online LCE algorithms are warm-started with data collected by a uniformly random policy over the entire environment.

Original: With sufﬁcient data the latent dynamics is accurate enough on most parts of the state space for control, therefore we do not observe a signiﬁcant difference between online CARL and VCARL.
Revised: With sufﬁcient data the latent dynamics is accurate enough on most parts of the state space for control, therefore we do not observe a signiﬁcant difference between online CARL and VCARL.

Original: To further illustrate the advantage of V-CARL over online CARL, we modify the experimental setting by gathering initial samples only from a speciﬁc region of the environment (see Appendix E. for more details).
Revised: To further illustrate the advantage of V-CARL over online CARL [Table] we modify the experimental setting by gathering initial samples only from a speciﬁc region of the environment (see Appendix E. for more details).

Original: Fig.
Revised: Fig.

Original: 3 shows the learning curves of online CARL and V-CARL in this case.
Revised: 3 shows the learning curves of online CARL and V-CARL in this case.

Original: As expected, with biased data, both algorithms experience a certain level of performance degradation, yet, V-CARL clearly outperforms online CARL — this veriﬁes our conjecture that control-aware
Revised: As expected, with biased data, both algorithms experience a certain level of performance degradation, yet, V-CARL clearly outperforms online CARL — this veriﬁes our conjecture that control-aware

Original: LCE models are more robust to initial data distribution and superior in policy optimization.
Revised: LCE models are more robust to initial data distribution and superior in policy optimization

Revised: [Table]

Revised: [Figure]

Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original:  C ONCLUSIONS
Revised:  C ONCLUSIONS

Original: 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.
Revised: 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.

Original: We proposed a LCE model called control-aware representation learning (CARL) that learns representations suitable for policy iteration (PI) style control algorithms.
Revised: We proposed a LCE model called control-aware representation learning (CARL) that learns representations suitable for policy iteration (PI) style control algorithms.

Original: We proposed three implementations of CARL that combine representation learning with model-based soft actor-critic (SAC), as the controller, in ofﬂine and online fashions.
Revised: We proposed three implementations of CARL that combine representation learning with model-based soft actor-critic (SAC), as the controller, in ofﬂine and online fashions.

Original: In the third implementation, called valueguided 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.
Revised: In the third implementation, called valueguided 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.

Original: We evaluated the proposed algorithms on benchmark tasks and compared them with several LCE baselines.
Revised: We evaluated the proposed algorithms on benchmark tasks and compared them with several LCE baselines.

Original: The experiments show the importance of SAC as the controller and of the online implementation
Revised: The experiments show the importance of SAC as the controller and of the online implementation.

Original: Future directions include  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.
Revised: 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.