Original: E NERGY - BASED P REDICTIVE R EPRESENTATIONS FOR P ARTIALLY O BSERVED R EINFORCEMENT L EARNING
Revised: E NERGY - BASED P REDICTIVE R EPRESENTATIONS FOR P ARTIALLY O BSERVED R EINFORCEMENT L EARNING

Original: Anonymous authors
Revised: Anonymous authors

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

Original: A BSTRACT
Revised: A BSTRACT

Original: In real world applications, it is usually necessary for a reinforcement learning algorithm to handle partial observability, which is not captured by a Markov decision process (MDP).
Revised: In real world applications, it is usually necessary for a reinforcement learning algorithm to handle partial observability, which is not captured by a Markov decision process (MDP).

Original: Although partially observable Markov decision processes (POMDPs) have been precisely motivated by this requirement, they raise significant computational and statistical hardness challenges in learning and planning.
Revised: Although partially observable Markov decision processes (POMDPs) have been precisely motivated by this requirement, they raise significant computational and statistical hardness challenges in learning and planning.

Original: In this work, we introduce the Energy-based Predictive Representation (EPR) to support a unified approach to practical reinforcement learning algorithm design for both the MDP and POMDP settings, which enables learning, exploration, and planning to be handled in a coherent way.
Revised: In this work, we introduce the Energy-based Predictive Representation (EPR) to support a unified approach to practical reinforcement learning algorithm design for both the MDP and POMDP settings, which enables learning, exploration, and planning to be handled in a coherent way.

Original: The proposed framework relies on a powerful neural energy-based model to extract a sufficient representation, from which Q -functions can be efficiently approximated.
Revised: The proposed framework relies on a powerful neural energy-based model to extract a sufficient representation, from which Q -functions can be efficiently approximated.

Original: With such a representation, confidence can be efficiently computed to allow optimism/pessimism in the face of uncertainty to be efficiently implemented in planning, enabling effective management of the exploration versus exploitation tradeoff.
Revised: With such a representation, confidence can be efficiently computed to allow optimism/pessimism in the face of uncertainty to be efficiently implemented in planning, enabling effective management of the exploration versus exploitation tradeoff.

Original: An experimental investigation shows that the proposed algorithm can surpass state-of-the-art performance in both
Revised: An experimental investigation shows that the proposed algorithm can surpass state-of-the-art performance in both

Original: MDP and POMDP settings in comparison to existing baselines.
Revised: MDP and POMDP settings in comparison to existing baselines.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Reinforcement learning (RL) based on Markov Decision Processes (MDPs) hashas proved to be extremely effective in several real world decision-making problems (Levine et al., 2016; Jiang et al., 2021).
Revised: Reinforcement learning (RL) based on Markov Decision Processes (MDPs) has proved to be extremely effective in several real world decision-making problems (Levine et al., 2016; Jiang et al., 2021).

Original: However, the success of most RL algorithms (Ren et al., 2022b; Zhang et al., 2022) relies heavily on the assumption that the environment state is fully observable to the agent.
Revised: However, the success of most RL algorithms (Ren et al., 2022b; Zhang et al., 2022) relies heavily on the assumption that the environment state is fully observable to the agent.

Original: In practice, such an assumption can be easily violated in the presence of observational noise.
Revised: In practice, such an assumption can be easily violated in the presence of observational noise.

Original: To address this issue,
Revised: To address this issue,

Original: Partially Observable Markov Decision Processes (POMDPs) ( ˚Astr ¨ om, 1965) have been proposed for capturing the inherent uncertainty about the state arising from partial observations.
Revised: Partially Observable Markov Decision Processes (POMDPs) ( ˚Astr ¨ om, 1965) have been proposed for capturing the inherent uncertainty about the state arising from partial observations.

Original: However, the flexibility of POMDPs creates significant statistical and computational hardness in terms of planning, exploration and learning.
Revised: However, the flexibility of POMDPs creates significant statistical and computational hardness in terms of planning, exploration and learning.

Original: In particular, i) , partial observability induces a nonMarkovian dependence over the entire history ; and ii) , the expanded spaces of observation sequences or state space distributions incur significant representation challenges.
Revised: In particular, i) , partial observability induces a nonMarkovian dependence over the entire history ; and ii) , the expanded spaces of observation sequences or state space distributions incur significant representation challenges.

Original: In fact, due to the full history dependence, it has been proved that the planning for even finite-horizon tabular POMDPs is NP-hard without additional structural assumptions (Papadimitriou & Tsitsiklis, 1987; Madani et al., 1998), and the sample complexity for learning POMDPs can be exponential with respect to the horizon (Jin et al., 2020a).
Revised: In fact, due to the full history dependence, it has been proved that the planning for even finite-horizon tabular POMDPs is NP-hard without additional structural assumptions (Papadimitriou & Tsitsiklis, 1987; Madani et al., 1998), and the sample complexity for learning POMDPs can be exponential with respect to the horizon (Jin et al., 2020a).

Original: These complexities only become more demanding in continuous state spaces and real-world scenarios.
Revised: These complexities only become more demanding in continuous state spaces and real-world scenarios.

Original: On the other hand, despite the theoretical hardness, the widely used sliding window policy parameterization has demonstrated impressive empirical performance (Mnih et al., 2013; Berner et al., 2019), indicating that there is sufficient structure in real-world POMDPs that can be exploited to bypass the aforementioned complexities.
Revised: On the other hand, despite the theoretical hardness, the widely used sliding window policy parameterization has demonstrated impressive empirical performance (Mnih et al., 2013; Berner et al., 2019), indicating that there is sufficient structure in real-world POMDPs that can be exploited to bypass the aforementioned complexities.

Original: Recently, observable POMDPs with invertible emissions have been investigated to justify the sliding window heuristic in tabular cases (Azizzadenesheli et al., 2016;
Revised: Recently, observable POMDPs with invertible emissions have been investigated to justify the sliding window heuristic in tabular cases (Azizzadenesheli et al., 2016;

Original: Guo et al., 2016; Jin et al., 2020a; Golowich et al., 2022a), which has been further extended with function approximation for large and continuous state POMDPs (Wang et al., 2022; Uehara et al., 2022).
Revised: Guo et al., 2016; Jin et al., 2020a; Golowich et al., 2022a), which has been further extended with function approximation for large and continuous state POMDPs (Wang et al., 2022; Uehara et al., 2022).

Original: Although these algorithms can exploit particular structure efficiently in terms of the sample complexity, they rely on unrealistic computation oracles, and are thus not applicable in practice.
Revised: Although these algorithms can exploit particular structure efficiently in terms of the sample complexity, they rely on unrealistic computation oracles, and are thus not applicable in practice.

Original: In this paper, we consider the following natural question:
Revised: In this paper, we consider the following natural question:

Original: How can one design efficient and practical algorithms for structured POMDPs
Revised: How can one design efficient and practical algorithms for structured POMDPs?

Original: In particular, we would like to exploit special structures that allows approximation to bypass inherent worst-case difficulties.
Revised: In particular, we would like to exploit special structures that allows approximation to bypass inherent worst-case difficulties.

Original: By “efficient” we mean considering learning , planning and exploration in a unified manner that can balance errors in each component and reduce unnecessary computation, while by “practical” we mean the algorithm retains sufficient flexibility and can be easily implemented and deployed in real-world scenarios.
Revised: By “efficient” we mean considering learning , planning and exploration in a unified manner that can balance errors in each component and reduce unnecessary computation, while by “practical” we mean the algorithm retains sufficient flexibility and can be easily implemented and deployed in real-world scenarios.

Original: There have been many attempts to address this question.
Revised: There have been many attempts to address this question.

Original: The most straightforward idea is to extend model-free RL methods, including policy gradient and Q -learning, with a memory-limited parametrization, e.g. , recurrent neural networks (Wierstra et al., 2007; Hausknecht & Stone, 2015; Zhu et al., 2017).
Revised: The most straightforward idea is to extend model-free RL methods, including policy gradient and Q -learning, with a memory-limited parametrization, e.g. , recurrent neural networks (Wierstra et al., 2007; Hausknecht & Stone, 2015; Zhu et al., 2017).

Original: Alternatively, in model-based RL (Kaelbling et al., 1998), an approximation of the latent dynamics can be estimated and a posterior over latent states ( i.e. , beliefs) maintained, in principle allowing an optimal policy to be extracted via dynamic programming upon beliefs.
Revised: Alternatively, in model-based RL (Kaelbling et al., 1998), an approximation of the latent dynamics can be estimated and a posterior over latent states ( i.e. , beliefs) maintained, in principle allowing an optimal policy to be extracted via dynamic programming upon beliefs.

Original: Following this idea, Deisenroth & Peters (2012) and Igl et al.
Revised: Following this idea, Deisenroth & Peters (2012) and Igl et al.

Original: ; Gregor et al.
Revised: ; Gregor et al.

Original: ; Zhang et al.
Revised: ; Zhang et al.

Original: ; Lee et al.
Revised: ; Lee et al.

Original: (2020) consider Gaussian process or deep model parametrizations, respectively.
Revised: (2020) consider Gaussian process or deep model parametrizations, respectively.

Original: Such methods are designed based on implicit assumptions about structure through the parameterization choices of the models.
Revised: Such methods are designed based on implicit assumptions about structure through the parameterization choices of the models.

Original: However, these approaches suffer from sub-optimal performance due to several compounding factors: i), approximation error from inaccurate parametrizations of the learnable components (policy, value function, model, belief), ii), a sub-optimal policy induced by approximated planning (through policy gradient or dynamic progamming), and iii), the neglect of exploration when interacting with the environment.
Revised: However, these approaches suffer from sub-optimal performance due to several compounding factors: i), approximation error from inaccurate parametrizations of the learnable components (policy, value function, model, belief), ii), a sub-optimal policy induced by approximated planning (through policy gradient or dynamic progamming), and iii), the neglect of exploration when interacting with the environment.

Original: As an alternative, spectral representation approaches provide an alternative strategy based on extracting a sufficient representation that can support learning, planning and exploration.
Revised: As an alternative, spectral representation approaches provide an alternative strategy based on extracting a sufficient representation that can support learning, planning and exploration.

Original: In this vein Azizzadenesheli et al.
Revised: In this vein Azizzadenesheli et al.

Original: (2016) investigate spectral methods (Anandkumar et al., 2014) for latent variable model estimation in POMDPs, but only consider tabular scenarios with finite state and action cases.
Revised: (2016) investigate spectral methods (Anandkumar et al., 2014) for latent variable model estimation in POMDPs, but only consider tabular scenarios with finite state and action cases.

Original: Predictive State Representations (PSR)
Revised: Predictive State Representations (PSR)

Original: (Littman & Sutton, 2001; Boots et al., 2011) also leverage spectral decomposition, but instead of recovering an underlying latent variable model, they learn an equivalent sufficient representation of belief.
Revised: (Littman & Sutton, 2001; Boots et al., 2011) also leverage spectral decomposition, but instead of recovering an underlying latent variable model, they learn an equivalent sufficient representation of belief.

Original: These methods have been extended to real-world settings with continuous observations and actions by exploiting kernel embeddings (Boots et al., 2013) or deep models (Downey et al., 2017; Venkatraman et al., 2017; Guo et al., 2018).
Revised: These methods have been extended to real-world settings with continuous observations and actions by exploiting kernel embeddings (Boots et al., 2013) or deep models (Downey et al., 2017; Venkatraman et al., 2017; Guo et al., 2018).

Original: However, efficient exploration and tractable planning with spectral representations has yet to be thoroughly developed (Zhan et al., 2022).
Revised: However, efficient exploration and tractable planning with spectral representations has yet to be thoroughly developed (Zhan et al., 2022).

Original: In this paper, we propose Energy-based Predictive Representation (EPR) to support efficient and tractable learning, planning, and exploration in POMDPs (and MDPs), as a solution to the aforementioned question.
Revised: In this paper, we propose Energy-based Predictive Representation (EPR) to support efficient and tractable learning, planning, and exploration in POMDPs (and MDPs), as a solution to the aforementioned question.

Original: More specifically:
Revised: More specifically:

Original: • We propose a flexible nonlinear energy-based model for induced belief-state MDPs without explicit parameterization of beliefs , providing a principled linear sufficient representation for the state-action value function.
Revised: • We propose a flexible nonlinear energy-based model for induced belief-state MDPs without explicit parameterization of beliefs , providing a principled linear sufficient representation for the state-action value function.

Original: • We reveal the connection between EPR and PSR, while also illustrating the differences, to demonstrate the modeling ability of the proposed EPR.
Revised: • We reveal the connection between EPR and PSR, while also illustrating the differences, to demonstrate the modeling ability of the proposed EPR.

Original: • We provide computationally-tractable learning and planning algorithms for EPR that implement the principles of optimism and pessimism in the face of uncertainty for online and offline RL, balancing exploration and exploitation.
Revised: • We provide computationally-tractable learning and planning algorithms for EPR that implement the principles of optimism and pessimism in the face of uncertainty for online and offline RL, balancing exploration and exploitation.

Original: • We conduct a comprehensive comparison to existing state-of-the-art RL algorithms in both MDP and POMDP benchmarks, demonstrating superior empirical performance of the proposed EPR.
Revised: • We conduct a comprehensive comparison to existing state-of-the-art RL algorithms in both MDP and POMDP benchmarks, demonstrating superior empirical performance of the proposed EPR.

Original:  P RELIMINARIES
Revised:  P RELIMINARIES

Original: In this section, we briefly introduce POMDPs and their degenerate case of MDPs, identifying the special structures that will be used to derive the proposed representation learning method.
Revised: In this section, we briefly introduce POMDPs and their degenerate case of MDPs, identifying the special structures that will be used to derive the proposed representation learning method.

Original: Partially Observable Markov Decision Processes.
Revised: Partially Observable Markov Decision Processes.

Original: Formally, we define a partially observable
Revised: Formally, we define a partially observable

Original: Markov decision process (POMDP) as a tuple P = ( S , A , O , r, H, µ, P, O ) , where H is a positive integer denoting the length of horizon; µ is the initial distribution of state, r : S ×
Revised: Markov decision process (POMDP) as a tuple P = ( S , A , O , r, H, µ, P, O ) , where H is a positive integer denoting the length of horizon; µ is the initial distribution of state, r : S ×

Original: A → [0 , 1] , the reward function, and S , A , O denote the state, action and observation space, respectively.
Revised: A → [0 , 1] , the reward function, and S , A , O denote the state, action and observation space, respectively.

Original: P ( ·| s, a ) :
Revised: P ( ·| s, a ) :

Original: S × A → ∆( S ) is the transition kernel, capturing the dynamics between states, and O ( ·| s ) : S → ∆( O ) is the emission kernel, where ∆ ( · ) denotes the set of probability measures over the support.
Revised: S × A → ∆( S ) is the transition kernel, capturing the dynamics between states, and O ( ·| s ) : S → ∆( O ) is the emission kernel, where ∆ ( · ) denotes the set of probability measures over the support.

Original: Initially, given a state s 1 ∼ µ ( s ) as a starting point, at each step h ∈ [0 , H ] , the agent takes an action a ∈ A , a new state s h +1 is generated s h +1 ∼ P ( ·| s h , a h ) , and the agent observes o h +1 ∼ O
Revised: Initially, given a state s 1 ∼ µ ( s ) as a starting point, at each step h ∈ [1 , H ] , the agent takes an action a ∈ A , a new state s h +1 is generated s h +1 ∼ P ( ·| s h , a h ) , and the agent observes o h +1 ∼ O

Original: ( ·| s h +1 ) and reward r ( s h
Revised: ( ·| s h +1 ) and reward r ( s h

Original: +1 , a h +1 ) .
Revised: +1 , a h +1 ) .

Original: Due to partial observability, the dependence between observations is non-Markovian, hence, we define a policy π
Revised: Due to partial observability, the dependence between observations is non-Markovian, hence, we define a policy π

Original: = { π t } where π t : O× ( A×O ) t → ∆( A ) to depend on the whole history, i.e. , x t = { o 0 , { a i , o i +1 } t − 1 i =0 } .
Revised: = { π t } where π t : O× ( A×O ) t → ∆( A ) to depend on the whole history, i.e. , x t = { o 0 , { a i , o i +1 } t − 1 i =0 } .

Original: The corresponding value for policy π can be defined as V π = E π (cid:104)(cid:80)
Revised: The corresponding value for policy π can be defined as V π = E π (cid:104)(cid:80)

Original: Hh =1 r ( s h , a h ) (cid:105) , and the objective is to find the optimal policy π ∗ = arg max π V π .
Revised: Hh =1 r ( s h , a h ) (cid:105) , and the objective is to find the optimal policy π ∗ = arg max π V π .

Original: Markov decision processes (MDPs) are a degenerate case of POMDPs, where S =
Revised: Markov decision processes (MDPs) are a degenerate case of POMDPs, where S =

Original: O
Revised: O

Original: and
Revised: and

Original: O ( o | s )
Revised: O ( o | s )

Original: =
Revised: =

Original: δ ( o = s ) , and can be specified as M = ( S , A , r, H, µ, P ) .
Revised: δ ( o = s ) , and can be specified as M = ( S , A , r, H, µ, P ) .

Original: One can also convert a
Revised: One can also convert a

Original: POMDP to an MDP by treating the whole history x t = { o 0 , { a i , o i +1 } t − 1 i =0 } as the state.
Revised: POMDP to an MDP by treating the whole history x t = { o 0 , { a i , o i +1 } t − 1 i =0 } as the state.

Original: Specifically, following (Kaelbling et al., 1998), we define the belief b : O × ( A × O ) t → ∆( S ) , ∀ t ∈ N + , which can be recursively defined as: b ( s 1 | o 1 ) =
Revised: Specifically, following (Kaelbling et al., 1998), we define the belief b : O × ( A × O ) t → ∆( S ) , ∀ t ∈ N + , which can be recursively defined as: b ( s 1 | o 1 ) =

Original: P ( s 1 | o 1 ) , and
Revised: P ( s 1 | o 1 ) , and

Original: [Equation]
Revised: [Equation]

Original: Each entry of the belief state describes the probability of the underlying state given the past history.
Revised: Each entry of the belief state describes the probability of the underlying state given the past history.

Original: Furthermore, with a slight abuse of notation, we use b t to denote the belief state at step t .
Revised: Furthermore, with a slight abuse of notation, we use b t to denote the belief state at step t .

Original: Then, one can construct the equivalent belief MDPM b =
Revised: Then, one can construct the equivalent belief MDP

Original: ( X , A , R h , H, µ b , T b ) with X denoting the set of possible histories, and
Revised: [Equation] , R h , H, µ b , T b ) with X denoting the set of possible histories, and

Original: [Equation]
Revised: [Equation]

Original: Therefore, the corresponding value function V πh ( b h ) and Q πh ( b h , a h ) for the belief MDP given a policy π can be defined as:
Revised: Therefore, the corresponding value function V πh ( b h ) and Q πh ( b h , a h ) for the belief MDP given a policy π can be defined as:

Original: [Equation]
Revised: [Equation]

Original: Following the MDP perspective, we also have the Bellman recursive equation:
Revised: Following the MDP perspective, we also have the Bellman recursive equation:

Original: [Equation]
Revised: [Equation]

Original: One can still apply a dynamic programming style approach to solve POMDPs according to (4), however since the belief depends on the entire history the number of possible beliefs can still be infinite even the number of states is finite.
Revised: One can still apply a dynamic programming style approach to solve POMDPs according to (4), however since the belief depends on the entire history the number of possible beliefs can still be infinite even the number of states is finite.

Original: To combat with these essential difficulties, we will leverage two particular structures, observability and linearity, as introduced below.
Revised: To combat with these essential difficulties, we will leverage two particular structures, observability and linearity, as introduced below.

Original: Observability in POMDPs.
Revised: Observability in POMDPs.

Original: It has been shown (Even-Dar et al., 2007; Golowich et al., 2022b;
Revised: It has been shown (Even-Dar et al., 2007; Golowich et al., 2022b;

Original: Uehara et al., 2022) that for POMDPs with an observability assumption, one can safely relax the history dependence with a short window, bypassing the exponential sample and planning complexity w.r.t.
Revised: Uehara et al., 2022) that for POMDPs with an observability assumption, one can safely relax the history dependence with a short window, bypassing the exponential sample and planning complexity w.r.t.

Original: horizon length (Golowich et al., 2022a;b).
Revised: horizon length (Golowich et al., 2022a;b).

Original: Specifically, the observability property for POMDPs is defined as follows.
Revised: Specifically, the observability property for POMDPs is defined as follows.

Original: Assumption 1 ((Even-Dar et al., 2007; Golowich et al., 2022b)) .
Revised: Assumption 1 ((Even-Dar et al., 2007; Golowich et al., 2022b)) .

Original: The POMDP with emission model O satisfies γ -observability if for arbitrary beliefs b and b ′ over states, ∥⟨ O, b ⟩ − ⟨ O, b ′ ⟩∥ 1 ⩾  ∥ 
Revised: The POMDP with emission model O satisfies γ -observability if for arbitrary beliefs b and b ′ over states, ∥⟨ O, b ⟩ − ⟨ O, b ′ ⟩∥ 1 ⩾ γ ∥ b − b ′

Original: b − b ′ ∥ 1 , where ⟨ O, b ⟩ := (cid:82) O ( o | s ) b ( s ) ds .
Revised: ∥ 1 , where ⟨ O, b ⟩ := (cid:82) O ( o | s ) b ( s ) ds .

Original: A key consequence of observability is that, the belief can be well approximated with a short history window (Golowich et al., 2022b), and one can construct an approximate MDP based on a finite belief history, which eliminates the exponential complexity induced by full history dependence.
Revised: A key consequence of observability is that, the belief can be well approximated with a short history window (Golowich et al., 2022b), and one can construct an approximate MDP based on a finite belief history, which eliminates the exponential complexity induced by full history dependence.

Original: Specifically, we denote L as the length of the window.
Revised: Specifically, we denote L as the length of the window.

Original: Then, defining x Lt = (cid:110)
Revised: Then, defining x Lt = (cid:110)

Original: o t − L , { a i , o i +1 } ti = t − L (cid:111)
Revised: o t − L , { a i , o i +1 } ti = t − L (cid:111)

Original: ∈
Revised: ∈

Original: X L , the approximated beliefs b L follow the same recursive definition as (1) but with only finite history x Lt starting from the uniform belief.
Revised: X L , the approximated beliefs b L follow the same recursive definition as (1) but with only finite history x Lt starting from the uniform belief.

Original: This immediately induces an approximate MDP M Lb
Revised: This immediately induces an approximate MDP M Lb

Original: = (cid:0)
Revised: = (cid:0)

Original: X L , A , R Lh , H, µ b , T Lb (cid:1) according to (3) with b L , instead of b .
Revised: X L , A , R Lh , H, µ b , T Lb (cid:1) according to (3) with b L , instead of b .

Original: Theorem 2.1 in Golowich et al.
Revised: Theorem 2.1 in Golowich et al.

Original: (2022a) proves that the approximation error of the finite-memory belief MDP is small for observable
Revised: (2022a) proves that the approximation error of the finite-memory belief MDP is small for observable

Original: POMDPs.
Revised: POMDPs.

Original: Hence, with slightly abuse of notation, we still use b to represent b L throughout the paper.
Revised: Hence, with slight abuse of notation, we still use b to represent b L throughout the paper.

Original: Linearity in MDPs.
Revised: Linearity in MDPs.

Original: To handle the complexity induced by large state spaces, linear/low-rank structures have been introduced in MDPs (Jin et al., 2020b; Agarwal et al., 2020) for effective function approximation, which leverages spectral factorization of the transition dynamics and reward:[Equation]
Revised: To handle the complexity induced by large state spaces, linear/low-rank structures have been introduced in MDPs (Jin et al., 2020b; Agarwal et al., 2020) for effective function approximation, which leverages spectral factorization of the transition dynamics and reward:

Original: [Equation]
Revised: [Equation]

Original: where ϕ : S ×
Revised: where ϕ : S ×

Original: A → H , µ : S → H are two feature maps to a Hilbert space H .
Revised: A → H , µ : S → H are two feature maps to a Hilbert space H .

Original: Under such an assumption, we can represent the state-action value function Q π for an arbitrary policy π by:
Revised: Under such an assumption, we can represent the state-action value function Q π for an arbitrary policy π by:

Original: [Equation]
Revised: [Equation]

Original: which implies that instead of a complicated function space defined on the raw state space, one can design a computationally efficient planning and sample efficient exploration algorithm in the space linearly spanned by ϕ .
Revised: which implies that instead of a complicated function space defined on the raw state space, one can design a computationally efficient planning and sample efficient exploration algorithm in the space linearly spanned by ϕ .

Original: In fact, from the correspondence between policy and Q -function as discussed in (Ren et al., 2022a), ϕ can be understood as representing primitives for skill set construction.
Revised: In fact, from the correspondence between policy and Q -function as discussed in (Ren et al., 2022a), ϕ can be understood as representing primitives for skill set construction.

Original: Efficient and practical algorithms have been designed for exploiting linearity in MDPs (Zhang et al., 2022; Qiu et al., 2022), which inspires us to exploit similar properties in POMDPs. 
Revised: Efficient and practical algorithms have been designed for exploiting linearity in MDPs (Zhang et al., 2022; Qiu et al., 2022), which inspires us to exploit similar properties in POMDPs. Energy-based Models.

Revised: Energy-based Models are one of the most flexible models to represent the conditional probability measure.

Revised: It takes the form of p ( y | x ) = exp( − f ( x, y )) /Z

Revised: ( x ) where f ( x, y ) , which can be parametrized by deep models, is the energy of ( x, y ) and Z ( x ) is a partition function that only depends on x to guarantee p ( y | x ) is a valid probability measure.

Original:   E NERGY - BASED P REDICTIVE R EPRESENTATION
Revised: When y is discrete, we have that p ( y | x ) = exp( − f ( x, y )) / (cid:80) y exp( − f ( x, y )) , which corresponds to the standard softmax probability where − f ( x, y ) is the softmax logits. We refer the interested readers to Song & Kingma (2021) for the training methods of energy-based models.  E NERGY - BASED P REDICTIVE R EPRESENTATION

Original: W propose Energy-based Predictive Representation (EPR) , which introduces linearity into finitehistory approximated POMDPs, allowing the complexity induced by large state spaces and long histories to be overcome, yielding improved efficiency for learning, planning and exploration.
Revised: We propose Energy-based Predictive Representation (EPR) , which introduces linearity into finitehistory approximated POMDPs, allowing the complexity induced by large state spaces and long histories to be overcome, yielding improved efficiency for learning, planning and exploration.

Original: We emphasize that the proposed method is also applicable to MDPs.
Revised: We emphasize that the proposed method is also applicable to MDPs.

Original: The approach builds upon recent progress in large-state MDPs (Zhang et al., 2022; Qiu et al., 2022) that leverages linear structure in the dynamics, P ( s ′ | s, a ) = ⟨ ϕ ( s, a ) , µ ( s ′ ) ⟩ , to obtain an efficient and practical framework for learning, planning and exploration.
Revised: The approach builds upon recent progress in large-state MDPs (Zhang et al., 2022; Qiu et al., 2022) that leverages linear structure in the dynamics, P ( s ′ | s, a ) = ⟨ ϕ ( s, a ) , µ ( s ′ ) ⟩ , to obtain an efficient and practical framework for learning, planning and exploration.

Original: Recall the construction of a finitememory belief MDP to approximate a POMDP discussed in Section 2, which avoids full history dependence.
Revised: Recall the construction of a finitememory belief MDP to approximate a POMDP discussed in Section 2, which avoids full history dependence.

Original: For such a constructed belief MDP, a natural idea is to apply linear MDP algorithms, i.e. , extracting the linear decomposition for T Lb ( b ′ | b, a ) =
Revised: For such a constructed belief MDP, a natural idea is to apply linear MDP algorithms, i.e. , extracting the linear decomposition for T Lb ( b ′ | b, a ) =

Original: ⟨ ϕ ( b, a ) , µ  b ′ ) ⟩ , to handle the hardnesses of POMDPs mentioned in Section 1.
Revised: ⟨ ϕ ( b, a ) , µ ( b ′ ) ⟩ , to handle the hardnesses of POMDPs mentioned in Section 1.

Original: However, there are several difficulties in such an straightforward extension: 
Revised: However, there are several difficulties in such a straightforward extension: i ,

Original: i , the set of beliefs is proportional to the number of states, which could be infinite; ii , the factorization of the transition dynamics (3) in the belief MDP isdifficult.
Revised: the set of beliefs is proportional to the number of states, which could be infinite ; ii , the factorization of the transition dynamics (3) in the belief MDP is

Revised: difficult.

Original: These difficulties hinder the extension of linear MDPs to observable POMDPs.
Revised: These difficulties hinder the extension of linear MDPs to observable POMDPs.

Original: However, note that we never explicitly require the beliefs and their dynamics, but only the representation ϕ ( b, a ) .
Revised: However, note that we never explicitly require the beliefs and their dynamics, but only the representation ϕ ( b, a ) .

Original: As beliefs are functions over finite-window histories, the representation can also be rewritten as ϕ ( x t , a t ) , which suggests that one might bypass the inherent difficulties by a reprameterization trick .
Revised: As beliefs are functions over finite-window histories, the representation can also be rewritten as ϕ ( x t , a t ) , which suggests that one might bypass the inherent difficulties by a reprameterization trick .

Original: Consider the energy-based parametrization (Arbel et al., 2020) for P ( o t +1 | b ( x t ) , a t ) 
Revised: Consider the energy-based parametrization (Arbel et al., 2020) for P ( o t +1 | b ( x t ) , a t ) where b ( x t ) is the belief for history x t :

Original: [Equation]
Revised: [Equation]

Original: where λ is a scalar, p ( o ) is a fixed distribution and the normalization condition enforces that the energybased model P ( o t +1 | b ( x t ) , a t ) is a valid distribution.
Revised: where λ is a scalar, p ( o ) is a fixed distribution and the normalization condition enforces that the energybased model P ( o t +1 | b ( x t ) , a t ) is a valid distribution.

Original: We avoid any explicit parametrization and computation of beliefs b , while preserving dependence through f and g , which will be learned jointly.
Revised: We avoid any explicit parametrization and computation of beliefs b , while preserving dependence through f and g , which will be learned jointly.

Original: Compared to standard parametrizations, we do not need to specify unnecessary model parameters for the transition dynamics P and emmission O , and bypass any learning and approximation of beliefs that induce compounding errors.
Revised: Compared to standard parametrizations, we do not need to specify unnecessary model parameters for the transition dynamics P and emmission O , and bypass any learning and approximation of beliefs that induce compounding errors.

Original: As a special case, we note that the observable Linear-Quadratic
Revised: As a special case, we note that the observable Linear-Quadratic

Original: Gaussian (LQG) actually follows (6) with a specific λ and p ( o ) .
Revised: Gaussian (LQG) actually follows (6) with a specific λ and p ( o ) .

Original: See Appendix B for details.
Revised: See Appendix C for details.

Original: Meanwhile, this approach also provides a linear factorization of T ( b t +1 | b t , a t ) almost for free.
Revised: Meanwhile, this approach also provides a linear factorization of T ( b t +1 | b t , a t ) almost for free.

Original: By viewing the proposed parameterization (6) as a kernel and following the random Fourier feature trick (Rahimi & Recht, 2007; Choromanski et al., 2020; Ren et al., 2022b), onecan write
Revised: By viewing the proposed parameterization (6) as a kernel and following the random Fourier feature trick (Rahimi & Recht, 2007; Choromanski et al., 2020; Ren et al., 2022b), one can write

Original: [Equation]
Revised: [Equation]

Original: ) where ω i ∼ N (0 , I d ) and
Revised: where ω i ∼ N (0 , I d ) and

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: which can be derived from the softmax random feature from Choromanski et al.
Revised: which can be derived from the softmax random feature from Choromanski et al.

Original:  Substituting (8) into (3) yields the factorization of T b as
Revised: We also provide a derivation in Appendix B. Substituting (8) into (3) yields the factorization of T b as

Original: [Equation]
Revised: [Equation]

Original: where µ ( b t +1 ) :
Revised: where µ ( b t +1 ) :

Original: = (cid:82) O 1 b ( ·| x t +1 )
Revised: = (cid:82) O 1 b ( ·| x t +1 )

Original: =
Revised: =

Original: b ( ·| x t ,a t ,o t +1 ) ψ ω ( o t +1 )
Revised: b ( ·| x t ,a t ,o t +1 ) ψ ω ( o t +1 )

Original: d o t +1 .
Revised: d o t +1 .

Original: With this formula, we obtain a valid linear representation ϕ
Revised: With this formula, we obtain a valid linear representation ϕ

Original: ( b t , a t ) as an Energy-based Predictive Representation (EPR) for a belief MDP without any explicit beliefs.
Revised: ( b t , a t ) as an Energy-based Predictive Representation (EPR) for a belief MDP without any explicit beliefs.

Original: To learn the EPR given data D :=
Revised: To learn the EPR given data D :=

Original: { o t − 1 , a t , r t } Ht =1 , we exploit maximum likelihood estimation (MLE) of (6),
Revised: { o t − 1 , a t , r t } Ht =1 , we exploit maximum likelihood estimation (MLE) of (6),

Original: [Equation]
Revised: [Equation]

Original: To ensure the constraints, we add a regularization term
Revised: To ensure the constraints, we add a regularization term

Original: [Equation]
Revised: [Equation]

Original: with o i ∼ p .
Revised: with o i ∼ p .

Original: The objective will be
Revised: The objective will be

Original: [Equation]
Revised: [Equation]

Original: In practice, we can further simplify the objective by normalizing the ˜f
Revised: In practice, we can further simplify the objective by normalizing the ˜

Original: ( x t , a t ) =
Revised: [Equation]

Original: f ( x t ,a t ) ∥ f ( x t ,a t ) ∥ 2 , obtaining the final objective
Revised: , obtaining the final objective

Original: [Equation]
Revised: [Equation]

Original: which reduces to a contrastive loss that can be optimized by stochastic gradient descent with a deep network parameterization of ˜ f and g .
Revised: which reduces to a contrastive loss that can be optimized by stochastic gradient descent with a deep network parameterization of ˜ f and g .

Original: We obtain negative samples { o i } mi =1 ∼ p ( o ) from a mixture of replay buffer and collected trajectories.
Revised: We obtain negative samples { o i } mi =1 ∼ p ( o ) from a mixture of replay buffer and collected trajectories.

Original: Before we introduce an exploration-exploitation mechanism with EPR in Section 3.1, we first discuss the relationship between the proposed EPR, predictive state representations (PSR)
Revised: Before we introduce an exploration-exploitation mechanism with EPR in Section 3.1, we first discuss the relationship between the proposed EPR, predictive state representations (PSR)

Original: (Littman & Sutton, 2001; Singh et al., 2004), and spectral dynamics embedding (SPEDE) (Ren et al., 2022b).
Revised: (Littman & Sutton, 2001; Singh et al., 2004), and spectral dynamics embedding (SPEDE) (Ren et al., 2022b).

Original: Connection to PSR (Littman & Sutton, 2001; Singh et al., 2004):
Revised: Connection to PSR (Littman & Sutton, 2001; Singh et al., 2004):

Original: Th predictive state representation (PSR) was also proposed for bypassing belief calculation by factorizing a variant of the transition dynamics operator.
Revised: The predictive state representation (PSR) was also proposed for bypassing belief calculation by factorizing a variant of the transition dynamics operator.

Original: Specifically, given the history ( x t , a t ) , the probability for observing a test, i.e. , the finite sequence of events y
Revised: Specifically, given the history ( x t , a t ) , the probability for observing a test, i.e. , the finite sequence of events y

Original: =
Revised: =

Original: ( a t +1 , o t +1 , · · · , a t + k , o t + k ) with k ∈ N , is p ( y | x ) :
Revised: ( a t +1 , o t +1 , · · · , a t + k , o t + k ) with k ∈ N , is p ( y | x ) :

Original: = p ( o t + k t +1 | x t , a t + k t ) .
Revised: = p ( o t + k t +1 | x t , a t + k t ) .

Original: For time step t , one can construct a set of core tests U
Revised: For time step t , one can construct a set of core tests U

Original: = { u i , . . .
Revised: = { u i , . . .

Original: , u k } as sufficient statistics for history x t , such that for any test τ , p ( τ | x t ) = ⟨ p ( U | x t ) , w τ ⟩ for some weights w τ, ∈ R | U t | .
Revised: , u k } as sufficient statistics for history x t , such that for any test τ , p ( τ | x t ) = ⟨ p ( U | x t ) , w τ ⟩ for some weights w τ, ∈ R | U t | .

Original: The forward dynamics can be represented in a PSR by Bayes’ rule: p ( τ | x t , a t , o t +1 ) =
Revised: The forward dynamics can be represented in a PSR by Bayes’ rule: p ( τ | x t , a t , o t +1 ) =

Original: w ⊤ ( τ,at,ot +1 ) p
Revised: w [Equation] w

Original: ( U | x t ) w ⊤ ( at,ot +1 ) p ( U | x t ) , which implies that a PSR updates with new observations and actions by repeating a calculation for each u i ∈
Revised: [Equation] ) p ( U | x t ) , which implies that a PSR updates with new observations and actions by repeating a calculation for each u i ∈

Original: U .
Revised: U .

Original: Although originally defined for tabular cases,
Revised: Although originally defined for tabular cases,

Original: PSRs have been extended to continuous observations by introducing kernels (Boots et al., 2013) or neural networks (Guo et al., 2018; Downey et al., 2017; Hefny et al., 2018).
Revised: PSRs have been extended to continuous observations by introducing kernels (Boots et al., 2013) or neural networks (Guo et al., 2018; Downey et al., 2017; Hefny et al., 2018).

Original: Obviously, the proposed EPR shares similarities with PSR.
Revised: Obviously, the proposed EPR shares similarities with PSR.

Original: Both factorize conditional distributions defined by the dynamics.
Revised: Both factorize conditional distributions defined by the dynamics.

Original: However, these representations are designed for different purposes, and thus, with different usages and updates.
Revised: However, these representations are designed for different purposes, and thus, with different usages and updates.

Original: Concretely, EPR is proposed for seeking a linear space that can represent the Q -function.
Revised: Concretely, EPR is proposed for seeking a linear space that can represent the Q -function.

Original: The representation is designed to preserve linearity with successive observations without the need for Bayesian updates, which induce extra nonlinearity in PSRs. 
Revised: The representation is designed to preserve linearity with successive observations without the need for Bayesian updates, which induce extra nonlinearity in PSRs. This d

Original: This linear property leads to efficient exploration and planning in EPR; while an efficient exploration and planning algorithm has not yet been discussed for PSR.
Revised: Random Policy Dataset linear property leads to efficient exploration and planning in EPR; while an efficient exploration and planning algorithm has not yet been discussed for PSR.

Original: Connection to SPEDE (Ren et al., 2022b):
Revised: Connection to SPEDE (Ren et al., 2022b):

Original: Linear random features have been proposed for solving planning in MDPs with nonlinear dynamics in (Ren et al., 2022b), where the transition operator is defined as T ( s ′ | s, a ) ∝ exp (cid:16) − ∥ s ′ − f ( s, a ) ∥ 22 / (2 σ 2 ) (cid:17) , corresponding to dynamics s ′ =
Revised: Linear random features have been proposed for solving planning in MDPs with nonlinear dynamics in (Ren et al., 2022b), where the transition operator is defined as T ( s ′ | s, a ) ∝ exp (cid:16) − ∥ s ′ − f ( s, a ) ∥ 22 / (2 σ 2 ) (cid:17) , corresponding to dynamics s ′ =

Original: f ( s, a ) + ϵ with Gaussian noise ϵ ∼ N (0 , σ ) .
Revised: f ( s, a ) + ϵ with Gaussian noise ϵ ∼ N (0 , σ ) .

Original: In addition to the generaliza-

Original: 1: Input: History Embedding f θ ( x, a ) ,

Original: Observation Embedding g θ ( o ) , Random Feature { ω i } ni = where ω i ∼ N (0 , I d ) , Initial Random Policy π 0 , Initial Dataset D = ∅ for online setting.

Original: 2: for Episode i = 1 , · · · , K do 3:

Original: Collect data { ( x i,j , a i,j , o i,j ) } Hj =1 with ˜ π i =

Original: ξπ i

Original: + (1 − ξ )

Original: π 0 , and add the data to D .

Original: Optimize f and g with (18) using the data from D .

Original: Obtain the representation ϕ

Original: ( x t , a t ) via (9) using { ω i } ni =1 .

Original: Add the bonus (19) to the reward and obtain the optimal policy π i +1 with the Q ( x t , a t ) parameterize as ϕ ( x t , a t ) and optimize via FQI.

Original: Optional) Extract policy by soft-AC from learned Q .

Original: 8: end for 9: Return π K . tion of EPR for POMDPs, even in an MDP, EPR considers a general energy-based model,T [Equation] s
Revised: In addition to the generalization of EPR for POMDPs, even in an MDP, EPR considers a general energy-based model,

Original: [Equation]

Original: s, a )

Original: [Equation]

Original: p

Original: [Equation]

Original:  s [Equation] ) exp [Equation]
Revised: T ( s ′ | s, a ) ∝ p ( s ′ ) exp (cid:0)

Original: f [Equation] s, a ) ⊤ ( g ( s ′ )
Revised: f ( s, a ) ⊤ ( g ( s ′ )

Original: + λf ( s, a )) (cid:1) for the dynamics, which is far more flexible than the Gaussian perturbation model considered in SPEDE.
Revised: + λf ( s, a )) (cid:1) for the dynamics, which is far more flexible than the Gaussian perturbation model considered in SPEDE.

Original: 3. O NLINE E XPLORATION AND O FFLINE P OLICY O PTIMIZATION WITH EPR
Revised: 3. O NLINE E XPLORATION AND O FFLINE P OLICY O PTIMIZATION WITH EPR

Original: With an EPR ϕ ( x t , a t ) learned for a POMDP, we can represent the Q -function linearly for the approximated belief MDP, and thus, achieve computationally efficient planning, while calculating confidence bounds for implementing the optimism/pessimism in the face of uncertainty.
Revised: With an EPR ϕ ( x t , a t ) learned for a POMDP, we can represent the Q -function linearly for the approximated belief MDP, and thus, achieve computationally efficient planning, while calculating confidence bounds for implementing the optimism/pessimism in the face of uncertainty.

Original: Exploration and Exploitation with Elliptical Confidence Bound.
Revised: Exploration and Exploitation with Elliptical Confidence Bound.

Original: Given the learned representation ϕ ( x t , a t ) , the confidence bounds can be calculated efficiently, which allows efficient implementation of optimism/pessimism in the face of uncertainty via upper/lower confidence bound (Abbasi-Yadkori et al., 2011; Jin et al., 2020b; Uehara et al., 2021).
Revised: Given the learned representation ϕ ( x t , a t ) , the confidence bounds can be calculated efficiently, which allows efficient implementation of optimism/pessimism in the face of uncertainty via upper/lower confidence bound (Abbasi-Yadkori et al., 2011; Jin et al., 2020b; Uehara et al., 2021).

Original: This is achieved simply by adding an additional elliptical bonus to the R ( x, a ) .
Revised: This is achieved simply by adding an additional elliptical bonus to the R ( x, a ) .

Original: Specifically, given the dataset we collect
Revised: Specifically, given the dataset we collect

Original: D = { ( x Li , a i , R i , o i +1 ) } ni =1 , and calculate the confidence bound as the bonus,
Revised: D = { ( x Li , a i , R i , o i +1 ) } ni =1 , and calculate the confidence bound as the bonus,

Original: [Equation]
Revised: [Equation]

Original: (19) where λ is a pre-specified hyperparameter, and Σ n = (cid:80) ni =1 λI + ϕ
Revised: (19) where λ is a pre-specified hyperparameter, and Σ n = (cid:80) ni =1 λI + ϕ

Original: ( x Li , a i ) ϕ
Revised: ( x Li , a i ) ϕ

Original: ( x Li , a i ) ⊤ .
Revised: ( x Li , a i ) ⊤ .

Original: One can then implement UCB/LCB by adding/subtracting the bonus to R ( x t , a t ) , and performing planning on the modified reward function.
Revised: One can then implement UCB/LCB by adding/subtracting the bonus to R ( x t , a t ) , and performing planning on the modified reward function.

Original: Planning with Obtained Representation.
Revised: Planning with Obtained Representation.

Original: Planning can be conducted by Bellman recursion within the linear space spanned by ϕ ( x Lt , a t ) without a bonus.
Revised: Planning can be conducted by Bellman recursion within the linear space spanned by ϕ ( x Lt , a t ) without a bonus.

Original: However, with an additional bonus term, the
Revised: However, with an additional bonus term, the

Original: Q π no longer lies in the linear space of ϕ , since
Revised: Q π no longer lies in the linear space of ϕ , since

Original: [Equation]
Revised: [Equation]

Revised: π ( x t +1 , a t +1 )] .

Original: As discussed in (Zhang et al., 2022), one can augment the feature space ψ ( x, a ) :
Revised: As discussed in (Zhang et al., 2022), one can augment the feature space ψ ( x, a ) :

Original: = { ϕ ( x, a ) , b ( x, a ) } to ensure the Q -functions can be linearly represented but with an extra O (cid:0) d 2 (cid:1) memory cost.
Revised: = { ϕ ( x, a ) , b ( x, a ) } to ensure the Q -functions can be linearly represented but with an extra O (cid:0) d 2 (cid:1) memory cost.

Original: In practice, we perform fitted Q iteration with a nonlinear component extending the linear parameterization, i.e. , Q ( x, a ) = { w 1 , w 2 } ⊤ (cid:2) ϕ
Revised: In practice, we perform fitted Q iteration with a nonlinear component extending the linear parameterization, i.e. , Q ( x, a ) = { w 1 , w 2 } ⊤ (cid:2) ϕ

Original: ( x, a ) , σ (cid:0) w ⊤ 3 ϕ
Revised: ( x, a ) , σ (cid:0) w ⊤ 3 ϕ

Original: ( x, a ) (cid:1)(cid:3) .
Revised: ( x, a ) (cid:1)(cid:3) .

Original: We provide an outline of our implementation of UCB in Algorithm 1.
Revised: We provide an outline of our implementation of UCB in Algorithm 1.

Original: LCB for pessimistic offline
Revised: LCB for pessimistic offline

Original: RL is similar but using a pre-collected dataset ˆ D without data collection iteration in Step 2, and with the bonus subtracted in Step 6.
Revised: RL is similar but using a pre-collected dataset ˆ D without data collection iteration in Step 2, and with the bonus subtracted in Step 6.

Original: Our algorithm follows the standard interaction paradigm between the agent and the environment, where for each episode, the agent executes the policy and logs the data to the dataset.
Revised: Our algorithm follows the standard interaction paradigm between the agent and the environment, where for each episode, the agent executes the policy and logs the data to the dataset.

Original: Then we perform representation learning and optimistic planning with the Q function parameterized upon the learned representation.
Revised: Then we perform representation learning and optimistic planning with the Q function parameterized upon the learned representation.

Original: Finally, we also extract a policy from the learned Q by soft actor-critic (Haarnoja et al., 2018).
Revised: Finally, we also extract a policy from the learned Q by soft actor-critic (Haarnoja et al., 2018).

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Partial Observability in Reinforcement Learning.
Revised: Partial Observability in Reinforcement Learning.

Original: Despite the essential hardness of POMDPs in terms of learning, planning and exploration (Papadimitriou & Tsitsiklis, 1987;
Revised: Despite the essential hardness of POMDPs in terms of learning, planning and exploration (Papadimitriou & Tsitsiklis, 1987;

Original: Madani et al., 1998;
Revised: Madani et al., 1998;

Original: Vlassis et al., 2012; Jin et al., 2020b), the study of reinforcement learning with partial observations, from both theoretical and empirical aspects, is still attractive due to its practical importance.
Revised: Vlassis et al., 2012; Jin et al., 2020b), the study of reinforcement learning with partial observations, from both theoretical and empirical aspects, is still attractive due to its practical importance.

Original: Algorithmically, model-based/-free algorithms have been extended to POMDPs, explicitly or implicitly exploiting structure.
Revised: Algorithmically, model-based/-free algorithms have been extended to POMDPs, explicitly or implicitly exploiting structure.

Original: Model-based RL algorithms parameterize and learn latent dynamics with an emission model explicitly, and planning through the simulation upon the learned models.
Revised: Model-based RL algorithms parameterize and learn latent dynamics with an emission model explicitly, and planning through the simulation upon the learned models.

Original: A variety of deep models have been proposed recently for better modeling (Watter et al., 2015; Karl et al., 2016; Igl et al., 2018; Zhang et al., 2019; Lee et al., 2020; Hafner et al., 2019a;b; 2020).
Revised: A variety of deep models have been proposed recently for better modeling (Watter et al., 2015; Karl et al., 2016; Igl et al., 2018; Zhang et al., 2019; Lee et al., 2020; Hafner et al., 2019a;b; 2020).

Original: Although deep models indeed provide better approximation ability, they also bring new challenges in terms of planning and exploration, which has not been fully handled.
Revised: Although deep models indeed provide better approximation ability, they also bring new challenges in terms of planning and exploration, which has not been fully handled.

Original: On the other hand, model-free RL algorithms have been extended for POMDPs by learning history dependent value functions and/or policies, through temporal-difference algorithms or policy gradients.
Revised: On the other hand, model-free RL algorithms have been extended for POMDPs by learning history dependent value functions and/or policies, through temporal-difference algorithms or policy gradients.

Original: For example, deep Q -learning (Mnih et al., 2013) concatenates 4 consecutive frames as the input of a deep neural Q -net, which is then improved by recurrent neural networks for longer windows (Bakker, 2001; Hausknecht & Stone, 2015; Zhu et al., 2017).
Revised: For example, deep Q -learning (Mnih et al., 2013) concatenates 4 consecutive frames as the input of a deep neural Q -net, which is then improved by recurrent neural networks for longer windows (Bakker, 2001; Hausknecht & Stone, 2015; Zhu et al., 2017).

Original: Recurrent neural networks have also been exploited for history dependent policies (Schmidhuber, 1990; Bakker, 2001; Wierstra et al., 2007;
Revised: Recurrent neural networks have also been exploited for history dependent policies (Schmidhuber, 1990; Bakker, 2001; Wierstra et al., 2007;

Original: Heess et al., 2015) in policy gradient algorithms as well as actor-critic approaches (Ni et al., 2021; Meng et al., 2021).
Revised: Heess et al., 2015) in policy gradient algorithms as well as actor-critic approaches (Ni et al., 2021; Meng et al., 2021).

Original: Model-free
Revised: Model-free

Original: RL for POMDPs bypasses the planning complexity of model-based RL algorithms.
Revised: RL for POMDPs bypasses the planning complexity of model-based RL algorithms.

Original: However, the difficulty in exploration remains, which leads to suboptimal performance in practice.
Revised: However, the difficulty in exploration remains, which leads to suboptimal performance in practice.

Original: By contrast, the proposed EPR not only can be efficiently learned, but is also equipped with simple yet principled planning and exploration methods, which has not been previously achieved.
Revised: By contrast, the proposed EPR not only can be efficiently learned, but is also equipped with simple yet principled planning and exploration methods, which has not been previously achieved.

Original: Representation Learning for RL.
Revised: Representation Learning for RL.

Original: Successful vision-based representation learning methods have been extended to RL for extracting compact and invariant state-only information from raw-pixels, e.g. , (Laskin et al., 2020a;b; Kostrikov et al., 2020).
Revised: Successful vision-based representation learning methods have been extended to RL for extracting compact and invariant state-only information from raw-pixels, e.g. , (Laskin et al., 2020a;b; Kostrikov et al., 2020).

Original: However, such vision-based features are not specially designed for capturing properties in POMDPs/MDPs essential for decision making.
Revised: However, such vision-based features are not specially designed for capturing properties in POMDPs/MDPs essential for decision making.

Original: To reveal structure that is particularly helpful for RL, many representation learning methods have been designed for different purposes, such as bi-simulation (Ferns et al., 2004; Gelada et al., 2019;
Revised: To reveal structure that is particularly helpful for RL, many representation learning methods have been designed for different purposes, such as bi-simulation (Ferns et al., 2004; Gelada et al., 2019;

Original: Zhang et al., 2020), successor features (Dayan, 1993; Barreto et al., 2017; Kulkarni et al., 2016), spectral decomposition of transition operators (Mahadevan & Maggioni, 2007; Wu et al., 2018; Duan et al., 2019), latent future prediction (Schwarzer et al., 2020; Stooke et al., 2021) and contrastive learning (Oord et al., 2018; Mazoure et al., 2020; Nachum & Yang, 2021;
Revised: Zhang et al., 2020), successor features (Dayan, 1993; Barreto et al., 2017; Kulkarni et al., 2016), spectral decomposition of transition operators (Mahadevan & Maggioni, 2007; Wu et al., 2018; Duan et al., 2019), latent future prediction (Schwarzer et al., 2020; Stooke et al., 2021) and contrastive learning (Oord et al., 2018; Mazoure et al., 2020; Nachum & Yang, 2021;

Original: Yang et al., 2021).
Revised: Yang et al., 2021).

Original: These representation methods ignore the requirement of planning tractability.
Revised: These representation methods ignore the requirement of planning tractability.

Original: Moreover, they are learning based on a pre-collected dataset, which ignores the exploration issue.
Revised: Moreover, they are learning based on a pre-collected dataset, which ignores the exploration issue.

Original: Features that are able to represent value functions are desirable for efficient planning and exploration.
Revised: Features that are able to represent value functions are desirable for efficient planning and exploration.

Original: Based on the linear MDPs structure (Jin et al., 2020b), several theoretical algorithms (Agarwal et al., 2020; Uehara et al., 2021) have been developed.
Revised: Based on the linear MDPs structure (Jin et al., 2020b), several theoretical algorithms (Agarwal et al., 2020; Uehara et al., 2021) have been developed.

Original: Ren et al.
Revised: Ren et al.

Original: (2022b);
Revised: (2022b);

Original: Zhang et al.
Revised: Zhang et al.

Original: ; Qiu et al.
Revised: ; Qiu et al.

Original: ; Ren et al.
Revised: ; Ren et al.

Original: (2022a) bridge the gap between theory and practice and bypass computational intractability via different techniques, demonstrating advantages empirically.
Revised: (2022a) bridge the gap between theory and practice and bypass computational intractability via different techniques, demonstrating advantages empirically.

Original: The proposed EPR is inspired from this class of representations, but extended to POMDPs, which is highly non-trivial.
Revised: The proposed EPR is inspired from this class of representations, but extended to POMDPs, which is highly non-trivial.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Our experiments investigate how our algorithm performs in robotic lomocation simulation environments.
Revised: Our experiments investigate how our algorithm performs in robotic lomocation simulation environments.

Original: We extensively evaluate the proposed approach on the Mojuco (Brockman et al., 2016) and
Revised: We extensively evaluate the proposed approach on the Mojuco (Brockman et al., 2016) and

Original: DeepMind Control Suites (Tassa et al., 2018).
Revised: DeepMind Control Suites (Tassa et al., 2018).

Original: We conduct experiments on both the fully observable
Revised: We conduct experiments on both the fully observable

Original: MDP and partially observable POMDP settings.
Revised: MDP and partially observable POMDP settings.

Original: 5. F ULLY O BSERVABLE MDP
Revised: 5. F ULLY O BSERVABLE MDP

Original: Dense-Reward Mujoco Tasks.
Revised: Dense-Reward Mujoco Tasks.

Original: We first conduct experiments in the fully observable MDP setting in Mujoco locomotion tasks.
Revised: We first conduct experiments in the fully observable MDP setting in Mujoco locomotion tasks.

Original: This is a test suite commonly used for both model-free and modelbased RL algorithms.
Revised: This is a test suite commonly used for both model-free and modelbased RL algorithms.

Original: We compare EPR with model-based RL baselines PETS (Chua et al., 2018) and ME-TRPO (Kurutach et al., 2018), and model-free RL baselines SAC (Haarnoja et al., 2018),
Revised: We compare EPR with model-based RL baselines PETS (Chua et al., 2018) and ME-TRPO (Kurutach et al., 2018), and model-free RL baselines SAC (Haarnoja et al., 2018),

Original: TRPO (Schulman et al., 2015) and PPO (Schulman et al., 2017).
Revised: TRPO (Schulman et al., 2015) and PPO (Schulman et al., 2017).

Original: In addition, we also compare to the representation learning RL baselines Deep Successor Feature (DeepSF)
Revised: In addition, we also compare to the representation learning RL baselines Deep Successor Feature (DeepSF)

Original: (Kulkarni et al., 2016) and SPEDE (Ren et al., 2022b).
Revised: (Kulkarni et al., 2016) and SPEDE (Ren et al., 2022b).

Original: We list the best model-based RL results (except for iLQR (Li &
Revised: We list the best model-based RL results (except for iLQR (Li &

Original: Todorov, 2004)) in MBBL (Wang et al., 2019) for comparison.
Revised: Todorov, 2004)) in MBBL (Wang et al., 2019) for comparison.

Original: All algorithms are run for 200K environment steps.
Revised: All algorithms are run for 200K environment steps.

Original: The results are averaged across four random seeds with a window size of 10K.
Revised: The results are averaged across four random seeds with a window size of 10K.

Original: We show that in Tab. 1, EPR significantly outperforms all the baselines including the strong previous
Revised: We show that in Tab. 1, EPR significantly outperforms all the baselines including the strong previous

Original: SoTA model-free algorithm SAC.
Revised: SoTA model-free algorithm SAC.

Original:  In particular, we observe that most model-based algorithms have a hard time learning the walk and hop behavior in the Walker and Hopper environments respectively.
Revised: [Table] [Table] In particular, we observe that most model-based algorithms have a hard time learning the walk and hop behavior in the Walker and Hopper environments respectively.

Original: We suspect that this is due to the fact that the quality of the data is bad at the initial data collection process (e.g., the agent often fall to the ground or has a hard time standing up).
Revised: We suspect that this is due to the fact that the quality of the data is bad at the initial data collection process (e.g., the agent often fall to the ground or has a hard time standing up).

Original: As a result, the behavior learned by most model-based[Table] [Table] algorithms can be suboptimal.
Revised: As a result, the behavior learned by most model-based algorithms can be suboptimal.

Original: For example, some model-based algorithms only learn to stand up without hopping in the Hopper environment.
Revised: For example, some model-based algorithms only learn to stand up without hopping in the Hopper environment.

Original: In contrast, EPR achieves SoTA performance in the
Revised: In contrast, EPR achieves SoTA performance in the

Original: Hopper task and Ant task, demonstrating the behavior of doing good exploration in the task domain.
Revised: Hopper task and Ant task, demonstrating the behavior of doing good exploration in the task domain.

Original: Sparse-Reward DM Control Tasks.
Revised: Sparse-Reward DM Control Tasks.

Original: Manually-designed dense reward functions are extremely hard to obtain, while it is difficult to gain access to a good dense reward function in practical real-robot settings.
Revised: Manually-designed dense reward functions are extremely hard to obtain, while it is difficult to gain access to a good dense reward function in practical real-robot settings.

Original: Thus, exploration in the sparse-reward settings is a key consideration for the success of
Revised: Thus, exploration in the sparse-reward settings is a key consideration for the success of

Original: RL in robotics settings.
Revised: RL in robotics settings.

Original: We test our algorithm EPR with SAC and PPO in such cases.
Revised: We test our algorithm EPR with SAC and PPO in such cases.

Original: Here we compare with DeepSF as an additional representation RL baseline.
Revised: Here we compare with DeepSF as an additional representation RL baseline.

Original: Note that the critic network used in SAC and PPO is deeper than EPR.
Revised: Note that the critic network used in SAC and PPO is deeper than EPR.

Original: From Tab. 2, we see that EPR achieves a particularly huge gain compared to SAC and PPO in sparse reward tasks walker-run-sparse .
Revised: From Tab. 2, we see that EPR achieves a particularly huge gain compared to SAC and PPO in sparse reward tasks walker-run-sparse .

Original: 5. P ARTIAL O BSERVABLE MDP COVERING V ELOCITY
Revised: 5. P ARTIAL O BSERVABLE MDP COVERING V ELOCITY

Original: Mujoco.
Revised: Mujoco.

Original: Often in practice, it is hard to recover a full observation of the states.
Revised: Often in practice, it is hard to recover a full observation of the states.

Original: Thus, the ability to handle a partially-observed MDP (POMDP) is also important if we can only recover partial observations.
Revised: Thus, the ability to handle a partially-observed MDP (POMDP) is also important if we can only recover partial observations.

Original: To conduct experiments in this setting, we mask the the velocities in the observations (replacing them by 0).
Revised: To conduct experiments in this setting, we mask the the velocities in the observations (replacing them by 0).

Original: We compare to algorithms with different embedding approaches that maps a given history sequence to a latent representation, which is used as the input for a SAC planner.
Revised: We compare to algorithms with different embedding approaches that maps a given history sequence to a latent representation, which is used as the input for a SAC planner.

Original: We consider four embedding methods as baselines:
Revised: We consider four embedding methods as baselines:

Original: Transformer (Trans), GRU, PSR (Guo et al., 2018), and finally a simple MLP baseline for sanity check, which concatenates the history sequence and directly maps that to a latent feature using a MLP.
Revised: Transformer (Trans), GRU, PSR (Guo et al., 2018), and finally a simple MLP baseline for sanity check, which concatenates the history sequence and directly maps that to a latent feature using a MLP.

Original: We find that this setting is very challenging and theperformance of all algorithms degrades comparied to the fully-observable setting. Nevertheless, the
Revised: We find that this setting is very challenging and the

Original: [Table] [Table]
Revised: [Table] [Table]

Revised: performance of all algorithms degrades comparied to the fully-observable setting.

Original:  proposed algorithm still achieves SoTA performance in tasks like Halfcheetah, Ant, SlimHumanoid.
Revised: Nevertheless, the proposed algorithm still achieves SoTA performance in tasks like Halfcheetah, Ant, SlimHumanoid.

Original: This demonstrates the capability of handling partial observability in EPR which can have an important effect in practice.
Revised: This demonstrates the capability of handling partial observability in EPR which can have an important effect in practice.

Original: DM Control Suite.
Revised: DM Control Suite.

Original: Correspondingly, we conduct POMDP experiments in the DM Control Suite.
Revised: Correspondingly, we conduct POMDP experiments in the DM Control Suite.

Original: However, we find that covering all the velocities is very challenging and thus we cover only the lastdimensions of the velocity.
Revised: However, we find that covering all the velocities is very challenging and thus we cover only the lastdimensions of the velocity.

Original: 5. I MAGE -B ASED E NVIRONMENTS
Revised: 5. I MAGE -B ASED E NVIRONMENTS

Original: To test the capability of our method on imagebased environments, we conduct an additional experiment on MetaWorld (Yu et al., 2020).
Revised: To test the capability of our method on imagebased environments, we conduct an additional experiment on MetaWorld (Yu et al., 2020).

Original: We choose one of the fetch-reach tasks and compare against the model-free algorithm
Revised: We choose one of the fetch-reach tasks and compare against the model-free algorithm

Original: SAC+AE (Yarats et al., 2021) and a popular representation learning method SPR (Schwarzer et al., 2020).
Revised: SAC+AE (Yarats et al., 2021) and a popular representation learning method SPR (Schwarzer et al., 2020).

Original: We show the results in Fig.
Revised: We show the results in Fig.

Original: 1 and note that the minimum distance between the current state and the goal is used as the evaluation metric
Revised: 1 and note that the minimum distance between the current state and the goal is used as the evaluation metric

Original: (the smaller distance means better performance).
Revised: (the smaller distance means better performance).

Original: We can see that EPR manages to reach the distant goal within 100K steps.
Revised: We can see that EPR manages to reach the distant goal within 100K steps.

Original: Comparing to SAC+AE,
Revised: Comparing to SAC+AE,

Original: EPR strictly dominate its performance.
Revised: EPR strictly dominate its performance.

Original: For SPR, although it learns faster at the beginning, EPR has better final performance.
Revised: For SPR, although it learns faster at the beginning, EPR has better final performance.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We exploit Energy-based Predictive Representation (EPR) for linearly representing value functions for arbitrary policies and supporting reinforcement learning in partially observed environments with finite memories.
Revised: We exploit Energy-based Predictive Representation (EPR) for linearly representing value functions for arbitrary policies and supporting reinforcement learning in partially observed environments with finite memories.

Original: The proposed EPR shows that planning and strategic exploration can be implemented efficiently.
Revised: The proposed EPR shows that planning and strategic exploration can be implemented efficiently.

Original: The coherent design of each component brings empirical advantages in RL benchmarks considering both the MDP and POMDP settings.
Revised: The coherent design of each component brings empirical advantages in RL benchmarks considering both the MDP and POMDP settings.

Original: Such superior performance makes the theoretical understanding of EPR more intriguing, which we leave as future work
Revised: Such superior performance makes the theoretical understanding of EPR more intriguing, which we leave as future work.

Original: A M ORE R ELATED W ORK
Revised: A M ORE R ELATED W ORK

Original: Provably RL for POMDPs.
Revised: Provably RL for POMDPs.

Original: Besides the statistical and computational hardness results for learning and planning upon POMDPs, most recent theoretical research focuses on overcoming the statistical complexity from the “curse of history” by considering tractable POMDPs (Krishnamurthy et al., 2016;
Revised: Besides the statistical and computational hardness results for learning and planning upon POMDPs, most recent theoretical research focuses on overcoming the statistical complexity from the “curse of history” by considering tractable POMDPs (Krishnamurthy et al., 2016;

Original: Azizzadenesheli et al., 2016; Guo et al., 2016; Jin et al., 2020a; Liu et al., 2022).
Revised: Azizzadenesheli et al., 2016; Guo et al., 2016; Jin et al., 2020a; Liu et al., 2022).

Original: Similarly to the observability structure we exploited in our algorithm, these work bypass the curse of history by different special structures, reducing the whole history dependency to finite-length memory.
Revised: Similarly to the observability structure we exploited in our algorithm, these work bypass the curse of history by different special structures, reducing the whole history dependency to finite-length memory.

Original: Recently, Uehara et al.
Revised: Recently, Uehara et al.

Original: ; Wang et al.
Revised: ; Wang et al.

Original: (2022) generalize these special structures with function approximation beyond tabular cases.
Revised: (2022) generalize these special structures with function approximation beyond tabular cases.

Original: Golowich et al.
Revised: Golowich et al.

Original: (2022a) consider the complexity planning and exploration together with learning, but only valid for tabular MDPs.
Revised: (2022a) consider the complexity planning and exploration together with learning, but only valid for tabular MDPs.

Revised: B D ERIVATION OF THE R ANDOM F EATURE IN EQUATION

Revised: We have that

Revised: [Equation]

Revised: [Equation]

Revised: where we have that

Original: B O BSERVABLE LQG AS EPR
Revised: [Equation] which concludes the proof for equation [Equation] C O BSERVABLE LQG AS EPR

Original: Follow the standard notations, the dynamics of Linear-Quadratic Gaussian is defined as
Revised: Follow the standard notations, the dynamics of Linear-Quadratic Gaussian is defined as

Original: [Equation]
Revised: [Equation]

Original: where w t and z t are Gaussian noise.
Revised: where w t and z t are Gaussian noise.

Original: Define the matrix
Revised: Define the matrix

Original: [Equation]
Revised: [Equation]

Original: and reduced observation
Revised: and reduced observation

Original: [Equation]
Revised: [Equation]

Original: By the observability condition of LQG, G L is full column rank, one can identify s 0 by
Revised: By the observability condition of LQG [Equation] G L is full column rank

Original: [Equation]
Revised: [Equation]

Original: Therefore, we have
Revised: one can identify s 0 by

Original: [Equation]
Revised: [Equation]

Original:  (cid:0) G ⊤ L G L (cid:1) −
Revised: Therefore, we have s 1 = As 0 + Ba 0 + w 0 =

Original: L (cid:88)
Revised: A 

Original: j = (cid:0) A ⊤ (cid:1) j − 1 C ⊤ ˜ o j 
Revised: [Equation] (cid:0) [Equation]

Original: 
Revised: 

Original: + Ba 1 + w 0 , s 

Revised: [Equation] 0 , s

Original: =
Revised: =

Original: As 1 + Ba 1 + w 1 = A (cid:0) G ⊤ L G L (cid:1) −
Revised: As 1 + Ba 1 + w 1 = A (cid:0) [Equation] −

Original: L (cid:88)
Revised: [Equation]

Original: j = (cid:0)
Revised: = [Equation]

Original: A ⊤ (cid:1) j − 1 C ⊤ ˜ o[-j   + ABa 
-]
Revised: o[-j   + ABa 
-]

Original: + Ba 2 + Aw 0 + w 1 , s L [-+
-]
Revised: [Equation] w 1 , s L [-+
-]+1 =

Original: =
Revised: As

Original: As L + Ba L + w L =
Revised: [Equation] Ba [Equation]

Original: A L  (cid:0) G ⊤ L G L (cid:1) −
Revised: A L  (cid:0) [Equation] (cid:1) −

Original: L (cid:88)
Revised: L (cid:88)

Original: j = (cid:0)
Revised: j = (cid:0)

Original: A ⊤ (cid:1) j − 1 C ⊤ ˜ o j 
Revised: A ⊤ (cid:1) j − 1 C ⊤ ˜ o j 

Original: 
Revised: 

Original: +
Revised: +

Original: L (cid:88)
Revised: L (cid:88)

Original: j =
Revised: j =

Original: A L − j Ba j +

Revised: A L − j Ba j +

Original: +
Revised: +

Original: L (cid:88)
Revised: L (cid:88)

Original: j =
Revised: j =

Original: A L − j w j , o L +1 =
Revised: A L − j w j , o L +1 =

Revised: Cs

Original: Cs t +1 +  zt =
Revised: [Equation] + [Equation] z

Original:  CA L  (cid:0) G ⊤ L G L (cid:1) −
Revised: [Equation] = CA L  (cid:0) [Equation]

Revised: ⊤

Revised: [Equation]

Original: L (cid:88)
Revised: L (cid:88)

Original: j = (cid:0)
Revised: j = (cid:0)

Original: A ⊤ (cid:1) j − 1 C ⊤ ˜ o j  
Revised: A ⊤ (cid:1) j − [Equation] C ⊤ ˜ o j  

Original:  + C
Revised: [Equation] C

Original: L (cid:88)
Revised: L (cid:88)

Original: j =
Revised: j =

Original: A L − j Ba j +1 + C
Revised: A L − j Ba j +1 + C

Original: L (cid:88)
Revised: L (cid:88)

Original: j =
Revised: j =

Original: A L − j w j + z t , which means o L +1 follows a Gaussian distribution with mean as a function of history x
Revised: A L − j w j + z t , which means o L +1 follows a Gaussian distribution with mean as a function of history x

Original: L = (cid:110)
Revised: L = (cid:110)

Original: ( o i − 1 , a i )
Revised: ( o i − 1 , a i )

Original: Li =1 (cid:111) and action a L +1 , and variance as a function of σ w , σ z , and ( A, B, C ) .
Revised: Li =1 (cid:111) and action a L +1 , and variance as a function of σ w , σ z , and ( A, B, C ) .

Original: Therefore, we have some function f A,B,C,σ w ,σ z and g A,B,C,σ w ,σ z , such that [Equation]
Revised: Therefore, we have some function f A,B,C,σ w ,σ z and g A,B,C,σ w ,σ z , such that g A,B,C,σ w ,σ z ( o L +1 ) =

Original:  A,B,C,σ w ,σ z ( x L , a L +1 ) + ξ, ξ ∼ N (0 , I ) .
Revised: f A,B,C,σ w ,σ z ( x L , a L +1 ) + ξ, ξ ∼ N (0 , I ) .

Original: On the other hand, we set λ = − 12 , and p ( o ) =
Revised: On the other hand, we set λ = − 12 , and p ( o ) =

Original: N (0 , I ) in (6), then, we obtain 
Revised: N (0 , I ) in (6), then, we obtain p ( o L +1 | x L , a L ) ∝ exp (cid:32)

Original: [Equation]
Revised: −∥

Original:  which reproduces the observable LQG with specific f A,B,C,σ w ,σ z and g A,B,C,σ w ,σ z .
Revised: g ( o L +1 ) − f ( x L , a L ) ∥ (cid:33) , which reproduces the observable LQG with specific f A,B,C,σ w ,σ z and g A,B,C,σ w ,σ z .

Original: C E XPERIMENT D ETAILS C. O NLINE S ETTING
Revised: D E XPERIMENT D ETAILS D. O NLINE S ETTING

Original: In Table [Table] we list all the hyperparameters and network architecture we use for our experiments.
Revised: In Table 9, we list all the hyperparameters and network architecture we use for our experiments.

Original: We see that we don’t use the additional exploration bonus term in the mojuco tasks.
Revised: We see that we don’t use the additional exploration bonus term in the mojuco tasks.

Original: But this is very helpful in DM control suite tasks, especially in those sparse-reward tasks.
Revised: But this is very helpful in DM control suite tasks, especially in those sparse-reward tasks.

Original: For evaluation in Mujoco, in each evaluation (every 5K steps) we test our algorithm for 10 episodes.
Revised: For evaluation in Mujoco, in each evaluation (every 5K steps) we test our algorithm for 10 episodes.

Original: We average the results over the last 4 evaluations and 4 random seeds.
Revised: We average the results over the last 4 evaluations and 4 random seeds.

Original: For Dreamer and Proto-RL, we change their network from CNN to 3-layer MLP and disable the image data augmentation part (since we test on the state space).
Revised: For Dreamer and Proto-RL, we change their network from CNN to 3-layer MLP and disable the image data augmentation part (since we test on the state space).

Revised: The architecture we used for the transformer is following the Trajectory

Revised: Transformer (Janner et al., 2021).

Revised: The attention used is the causal attention.

Original: We tried to tune some of their hyperparameter (e.g., exploration steps in Proto-RL) and report the best number across our runs.
Revised: We tried to tune some of their hyperparameter (e.g., exploration steps in Proto-RL) and report the best number across our runs.

Original: However, due to the short time, it is also possible that we didn’t tune the hyperparameter enough
Revised: However, due to the short time, it is also possible that we didn’t tune the hyperparameter enough.

Original: [Table]

Original: [Table] L EARNING C 
Revised: D. L EARNING C URVES

Original: [Table] We provide the performance curves for online DM Control Suite experiments in Figure 2.
Revised: We provide the performance curves for online DM Control Suite experiments in Figure 2.

Original: As we can see in the figures, the proposed EPR converges faster and achieve the state-of-the-art performances in most of the environments, demonstrating the sample efficiency and the ability to balance of exploration vs. exploitation of EPR.
Revised: As we can see in the figures, the proposed EPR converges faster and achieve the state-of-the-art performances in most of the environments, demonstrating the sample efficiency and the ability to balance of exploration vs. exploitation of EPR.

Original: We also provide additional curves for POMDP setting in
Revised: We also provide additional curves for POMDP setting in

Original: Figure 3.
Revised: Figure 3.

Original: C. I MAGE -B ASED E XPERIMENTS
Revised: D. I MAGE -B ASED E XPERIMENTS

Original: We provide the details of metaworld image-based experiments here.
Revised: We provide the details of metaworld image-based experiments here.

Original: We first provide an illustration of the reach environment in Figure 4.
Revised: We first provide an illustration of the reach environment in Figure 4.

Original: We then provide some more experiment details in the following section.
Revised: We then provide some more experiment details in the following section.

Original: [Table]
Revised: [Table]

Original: [Table] 
Revised: [Table] [Table] [Figure] [Table]

Original:  [Table]
Revised: [Figure] [Table] [Figure] [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Original: [Table] 
Revised: [Table] [Figure] [Table] [Table]