Energy-based Predictive Representation for Reinforcement Learning
Authors: Tianjun Zhang
Tongzheng Ren
Chenjun Xiao
Wenli Xiao
Joseph E. Gonzalez
Dale Schuurmans
Bo Dai
Venue: Submitted to ICLR 2023 (2023)
Keywords: Energy-based Models
Predictive State Representation
Partially Observable Markov Decision Process
Reinforcement Learning

Abstract: In real world applications, it is usually necessary for a reinforcement learning algorithm to handle the partial observability beyond Markov decision processes (MDPs). Although the partially observable Markov decision process (POMDP) has been precisely motivated for this requirement, such a formulation raises significant computational and statistical hardness challenges in learning and planning. In this work, we introduce the Energy-based Predictive Representation (EPR), which leads to a unified framework for practical reinforcement learning algorithm design in both MDPs and POMDPs settings, to handle the learning, exploration, and planning in a coherent way. The proposed approach relies on the powerful neural energy-based model to extract sufficient representation, from which Q-functions can be efficiently approximated. With such a representation, we develop an efficient approach for computing confidence, which allows optimism/pessimism in the face of uncertainty to be efficiently implemented in planning, hence managing the exploration versus exploitation tradeoff. An experimental investigation shows that the proposed algorithm can surpass state-of-the-art performance in both MDP and POMDP settings in comparison to existing baselines.

Under review as a conference paper at ICLR 2023
ENERGY-BASED PREDICTIVE REPRESENTATIONS FOR PARTIALLY OBSERVED REINFORCEMENT LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
In real world applications, usually necessary for reinforcement learning algorithm to handle partial observability, which is not captured by Markov decision process (MDP). Although partially observable Markov decision pro- cesses (POMDPs) have been precisely motivated by this requirement, they raise significant computational and statistical hardness challenges in learning and plan- ning. In this work, we introduce the Energy-hased 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. The proposed framework relies on a powerful neural energy- based model to extract sufficient representation, from which Q-functions can be efficiently approximated. With such a representation, confidence can be efficiently computed allow optimism/ppssimism in the face of uncertainty to be efficiently implemented in planning, enabling effective manage- ment of the exploration versus exploitation tradeoff. An experimental investigation shows that the proposed algorithm can surpass tfte-oofthhe-rr performance in both MDP and POMDP settings in comparison to existing baselines.
INTRODUCTION
Reinforcement learning (RL) based on Markov Decision Processes (MDPs) has proved be extremely effective in several real world decision-making problems (Levine et al., 2016; Jiang et al., 2021), 1). However, the success of most RI. algorithms (Ren et al., 2022b; Zhang et al., 2022) relies heavily on the assumption that the environment state is fully observable to the agent. In practice, such an assumption can be easily violated in the presence of observational noise. To address this issue, Partially Observable Markov Decision Processes (POMDPs) (Aström, 1965) have been proposed for capturing the inherent uncertainty about the state arising from partial observations.
However, the flexibility of POMDPs creates significant statistical and computational hardness in terms of planning, exploration and learning. In particular, i), partial observability induces a non- Markovian dependence over the entire history; and ii), the expanded spaces of observation sequences or state space distributions incur significant representation challenges. 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). These complexities only become more demanding in continuous state spaces and real-world scenarios.
On the other hand, despite the theoretical hardness, the widely used sliding window policy parameter- ization 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. Recently, observable POMDPs with invertible emissions have been investigated to justify the sliding window heuristic in tabular cases (Azizzadenesheli et al., 2016; 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; Uchara et al., 2022). 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. In this paper, we consider the following natural question:
How can one design efficient and practical algorithms for structured POMDPs?

---

Under review as a conference paper at ICLR 2023
In particular, would like to exploit special structures that allows approximation to bypass inherent worst-case difficulties. By "efficient" we mean considering learning, planning and exploration in a unified manner that can balance emrоrs in each component and reduce unnecessary computation, while by "practical" we mean the algorithm retains sufficient flexibility and can he easily implemented and deployed in real-world scenarios.
There have been many attempts to address this question. The most straightforward idea is to extend model-free RL methods, including policy gradient and Q-learning, with memory-limited parametrization, e.g,, recurrent neural networks (Wierstra et al., 2007; Hausknecht & Stone, 2015; Zhu at. 2017). 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. Following this idea, Deisenroth & Peters (2012) and Igl et al. (2018); Gregor et al. (2019); Zhang et al. (2019): Lee et al. (2020) consider Gaussian process or deep model parametrizations, respectively. Such methods are designed based on implicit assumptions about structure through the parameterization choices of the models. 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.
As alternative, spectral representation approaches provide alternative strategy based on ex- tracting sufficient representation that can support learning, planning and exploration. In this vein Azizzadenesheli al. (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. Predictive State Representations (PSR) (Littman & Sutton, 2001; Boots et al., 2011) I) also leverage spectral decomposition, but instead of recovering an underlying latent variable model, they learn an equivalent sufficient representation of belief. 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). However, efficient exploration and tractable planning with spectral representations has yet to be thoroughly developed (Zhan et al., 2022).
In this paper, propose Energy-based based Predictive Representation (EPR) to support efficient and tractable learning, planning, and exploration in POMDPs (and MDPs), as a solution to the aforemen- tioned question. More specifically:
We propose flexible nonlinear energy-based model for induced belief-state MDPs without explicit parameterization of beliefs, providing principled linear sufficient representation for the state-action value function.
We reveal the connection between EPR and PSR, while also illustrating the differences, to demonstrate the modeling ability of the proposed EPR.
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.
We provide computationalll--raaaao
We conduct comprehensive comparison to existing state-of-the-arr RI algorithms in both MDP and POMDP benchmarks, demonstrating superior empirical performance of the proposed EPR.
2 PRELIMINARIES
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.
Partially Observable Markoy Decision Processes. Formally, we define a partially observable Markov decision process (POMDP)
tuple P = AS H. μ., P, O). where H is positive integer denoting the length of horizon; μ is the initial distribution of state, r : x A - [0, 1]. the reward function, and S, .A, O denote the state, action and observation space, respectively. p(- x A + ASS) is the transition kernel, capturing the dynamics between states, and ((+1s) : S (OO) is the emission kernel, where A (-) denotes the set of probability measures over the support.
Initially, given a state $1 (ss) as a starting point, at each step h E [1, H], the agent takes an action E A, a new state 8 is generated $1 P(: ((-18h; an), and the agent observes 01+1 O(118441)
2

---

Under review as a conference paper at ICLR 2023
and reward r(5h+1, ah+1). Due to partial observability,, the dependence between observations is Non-Markovian, hence, we define a policy x = {Te} where TT O (Ax + (AA) to depend on the whole history, i.c., He = foo, {a,, 0+41) ) The corresponding value for policy я can be defined as V
Σ h=1 r(sh, an) and the objective is to find the optimal policy Nr" = arg max; V
Markov decision processes (MDPs) are a degenerate case of POMDPs, where S = O and O(o]s) = 6(0 = s). and can be specified as M = (S, A, r, H, μ, P). One can also convert a POMDP to an MDP by treating the whole history zt = foo. {ai, 0++1 o)) as the state. Specifically, following (Kaelbling et al., 1998), we define the belief b: O x (A x O)² A(S), VI Nt, which can be recursively defined as: ((81(01) = ((s1|01), and
b(51+1(2++1) = (bsss//e) (14++ dor+1"
(1)
Each entry of the belief state describes the probability of the underlying state given the past history. Furthermore, with slight abuse of notation, we use b, denote the belief state step Then, one can construct the equivalent belief MDP M₁ = (X, RR. H, pb. Tb) with X' denoting the set of
possible histories, and
p4b := - b(s)o)))((o1)da R((b,a) = ff ( (((s))((ss,))dds
To (be+1 [b,, 0r) = / 10 P(oc+ lb,, at) doc+1.
(3)
Therefore, the corresponding value function V (bb) and Q7 (bh,,a for the belief MDP given a policy ι can be defined as:
(be,,, (bn) = EE Phe (e) (h, an) = E Σ (be ax) an
Following the MDP perspective, we also have the Bellman recursive equation:
VØ (b,) E,, 1Q% ax)], @%, a.)) = Ry a:.)) Er, (
(4)
One can still apply a dynamic programming style approach 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. To combat with these essential difficulties, we will leverage two particular structures, observability and linearity, as introduced below.
Observability in POMDPs. It has been shown (Even-Dar et al., 2007; Golowich et al., 2022b; Uchara 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. horizon length (Golowich et al., 2022a;b). Specifically, the observability property for POMDPs is defined as follows.
Assumption 1 ((Even-Daree et al., 2007; Golowich et al., 2022b)). The POMDP with emission model O satisfies o-bbervabilityy arbitrary beliefs and b' over states, (1(O,) - (0, > > 1/6 1' where b) f O(o is ()sss)ss.
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 finite belief history, which eliminates the exponential complexity induced by full history dependence. Specifically,
we denote L as the length of the window. Then, defining =t oo r., {ai, 0:4 1 } E X24,
the approximated beliefs b² follow the same recursive definition as (1) but with only finite history I. starting from the uniform belief. This immediately induces an approximate MDP Mí (XLL_ A, R/, H. ub, T according to (3) with b², instead of b. Theorem 2.1 in Golowich et al. (2022a) proves that the approximation error of the finite- memory belief MDP is small for observable POMDPs. Hence, with slight abuse of notation, we still use b to represent b² throughout the paper.
Linearity in MDPs. 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:
P(s'|s, a) = (ø(s, a), p(s')), r(s, a) = (ó(s, a).8).
(5)
where o : S A + H. μ : S > H are two feature maps to a Hilbert space H. Under such an assumption, we can represent the state-action value function QF for an arbitrary policy n by:
Q" r(s,a) + + V (s') a)ds' (d(s, 0 + 7 V (d))(())dds))
3

---

Under review as a conference paper at ICLR 2023
which implies that instead of 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 Ø. 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. 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. Energy-based Models are one of the most flexible models to represent the conditional probability measure. It takes the form of p(y]z) exp( ( fxx, (2)/ where f(r, y). which can be parametrized by deep models, is the energy of (zc, y) and (xx) a partition function that only depends on to guarantee p(y(ar) is valid probability measure. When y is discrete, we have that p(y[z) = exp(- f(x, /))/ E, exp(-f y)), which corresponds to the standard softmax probability where fz, y) is the softmax logits. We refer the interested readers to Song & Kingma (2021) for the training methods of energy-based models.
ENERGY-BASED PREDICTIVE REPRESENTATION
We propose Energy-base based Predictive Representation (EPR), which introduces linearity into finite- history approximated POMDPs, allowing the complexity induced by large state spaces and long histories to be overcome, yielding improved efficiency for learning, planning and exploration. We emphasize that the proposed method is also applicable to MDPs.
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'|a [s, a) = (ø(s, a), uu(')); ()) to obtain an efficient and practical framework for learning, planning and exploration. Recall the construction of finite- memory belief MDP to approximate a POMDP discussed in Section 2. which avoids full history dependence. For such a constructed belief MDP, natural idea is apply linear MDP algorithms, i.e., extracting the linear decomposition for Th (b'(b,, a) = (ø(b, a), n(b'),, to handle the hardnesses of POMDPs mentioned in Section 1. However, there are several difficulties in such straightforward extension:
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 is difficult.
These difficulties hinder the extension of linear MDPs to observable POMDPs. However, note that never explicitly require the beliefs and their dynamics, but only the representation Φ(b, a). As beliefs are functions over wnidew histories, the representation can also be rewritten di at), which suggests that one might bypass the inherent difficulties by a reprameterization trick. Consider the energy-basee parametrization (Arbel al., 2020) for P(o++1]((z)) ae) where b(xz) is the belief for history Ze:
P(ov+11((z)). a.) p(o exp (f ar) (g(o:+1) + (((1))a))))
(6)
Bor+s (exp (f a (g(or+1) + (z),, = 1, (2,, e XL,
(7)
where A is scalar, p(o) is fixed distribution and the normalization condition enforces that the energy- based model (000+116(x))).).))) a;) is a valid distribution. We avoid any explicit parametrization and computation of beliefs h, while preserving dependence through / and g, which will be learned jointly. 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. As a special case, we note that the observable Linear- Quudraat Gaussian (LQG) actually follows (6) with a specific À and p(o). See Appendix C for details.
Meanwhile, this approach also provides a linear factorization of T(be+ a:) almost for free. By viewing the proposed parameterization (6) as kernel and following the random Fourier feature trick (Rahimi & Recht, 2007: Choromanski et al., 2020; Ren et al., 2022b), one can write
P(ot+1 az) = En [dw az) (0+1),
(8)
where wi ~ N/0, Ta) and
(())) = exp ((~ [ (1te + wi f(re,
im
(((++1)= = p(ott1) exp (wt g(o++1)= |1a(04+1))
(10)
2
4

---

Under review as a conference paper at ICLR 2023
which can be derived from the softmax random feature from Choromanski et al. (2020). We also provide a derivation in Appendix B. Substituting (8) into (3) yields the factorization T3 as
To (be+1lbe bb, ar) = =So 1 (((l Eas (o =lE (I ae)u(bb+1)),
do
(11)
(12)
where (u(b+1) := Jo 1 (e o Vhe (0t+1))
do With this formula, we obtain a valid linear representation o(b,, 0t) as an Energy-based Predictive Representation (EPR) for a belief MDP without any explicit beliefs.
To learn the EPR given data D (1, exploit maximum likelihood estima- tion (MLE) of (6),
min fi
-ED [f a1) (g(ot+1) (f(r,
(13)
S.L.,
Bp(or+1) [exp (f( at) (g(o++1) ((z))))))) = 1, E xx.
(14)
To ensure the constraints, we add regularization term
(log (E, [exp([(x,, a (g(0) + Xf(x.,
(15)
log 1
a.) (g(o.) + Xf(x.,
(16)
2=1
with or p. The objective will be
mins. ED
[(1t, az) (glor+1) af))
- m
m Σ exp((((xr, ar) (g(o,) + A Af(xe,
)
+
log
(17)
In practice, we can further simplify the objective by normalizing the f(x,, a:) = (anaa) ing the final objective
TF( obtain-
minj. ED -i a.) g((o)+)) + + a (log EE exp(i(sr, at) g(o;)) + 2) 27
(18)
which reduces to contrastive loss that can be optimized by stochastic gradient descent with a deep network parameterization of ʃ and We obtain negative samples {o:) ~ p(o) from mixture of replay buffer and collected trajectories.
Before we introduce an exploration-explottaiie mechanism with EPR in Section 3.1, we first discuss the relationship between the proposed EPR, predictive state representations (PSR) (Littman & Sutton, 2001; Singh et al., 2004), and spectral dynamics embedding (SPEDE) (Ren et al., 2022b).
Connection to PSR (Littman & Sutton, 2001; Singh al., 2004): The predictive state rep- resentation (PSR) was also proposed for bypassing belief calculation by factorizing a variant of the transition dynamics operator. Specifically, given the history (Ir,, az), the probability for ob- serving test, i.e., the finite sequence of events y (0++1, 0++1, 0r+k; 0+-)) (t+k; with k E N, is p(y|z) p(o ( + For time step one can construct set of core tests U = {ui,..., up. as sufficient statistics for history It, such that for any test т, p(T[2z) = (p(U) [1)) wx) for some weights Wy. E R!üi The forward dynamics can be represented
a PSR by Bayes' rule:
p(T Fe, On 1+) = ( which implies that a PSR updates with new observations and "(arioe)1) P(U ((U)=)))
actions by repeating calculation for each u; E U. Although originally defined for tabular cases, 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).
Obviously, the proposed EPR shares similarities with PSR. Both factorize conditional distributions defined by the dynamics. However, these representations are designed for different purposes, and thus, with different usages and updates. Concretely, EPR is proposed for seeking a linear space that can represent the funntion.. The representation is designed to preserve linearity with successive observations without the need for Bayesian updates, which induce extra nonlinearity in PSRs. This
5

---

Under review as a conference paper at ICLR 2023
Algorithm 1 Energy-based Predictive Representation
1: Input: History Embedding f(x, a), Observation Embedding g Random Feature {wi ) where wi ~ (0, Ia), Initial Random Policy To. Initial Dataset D @ for online setting.
2: for Episode 1. K do
3: Collect data { di.j. Ouj rig)) with i, - + (1 TT, and add the data to D.
4: Optimize f and with (18) using the data from D.
5: Obtain the representation (bI2,, (z) via (9) using [w:)
6: Add the bonus (19) to the reward and obtain the optimal policy 4++1 with the (T o:)
parameterize as ofIe, ae) and optimize via FQI.
7:
(Optional) Extract policy by soft-AC from learned Q.
8: end for
9: Return
linear property leads to efficient exploration and planning in EPR; while an efficient exploration and planning algorithm has not yet been discussed for PSR.
Connection to SPEDE (Ren et al., 2022b): Linear random features have been proposed for solving planning in MDPs with nonlinear dynamies in (Ren et al., 2022b), where the transi- tion operator is defined as T (s' a) exp - f(s, a)ll 1012/(22)) /(02)) corresponding to dy- namics s' = [(s a) + with Gaussian noise ~ .N o0, In addition to the generaliza- tion of EPR for POMDPs, even in MDP, EPR considers general energy-based model, T(s'ls, a) x p(s") exp (f(s, a) T (g(s')+ Af(s. for the dynamics, which is far more flexible than the Gaussian perturbation model considered SPEDE.
3.1 ONLINE EXPLORATION AND OFFLINE POLICY OPTIMIZATION WITH EPR
With an EPR o(I,, az) learned for 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 optimisn/pessimismm in the face of uncertainty.
Exploration and Exploitation with Elliptical Confidence Bound. Given the learned repre- sentation o(x,, ot), 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; Uchara et al., 2021). This achieved simply by adding an additional elliptical bonus to the R(z, a). Specifically, given the dataset we collect D ( ; a, R2 11+1) and calculate the confidence bound as the bonus,
b(I,, az) = V a,) E 10,,,1r a+)
(19)
where À is a pre-specified hyperparameter, and Σ,, = Ei= AT + alat 0z az)o(at.a)) One can then implement UCB/LCB by adding/subtracting the bonus to R(x,; at), and performing planning on the modified reward function.
Planning with Obtained Representation. Planning can be conducted by Bellman recursion within the linear space spanned by oa ,a)) without a bonus. However, with an additional bonus term, the Q" no longer lies in the linear space of Φ, since
Q" (zt: at) = R(z az) + b(a);, - -r, IQ" (I1+1, 1-1)
As discussed in (Zhang et al., 2022). one can augment the feature space v (x, a) == {0(z, a), b(x, a)} to ensure the f-fuuctions can be linearly represented but with an extra O (d) memory cost. In prac- tice, perform fitted Q iteration with nonlinear component extending the linear parameterization, i.e. Q(x, a) {wn, w²] (Ø(2, a). a (w:] (xx, a))]
We provide an outline of our implementation of UCB in Algorithm 1. LCB for pessimistic offline RL is similar but using pre-collected dataset D without data collection iteration in Step 2, and with the bonus subtracted Step 6. 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. Then we perform representation learning and optimistic planning with the Q function parameterized upon the learned representation. Finally, also extract policy from the learned Q by soft actor-critic critic (Haarnoja et al., 2018).
4 RELATED WORK
Partial Observability in Reinforcement Learning. Despite the essential hardness of POMDPs in terms of learning, planning and exploration (Papadimitriou & Tsitsiklis, 1987; Madani al., 1998;

---

Under review as a conference paper at ICLR 2023
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.
Algorithmically, model-based/-free algorithms have been extended to POMDPs, explicitly implic- itly exploiting structure. Model-based RL algorithms parameterize and learn latent dynamics with an emission model explicitly, and planning through the simulation upon the learned models. 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 al., 2019a;b; 2020). 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. On the other hand, model-free RL algorithms have been extended for POMDPs by learning history dependent value functions and/or policies, through temporal-differee
algorithms or policy gradients. For example, deep Q-learning (Mnih et 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; et al., 2017). Recurrent neural networks have also been exploited for history dependent policies (Schmidhuber, 1990; Bakker, 2001; Wierstra et al., 2007; Heess et al., 2015) in policy gradient algorithms as well as actor-critic approaches (Ni et al., 2021; Meng et al., 2021). Model-free RL for POMDPs bypasses the planning complexity of model- difficulty in exploration remains, which leads to suboptimal performance in practice. 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.
RL algorithmss However, the
Representation Learning for RL. Successful vision-based representation learning methods have been extended to RL for extracting compact and invariant state-only information from raw- ruw-pxxels, e.g., (Laskin et al., 2020a;b; Kostrikov et al., 2020). However, such vision-based features not specially designed for capturing properties in POMDPs/MDPs essential for decision making. To reveal structure that is particularly helpful for RI,, many representation learning methods have been designed for different purposes, such as bi-simulation (Ferns et al., 2004; Gelada et al., 2019; 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; Yang et al., 2021). These representation methods ignore the requirement of planning tractability. Moreover, they are learning based on pre-collected dataset, which ignores the exploration issue.
Features that are able to represent value functions are desirable for efficient planning and exploration. Based on the linear MDPs structure (Jin et al., 2020b), several theoretical algorithms (Agarwal et al., 2020; Uehara et al., 2021) have been developed. Ren et al. (2022b); Zhang et al. (2022); Qiu et al. (2022); Ren et al. (2022a) bridge the gap between theory and practice and bypass computational intractability via different techniques, demonstrating advantages empirically. The proposed EPR is inspired from this class of representations, but extended to POMDPs, which is highly non-trivial.
EXPERIMENTS
Our experiments investigate how our algorithm performs in robotic lomocation simulation environ- ments. We extensively evaluate the proposed approach on the Mojuco (Brockman et al., 2016) and DeepMind Control Suites (Tassa et al., 2018). We conduct experiments on both the fully observable MDP and partially observable POMDP settings.
5.1 FULLY OBSERVABLE MDP
Dense-Reward Mujoco Tasks. We first conduct experiments in the fully observable MDP setting in Mujoco locomotion tasks. This is a test suite commonly used for both model-free and model- based RL algorithms. We compare EPR with model-based RL basclines PETS (Chua et al., 2018) and ME-TRPO (Kurutach et al., 2018), and model-free RL. baselines SAC (Haarnoja al., 2018), TRPO (Schulman et al., 2015) and PPO (Schulman et al., 2017). In addition, we also compare to the representation learning RL baselines Deep Successor Feature (DeepSF) (Kulkarni et al., 2016) and SPEDE (Ren et al., 2022b). We list the best model-based RL results (except for iLQR (Li & Todorov, 2004)) in MBBL (Wang et al., 2019) for comparison. All algorithms are run for 200K environment steps. The results are averaged across four random seeds with window size of 10K. We show that in Tab. 1, EPR significantly outperforms all the baselines including the strong previous SoTA model-free algorithm SAC.
7

---

Under review as a conference paper at ICLR 2023
Table I: Performance various MuloCo control tasks. All the results are averaged across 4 random seeds and a window size of 10K. Results marked with adopted from MBBL.. EPR achieves strong performance compared with baselines.
HaltChoctah
Reacher
Humanoid-ET
Pendulum
I-Peodulum
ME-TRPO"
2283,74:000
-13.445.2
72.9=8.9
177.3
8666
Model- RL.
PETS RS"
966.91477.6
400.1...9
102.6.1022.6
167.91.35.8
12.11125.1
PETS-CEM'
2799.3+1899.9
-12.345.22
110.8:990.1
167.44::3.0
-20.54:28.9
Best MBBL
3639.0=t:135.88
-.1120.1
1577.0=15044
177.3+199
0.0=0.0
PPO"
17,2484.4
-17.24:0.9
451.44499.1
163.448.00
.40.84221.00
Model-Froe RL
TRPO
12.0:485.5
-10.140.6
289.8:45.2
166.747.3
-27.6++15.8
SAC" (3-layer)
4000.74-022.1
-6.420.5
4794
168.2499
-0,2+0011
DoepSF SPEDE
1180,4
-16.843.6
168.6+5.11
168.6+5..1
-0.2+0.33
42.6,3+192.6
-7,221111
886.9:195.2
169.5::0.6
0.0=0.0
Representation RL
5107.6.1.944
5.611..
131.611311.3
169.414.11
0.0
Ant-ET
Hopper-ET ET
S-HHmannid--E
CariPole
Walker-ET
ME-TRFO'
42.6:221.1
1272.5+::00.9
-154.9:33443
160.1::69.
16009.3
PETS-RS
130.0 148.1
205.8+36.5
120.2+1822
195 0-2-2800
312.5=493.44
Model-Based RL
PETS-CEM"
81.6+145.88
129.3+33600
15571-1
195 5+:300
260.2=536.9
Best MBBL
275.41:309.1
1272.5+::00.9
1084.3+77.0
200.0:1:.00
312.5=493.44
PPO TRPO
80.1
758.0=62.0
454.34+36.7
86.5=188
306.1=17.22
Model-Five RL
116.81:47.33
337.4:133,5
281.3:1:099
47,3=15.7
229.5
SAC" (3-layer)
2012.7:1.71.3
1815.51165..1
834.6.1.11.11
199.4100.4
2216.4.1668.7
DeepSI
768.1+:44.1
548.9:2253.3
533.8::154.9
194.5+:5
165.6 127.9
SPEDE
806.2.1.00.
732.2.1.263.9
986.4-1154.7
138.2:1.39.5
501.61204.0
Representation RI.
EPR
4081,3:1.73.99
2191.41.502.8
1326.1-1.0.8
200.8:10.11
7975,41751.3
Table 2:
Performance on various Deepmind Suite Control tasks. All the results are averaged across four random seeds and window size 10K. Comparing with SAC, our methed achieves even better performance on sparse-reward tasks.
cheetah.rum
walker.run
walker.
humancid.
Model-Based RL
Dreamer
542.0 + 27.7
799.9
337.7:267.2
95.4=54.77
1.0:0.
PPO
227.7+557.9
5.4
$1.641155
0.0+:0.0
11100.0
Model-Free RL
SAC (2-layer)
222.2++-1.0
32.4+:27.88
183.0::2..4
$3.5+699.3
1.3+0.1
SAC (3-layer)
595.2:1:9.00
419.5:73.3
700.9:336.6
311,5+36144
12400.1
DocpSF
295.3::43.5
0.0:00.0
27.9:2222
0.1+:0.1
0.940.1
Representation RL Protn RL.
305.5:1.7.99
0.0:10.0
433.51156.8
46.9113.11
0.340.66
EPR
611.6::3.5
36..8-30066
79288:35.7
701.:30044
11.5+:544
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. We suspect that this is due to the fact that the quality of the data is bad at the initial data collection process (c.g., the agent often fall to the ground or has hard time standing up). As result, the behavior learned by most model- based algorithms can be suboptimal. For example, some model-based algorithms only learn to stand up without hopping in the Hopper environment. In contrast, EPR achieves SoTA performance in the Hopper task and Ant task, demonstrating the behavior of doing good exploration in the task domain.
Sparse-Reward DM Control Tasks. Manually-designee dense reward functions are extremely hard to obtain, while it is difficult to gain access to good dense reward function in practical real-robot settings. Thus, exploration in the sparse-reward settings is a key consideration for the success of RL in robotics settings. We test algorithm EPR with SAC and PPO in such cases. Here compare with DeepSF as an additional representation RL baseline. Note that the critic network used in SAC and PPO is deeper than EPR. From Tab. 2, we see that EPR achieves a particularly huge gain compared to SAC and PPO in sparse reward tasks wa ker- run--ruunsp...e....
5.2 PARTIAL OBSERVABLE MDP COVERING VELOCITY
Mujoco. Often in practice, it is hard to recover a full observation of the states. Thus, the ability to handle a partially-observer observations. To conduct experiments this setting, we mask the the velocities in the observations (replacing them by 0). We compare to algorithms with different embedding approaches that maps given history sequence to a latent representation, which used as the input for a SAC planner. We consider four embedding methods baselines: 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 latent feature using ML.. We find that this setting is very challenging and the
MDP (POMDP) is also important if we can only recover partial
8

---

Under review as a conference paper at ICLR 2023
Table 3: Performance various MuloCo control tasks. All the results are averaged random seeds and a window size of 10K. Results marked with adopted from MBBL.. EPR achieves strong performance compared with baselines.
HalfChectah
Huminoid-ET
Walker-ET
EPR
3441,6 + 993.0
865.6
107.3
416.6 + 145.6 862.4 1 355.3
Representation RL
PSR
2679.75 1: 386
534.4 36.6
MLP
1612.0 +: 223
614.15 +: 67.6
236.5 +: 65.6 388.7
Model Free RL
1: 227.2
GRU
1664.3 = 431.2
467.7 43.4
1020.6 + 364.9
Ant-ET
Hopper-FT
Hurmanid--FT
EPR
1508.1 594.4
1059.4 1: 582.5
805.7 :: 65.9
Representation RL
166.6
818.8
65.2
MLP
1262.0 = 68.7
260.6 1: 63.7
294.17 49.4
Free RL
Trans
928.5 44.
470.8 1: 50.3
447.9 112.6
GRU
1190.8 = 79.4
777.5 + 113.3
485.9 = 25.8
Table 4: Performance of on various Deepmind Suite Control tasks. All the results averaged across four random seeds and window size of 10K. Comparing with SAC, method achieves even better performance on aparse-reward tasks.
chectah.run
chettahrun.sparse
walkerrrn
walker.run.sperse
humanoid.run
Explicir Model
SLAC
30.1
+ 3.4
0.0 + 0.0
MI.P
743.3 1: 7.2
0.0 1: 0.0
279.8 +: 190.6
0.0 1: 0.0
1.2 0.1
Model-Free RL
Trans
379.6 - 80
0.0
68.06 + 39.9
0.0 + 0.0
0.92
173.7 25.7
$7.4 :: 7.4 509.8:124.4
0.0 1:
0.89
Representation RL
EPR
526,5:161.1
411.0::51.6
460.3-1:1.66
6.1
performance of all algorithms degrades comparied the fully-observable setting. Nevertheless, the proposed algorithm still achieves SoTA performance in tasks like Halfcheetah, Ant, SlimHumanoid. This demonstrates the capability of handling partial observability in EPR which can have an important effect in practice.
DM Control Suite. Correspondingly, we conduct POMDP experiments in the DM Control Suite. However, we find that covering all the velocities is very challenging and thus we cover only the last dimensions of the velocity.
5.3 IMAGE-BASED ENVIRONMENTS
To test the capability of our method on image- based environments, we conduct an additional experiment on MetaWorld (Yu et al., 2020).
Fetch-Reach-lmage
(ours)
soc a2
SPR
We choose one of the fetch-reach tasks and compare against the model-free algorithm SAC+AE (Yarats et al., 2021) and popular repre- sentation learning method SPR (Schwarzer et al., 2020). We show the results in Fig. and note that the minimum distance between the current state and the goal is used the evaluation metric (the smaller distance means better performance). We can see that EPR manages to reach the distant goal within 100K steps. Comparing to SAC+AE, EPR strictly dominate its performance. For SPR, although it learns faster at the beginning, EPR has better final performance.
Ninimum distance
0.15
0.05
1.0
Erviromment
Figure 1: EPR in image-based environment: EPR gets a good performance compared to all baselines (e.g. SPR and SAC+AF).
6 CONCLUSION
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. The proposed EPR shows that planning and strategic exploration can be implemented efficiently. The coherent design of each component brings empirical advantages in RL benchmarks considering both the MDP and POMDP settings. Such superior performance makes the theoretical understanding of EPR more intriguing, which we leave future work.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Yasin Abbasi-Yadkori, Dávid Pál, and Csaba Szepesvári. Improved algorithms for linear stochastic bandits. Advances neural information processing systems, 24, 2011.
Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. Flambe: Structural complexity and representation learning of low rank mdps. Advances in neural information processing systems, 33:20095-2000 20107, 2020.
Animashree Anandkumar, Rong Ge, Daniel Hsu, Sham M Kakade, and Matus Telgarsky. Tensor decompositions for learning latent variable models. Journal of machine learning research, 15: 2773-2832, 2014.
Michael Arbel, Liang Zhou, and Arthur Gretton. Generalized energy based models. arXiv preprint arXiv:2003.0500 050033, 2020.
Karl Johan Aström. Optimal control of markov processes with incomplete state information. Journal of mathematical analysis and applications, 10(1):174-20) 1965.
Kamyar Azizzadenesheli, Alessandro Lazaric, and Animashree Anandkumar. Reinforcement learning of pomdps using spectral methods. In Conference on Learning Theory, pp. 193-256. PMLR, 2016.
Bram Bakker. Reinforcement learning with long short-term memory. Advances neural information processing systems, 14, 2001.
André Barreto, Will Dabney, Rémi Munos, Jonathan J Hunt, Tom Schaul, Hado van Hasselt, and David Silver. Successor features for transfer in reinforcement learning. Advances in neural information processing systems, 30, 2017.
Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemyslaw Debiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, et al. Dota 2 with large scale deep reinforcement learning. arXiv preprint arXiv:1912.064
2019.
Byron Boots, Sajid M Siddiqi, and Geoffrey J Gordon. Closing the learning-planning loop with predictive state representations. The International Journal of Robotics Research, 30(7):954-966, 2011.
Byron Boots, Geoffrey Gordon, and Arthur Gretton. Hilbert space embeddings of predictive state representations. arXiv preprint arXiv: 130.66130996/99 2013.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint arXiv:/606.013: 2016.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane. Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. In International Conference on Learning Representations, 2020.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine Deep reinforcement learning in handful of trials using probabilistic dynamics models. Advances in neural information processing systems, 31, 2018
Peter Dayan. Improving generalization for temporal difference learning: The successor representation. Neural computation, 5(4):613 1993.
Marc Peter Deisenroth and Jan Peters. Solving nonlinear continuous state-action- pomdps for mechanical systems with gaussian noise. 2012.
Carlton Downey, Ahmed Hefny, Byron Boots, Geoffrey J Gordon, and Boyue Li. Predictive state recurrent neural networks. Advances in Neural Information Processing Systems, 30. 2017.
Yaqi Duan, Tracy Ke, and Mengdi Wang. State aggregation learning from markov transition data. Advances in Neural Information Processing Systems, 32, 2019.
10

---

Under review as a conference paper at ICLR 2023
Eyal Even-Dar, Sham M Kakade, and Yishay Mansour. The value of observation for monitoring dynamic systems. In ICCA, pp. 2474-2479, 2007.
Norm Ferns, Prakash Panangaden, and Doina Precup. Metrics for finite markov decision processes. In UAI, volume pp. 162-169, 2004.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International Conference on Machine Learning, pp. 2170-2179. PMLR, 2019.
Noah Golowich, Ankur Moitra, and Dhruv Rohatgi. Learning in observable pomdps, without computationally intractable oracles. arXiv preprint arXiv:2206.03446, 2022a.
Noah Golowich, Ankur Moitra, and Dhruv Rohatgi Planning in observable pomdps in quasipolyno- mial time. arXiv preprint (7Xiv:2201.04733: 2022b.
Karol Gregor, Danilo Jimenez Rezende, Frederic Besse, Yan Wu, Hamza Merzic, and Aaron van den Oord. Shaping belief states with generative environment models for rl. Advances in Neural Information Processing Systems, 32, 2019.
Zhaohan Daniel Guo, Shayan Doroudi, and Emma Brunskill. pac algorithm for episodic pomdps. In Artificial Intelligence and Statistics, pp. 510-518. PMLR, 2016.
Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Bernardo A Pires, and Rémi Munos. Neural predictive belief representations. arXiv preprint 18Xiv: 1.06407, 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International conference on machine learning, pp. 1861-1870. PMLR, 2018.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint 1912v 2019a.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference machine learning, pp. 2555-2565 PMLR, 2019b.
Danijar Hafner, Timothy L Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02/99: 2020.
Matthew Hausknecht and Peter Stone. Deep recurrent q-learming for partially observable mdps. In 2015 aaui fall symposium series, 2015.
Nicolas Heess, Jonathan J Hunt, Timothy P Lillicrap, and David Silver. Memory-based control with recurrent neural networks. arXiv preprint arXiv: 1512.04455, 2015.
Ahmed Hefny, Zita Marinho, Wen Sun, Siddhartha Srinivasa, and Geoffrey Gordon. Recurrent predictive state policy networks. In International Conference on Machine Learning, pp. 1949- 1958. PMLR, 2018.
Maximilian Igl, Luisa Zintgraf, Tuan Anh Le, Frank Wood, and Shimon Whiteson. Deep variational reinforcement learning for pomdps. International Conference on Machine Learning, pp. 2117- 2126. PMLR, 2018.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as big sequence modeling problem. Advances in neural information processing systems, 34, 2021.
Haoming Jiang. Bo Dai, Mengjiao Yang, Tuo Zhao, and Wei Wei. Towards automatic evaluation of dialog systems: A model-free off-policy evaluation approach. Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 7419-7451, 2021.
Chi Jin, Sham Kakade, Akshay Krishnamurthy, and Qinghua Liu. Sample-efficient reinforcement learning of undercomplete pomdps. Advances in Neural Information Processing Systems, 33: 18530-18539, 2020a.
11

---

Under review as a conference paper at ICLR 2023
Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on Learning Theory, pp. 2137-2143. PMLR. 2020b.
Leslie Pack Kaelbling, Michael Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence, (01(1-2)99-134, 1998.
Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick Van der Smagt. Deep varia- tional bayes filters: Unsupervised learning of state space models from raw data. arXiv preprini arXiv:1605.06432. 2016.
Iyya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004.13649 2020.
Akshay Krishnamurthy, Alekh Agarwal, and John Langford. Pac reinforcement learning with rich observations. Advances in Neural Information Processing Systems, 29, 2016.
Tejas Kulkarni, Ardavan Saeedi, Simanta Gautam, and Samuel J Gershman. Deep successor reinforcement learning. arXiv preprint arXiv: 02396, 2016.
Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Pieter Abbeel. Model-ensemble trust-region policy optimization. arXiv preprint arXiv: 1802. /0592, 2018.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning, pp. 5639-5650. PMLR, 2020a.
Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Rein- forcement learning with augmented data. Advances in neural information processing systems, 33: 19884-19895, 2020b.
Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine, Stochastic latent actor-critic: Deep reinforcement learning with latent variable model. Advances in Neural Information Processing Systems, 33:741-752, 2020.
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuomotor policies. The Journal of Machine Learning Research, 17(1):13341373,, 2016.
Weiwei Li and Emanuel Todorov, Iterative linear quadratic regulator design for nonlinear biological movement systems. In ICINCO (1), pp. 222-229. Citeseer, 2004.
Michael Littman and Richard S Sutton. Predictive representations of state. Advances neural information processing systems, 14, 2001.
Qinghua Liu, Alan Chung, Csaba Szepesvári, and Chi Jin. When partially observable reinforcement learning not scary? arXiv preprint arXiv:2204.08967, 2022.
Omid Madani, Hanks, and Condon. On the computability of infinite-horizon partially observable markov decision processes. In AAA198 Fall Symposium on Planning with POMDPs, Orlando, FL. Citeseer, 1998.
Sridhar Mahadevan and Mauro Maggioni. Proto-value functions: A laplacian framework for learning representation and control in markov decision processes. Journal of Machine Learning Research, 8(10), 2007.
Bogdan Mazoure, Remi Tachet des Combes, Thang Long Doan, Philip Bachman, and R Devon Hjelm. Deep reinforcement and infomax learning. Advances in Neural Information Processing Systems, 0123686-3698,0 2020.
Lingheng Meng, Rob Gorbet, and Dana Kulić. Memory-based deep reinforcement learning for pomdps. In 2021 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 5619-5626. IEBE, 2021.
12

---

Under review as a conference paper at ICLR 2023
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
Ofir Nachum and Mengjiao Yang. Provable representation learning for imitation with contrastive fourier features. Advances in Neural Information Processing Systems, 34:30100-300 2021.
Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhutdinov. Recurrent model-free rl is a strong baseline for many pomdps. arXiv preprint arXiv:2/ 2021.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint 18X7038807,03748, 2018.
Christos Papadimitriou and John N Tsitsiklis. The complexity of markov decision processes. Mathematics of operations research, 12(3):441-435 1987.
Shuang Qiu, Lingxiao Wang, Chenjia Bai, Zhuoran Yang, and Zhaoran Wang. Contrastive ucb: Prov- ably efficient contrastive self-supervised learning in online reinforcement learning. In International Conference Machine Learning, pp. 18168-18 18210. PMLR, 2022.
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. Advances in neural information processing systems, 20, 2007.
Tongzheng Ren, Tianjun Zhang, Lisa Lee, Joseph E Gonzalez, Dale Schuurmans, and Bo Dai. Spectral decomposition representation for reinforcement learning. arXiv preprint arXiv:2208.095155 2022a.
Tongzheng Ren, Tianjun Zhang, Csaba Szepesvári, and Bo Dai. A free lunch from the noise: Provable and practical exploration for representation learning. In Uncertainty in Artificial Intelligence, pp. 1686-1696 PMLR, 2022b.
Jürgen Schmidhuber.
Reinforcement learning in markovian and non-markovian environments.
Advances in neural information processing systems, 3, 1990.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference machine learning, pp. 1889-1897. PMLR, 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov, Proximal policy optimization algorithms. arXiv preprint 1777v:1707.06347, 2017.
Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-efficient reinforcement learning with self-predictive representations. arXiv preprint arXiv:2007.05999, 2020.
Satinder Singh, Michael R James, and Matthew R Rudary. Predictive state representations: new theory for modeling dynamical systems. In Proceedings of the 20th conference on Uncertainty in artificial intelligence, pp. 512-519, 2004.
Yang Song and Diederik P Kingma. How to train your energy- based models. arXiv preprint arXiv:2101.03288, 2021.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In International Conference Machine Learning, pp. 9870-9879 PMLR, 2021.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv:1801.000690 2018
Masatoshi Uchara, Xuezhou Zhang, and Wen Sun. Representation learning for online and offline rl in low-rank mdps. arXiv preprint arXiv:2/ 10.04652, 2021.
Masatoshi Uehara, Ayush Sekhari, Jason D Lee, Nathan Kallus, and Wen Sun. Provably efficient re- inforcement learning in partially observable dynamical systems. arXiv preprint arXiv:2206. 12020, 2022.
13

---

Under review as a conference paper at ICLR 2023
Arun Venkatraman, Nicholas Rhinehart, Wen Sun, Lerrel Pinto, Martial Hebert, Byron Boots, Kris Kitani, and J Bagnell. Predictive-stir statt decoders: Encoding the future into recurrent networks. Advances in Neural Information Processing Systems, 30, 2017.
Nikos Vlassis, Michael Littman, and David Barber. On the computational complexity of stochastic controller optimization in pomdps. ACM Transactions on Computation Theory (TOCT), 4(4):1-8, 2012.
Lingxiao Wang. Qi Cai, Zhuoran Yang, and Zhaoran Wang. Embed control partially observed sys- tems: Representation learning with provable sample efficiency, arXiv preprint arXiv:2205. 13476, 2022.
Tingwu Wang, Xuchan Bao, Ignasi Clavera, Jerrick Hoang, Yeming Wen, Eric Langlois, Shunshi Zhang. Guodong Zhang, Pieter Abbeel, and Jimmy Ba. Benchmarking model-basee reinforcement learning. arXiv preprint arXiv: (ane:::07700055 2019.
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. Advances neural information processing systems, 28, 2015.
Daan Wierstra, Alexander Foerster, Jan Peters, and Juergen Schmidhuber. Solving deep memory pomdps with recurrent policy gradients. In International conference on artificial neural networks, pp. 697-706. Springer, 2007.
Yifan Wu, George Tucker, and Ofir Nachum. The laplacian in rl: Learning representations with 2018.
efficient approximations. arXiv preprint
Mengjiao Yang, Sergey Levine, and Ofir Nachum. Trail: Near-optimal imitation learning with suboptimal data. arXiv preprint arXiv:2 14770, 2021.
Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improving sample efficiency in model-free reinforcement learning from images. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 10674-1068) 2021.
Tianhe Yu, Deirdre Quillen, Zhampeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pp. 1094-1100. 100. PMLR, 2020.
Wenhao Zhan, Masatoshi Uehara, Wen Sun, and Jason D Lee. Pac reinforcement learning for predictive state representations. arXiv preprint arXiv:2207.05738, 2022.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006. 10742, 2020.
Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew Johnson, and Sergey Levine. Solar: Deep structured representations for model-based reinforcement learning. In International Conference on Machine Learning, pp. 1444-7453. PMLR, 2019.
Tianjun Zhang, Tongzheng Ren, Mengjiao Yang. Joseph Gonzalez, Dale Schuurmans, and Bo Dai. Making linear mdps practical via contrastive representation learning. In International Conference on Machine Learning, pp. 26447-26466. PMLR, 2022.
Pengfei Zhu, Xin Li, Pascal Poupart, and Guanghui Miao. On improving deep reinforcement learning for pomdps. arXiv preprint arXiv: 1704.07978 2017.
14

---

Under review as a conference paper at ICLR 2023
A MORE RELATED WORK
Provably RL for POMDPs. Besides the statistical and computational hardness results for learning and planning upon POMDPs, most recent theoretical research focuses overcoming the statistical complexity from the "curse of history" by considering tractable POMDPs (Krishnamurthy et al., 2016; Azizzadenesheli al., 2016; Guo al., 2016; Jin al., 2020a; Liu al., 2022). 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. Recently, Uehara et al. (2022); Wang et al. (2022) generalize these special structures with function approximation beyond tabular cases, Golowich et al. (2022a) consider the complexity planning and exploration together with learning, but only valid for tabular MDPs
B DERIVATION OF THE RANDOM FEATURE IN EQUATION 8
We have that P(or-1lZt, as) =p(0 exp (f(x), ar) (g(0c+1) + +A((z(,a)))
=p(or+1) exp ((^ ((~2))
1 2.) [[(())))))3) exp
[9(0(+1))l*))
[[f(zz, a.) + +((o+=1)l))
exp
2
2
where we have that
[ ar) + g(o
exp
2
[1((0),))) at) + (lottt))?) / exp (- I| (f(rz, + s(o++1))l? de
= (2m))= 4/2 exp
-4/2
2
2
=(2m) 4/2 / exp (- ||wll
-22m)-4/2
+ + (f(s,,)))))))))))))))))))))))( + g(ove1))) dw
2
-Ewn,N1o.ta [exp (w T f(z,, ad)) exp (w s(ou+1))], which concludes the proof for equation 8.
C
OBSERVABLE LQG AS EPR
Follow the standard notations, the dynamics of sinear-Quadratic Gaussian is defined as =As + Ba, + wt,
=Cs,, + 2t,
where w. and 2, are Gaussian noise. Define the matrix
GL = [CT, CCA
(CA
and reduced observation
1-2
1-2 A W
o, = 24 C
Σ
* Ba
k
k=0
By the observability condition of LQG, Gr. is full column rank, one can identify 8o by
50 = (G.
Σ (AT)²-1 cToz
15

---

Under review as a conference paper at ICLR 2023
Therefore, have
$1 =Aso + Bag wo A (G[GL) E (AT) - C.
+ Ba1 wo:
82 = + Baj + w1 = A
(G]
-1
Σ (AT) CToz
+ +BBa + BB2 + Awo + wi,
j=l
81,11 =As, Bag 1 At (c..G.)
-1 Σ
(AT))- CTo ((AT)--CTT
+ Σ
A BBjiBaj+1 Σ
A4-1
5=1
j=0
1=0
L CTaj + +C
L
L
01+1 -0
Σ
B Bag+1 + + E
Al1J + 2t:
j=1
J=0
3=0
which means 01++1 follows a Gaussian distribution with mean as function of history IT. (i)- and action az+1, and variance as a function of ows 0z. and B.C) There-
fore, we have some function (A,,,C,ow.o, and B.a..,, such that JA,,BC.ow.o. (oL+1) SA,B,C.o aa (LL,"L+1) + {~N(0,1)
On the other hand, we set À - 1/2 and p(o) =N =O.0, in (6), then, we obtain |Io(oo+1) f(xL,a)))
p IT, |
x exp
2
which reproduces the observable LQG with specific SA,BBCoowaa and JA, B.C.C..a....
D EXPERIMENT DETAILS
D.1 ONLINE SETTING
In Table 9, we list all the hyperparameters and network architecture we use for our experiments. We that don't 't use the additional exploration bonus term in the mojuco tasks. But this is very helpful in DM control suite tasks, especially in those sparse-reward tasks.
For evaluation in Mujoco, in each evaluation (every 5K steps) test algorithm for 10 episodes. We average the results over the last 4 evaluations and random seeds. For Dreamer and Proto-RL, we change their network from CNN 3-layer MLP and disable the image data augmentation part (since we test on the state space). The architecture we used for the transformer is following the Trajectory Transformer (Janner et al., 2021). The attention used is the causal attention. We tried to tune some of their hy terrparamettrr (c.g., exploration steps in Proto-RL) and report the best number across our runs. However, due to the short time, is also possible that we didn't tune the hyperparameter enough.
D.2 LEARNING CURVES
We provide the performance curves for online DM Control Suite experiments in Figure 2. As see in the figures, the proposed EPR converges faster and achieve the state-of-the-ar performances in most of the environments, demonstrating the sample efficiency and the ability to balance of exploration vs. exploitation of EPR. We also provide additional curves for POMDP setting in Figure 3.
D.3 IMAGE-BASED EXPERIMENTS
We provide the details of metaworld image- masee experiments here. We first provide an illustration of the reach environment in Figure 4. then provide some more experiment details the following section.
16

---

Under review as a conference paper at ICLR 2023
Table 5: Hyperparameters used for EPR in all the environments in MuJoCo and DM Control Suite.
Hyperparameter Value
Bonus Coefficient (MuloCo)
0.0
Bonus Coefficient (DM Control)
5.0
Actor Ir
0.0003
Model Ir
0.0003
Actor Network Size (MuloCo)
(256, 256)
Actor Network Size (DM Control)
(1024, 1024)
ERP Embedding Network Size (MuloCo)
(1024, 1024, 1024)
ERP Embedding Network Size (DM Control)
(1024, 1024, 1024)
Critic Network Size (MuloCo)
(1024, 1)
Critic Network Size (DM Control)
(1024, 1)
Discount
0,99
Target Update Tau
0.005
Model Update Tau
0.005
Batch Size
256
EPR
Deep5F
SAC (2-layrr
SAC (3-ayer)
PPO (3-layer)
cheetah,nu
cheetah,
Walker Sparse
600
Ratur
300
200
200
100
100
0.0
0.5
1.3
1.5
2.0
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
15
2.0
165
105
Environment Steps
1e5
humanoid
15.0
800
10.0
600
7.5
200
2.5
0.0
1.0
Steps
Figure 2: Performance Curves for online DM Control Suite.
Table 6: Settings of adapted OpenAl Environment.
Hyperparameter Value
Maximum Episode Steps
50
Reward Type
'sparse"
Observation Size
(3, 64, 64)
Fixed Goal Position
(1.27, 0.90, 0.66)
17

---

Under review as a conference paper at ICLR 2023
Cheelah run
Chestah_nnn _u____r____eeeeeee___________e__
Humanoid
MLP PSR
400
PSR
PSR
EPR
Trans
300
Trans
Trans
100
100
200
0.5
2.0
training sleps (1e5)
training steps (1c5)
traniing (1e5)
Walker_nur run
Walker Bonker__ru__eeee uun_speree
600
500
M.P
400
Trame
EPR
EPR
900
0.0
1.0
1.5
2,0
0.5
1.0
1.5
2.0
training steps (1e5)
training steps (1e5)
Figure 3: Performance Curves for online POMDP DM Control Suite.
Figure 4: Reach environment: Using robot arm to reach specific position.
Table Hyperparameters used for EPR in FetchReachlmage.
Hyperparameter Value
Bonus Coefficient (MuloCo)
0.0
Bonus Coefficient (DM Control)
5.0
Actor Ir
0.0003
Model
0.0003
Actor Network Size (MuloCo)
(256, 256)
Actor Network Size (DM Control)
(1024, 1024)
ERP Embedding Network Size (MuJoCo)
(1024, 1024, 1024)
ERP Embedding Network Size (DM Control)
(1024, 1024, 1024)
Critic Network Size (MuloCo)
(1024, 1)
Critic Network Size (DM Control)
(1024, 1)
Discount
0,99
Target Update Tau
0,005
Model Update Tau
0.005
Batch Size
18

---

Under review as a conference paper at ICLR 2023
Table 8: Hyperparameters used for SPR in FetchReachlmage.
Hyperparameter Value
Ir
0.0001
Dropout
0.5
Discount
0.99
Batch Size
32
Augmentation
off
Target Update Tau
0.005
Model Update Tau
0.005
Batch Size
256
Update
Distributional Q
Dueling
True
Optimizer
Adam
Optimizer: learning rate
0,0001
Max gradient norm
10
Priority exponent
0.5
Noisy nets parameter
0.5
Min replay size for sampling
2000
Replay period every
1 step
Updates per step
2
Multi-step return length
Q network: channels
32, 64, 64 8 x 8, 4x8,3x
Q network: filter size
Q network: stride
2,
Q network: hidden units
256
Non-linearity
ReLU
Target network: update period
1
À (SPR loss coefficient
2
K (Prediction Depth)
5
Table 9: Hyperparameters used for SAC-AE in FetchReachlmage.
Hyperparameter Value
Critic Ir
0.001
Actor Ir
0.001
Discount
0.99
Batch Size
128
Critic Q- Q-function soft-update rate TQ
0.01
Critic encoder soft-update rate Tene
0.05
Critic target update frequency
2
Actor update frequency
2
Actor standard deviation bounds
[-10,2] 2]
Autoencoder learning rate
0.001
Temperature learning rate
0.0001
Temperature Adam's Init temperature
0.5
0.1
19


Reviews:
[{"review_id": "xbYtYJoo9Z", "reviewer": "Reviewer_TDvH", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper offers a method for reinforcement learning under partial observability by offering a new form of representing the environment model. Specifically, this representation takes advantage of the structure of the environment and linear approximation of value/Q functions. The authors showed the superiority of their method on some MuJoCo tasks.", "summary_of_the_review": "Overall, I believe this is a very interesting and promising idea, but I wish to see more standard experiments and/or more intuition about the type of tasks/environments where EPR is useful.  ", "strength_and_weaknesses": "Strength:\nThe ideas behind taking advantage of the environmental structure were very interesting. Moreover, deep networks are still not performing well in domains with uncertainty, making this line of work very worthy of more exploration. Finally, I think the authors did a very good job of explaining the connection between EPR and other existing representations (although these connections could be moved to the appendix to make room for explaining the domains that EPR works/does not work) \n\nWeakness: \nI think the paper could be significantly improved if the authors give more intuition about the type of tasks in which their method works and tasks where it does not, especially classic POMDP tasks such as RockSample or classic localization problems. Particularly, it looks like the belief should be concentrated around a state and its surrounding (as opposed to multi-modal), so the method works well. \n\nI also found the choice of masking velocity to make the tasks POMDP in experiments a little bit strange. What was the rationale behind that? Generally, while I do believe that taking advantage of real problem structures is totally okay and in fact useful, uncertainty about agents' own sensing is not very much aligned with reality. There is actually a very nice benchmark for POMDPs based on MuJoCo tasks in \"Recurrent Model-Free RL Can Be a Strong Baseline for Many POMDPs\" by Ni et al, ICML 2022 (repo: https://github.com/twni2016/pomdp-baselines). They have defined different types of Partial Observabilty and I think using even one type of these test-bench marks is extremely valuable for evaluation. Also, it looks like some baseline methods are provided by them. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "NtElzHYonlV", "reviewer": "Reviewer_m9BZ", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper aims at using \"Energy-based Predictive Representation (EPR)\" to create 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\".", "summary_of_the_review": "The paper is not clear.", "strength_and_weaknesses": "Strengths:\n- The claims are very ambitious\n\nWeaknesses:\n- The paper does not provide a clear formalization and description of the algorithm", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "3KWPadK4iN", "reviewer": "Reviewer_4vmU", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a new algorithm dealing with reinforcement learning (RL) in partially observable environments by linearly representing action-value functions. To learn linear action-value functions, the energy-based model is applied to represent a valid transition function. This energy-based method bypasses the explicit computation of beliefs. Then the learned energy-based transition on beliefs can be linearly factorized by applying the existing random Fourier feature trick. The action-value function and additional reward are represented by using the random feature, and numerical results show promising results on exploration, exploitation, and planning in RL with MDPs/POMDPs.", "summary_of_the_review": "I have several concerns about the core idea as well as clarity, novelty, and reproducibility, and I hope the authors can address this in their rebuttal. \n", "strength_and_weaknesses": "The authors dealt with the formulation of RL in POMDPs with various interesting concepts. However, several parts of the paper should be clarified, and additional experiments should be performed. It was hard to check the reproducibility. Please refer to the detailed comments below.\n\n\nComments\n\n1. The core idea of linear value-function representation is not well-motivated. In this paper, the following statement is written: \"To handle the complexity induced by large state spaces, linear/low-rank structures have been introduced in MDPs.\" However, several questions arise.\n- Is linear representation guaranteed to be effective in POMDPs? Since linear representation lacks representation ability compared with nonlinear counterparts, it can be inefficient in POMDPs, and even in MDPs with large spaces.\n- Does the proposed linear representation enjoy better Q-function convergence?\n- This statement contradicts the core idea of linear representation. “In practice, we perform fitted Q iteration with a nonlinear component extending the linear parameterization.” (page 6)\n\n2. Regarding energy-based parameterization\n- “We avoid any explicit parameterization and computation of beliefs” (page 4). To support the superiority of the energy-based model over explicit parameterization, we require explicit parameterization methods including e.g., DVRL (Igl et al., 2018) or SLAC (Lee et al., 2020).\n- The background of energy-based parameterization should be explained in detail (e.g., in Preliminaries) since it is one of the core components in the proposed method.\n\n3. Regarding the random Fourier feature trick\n- It would be better to explicitly write the full derivation of Eqs. 9 and 10 for the readers since this is also another core component of the proposed method.\n- What is the sample number $n$ in the random Fourier feature trick? Is the performance robust to  $n$?\n\n4. Regarding exploration\n- Is extracting the random feature $\\phi$ the best method for generating additional bonus reward? Any other intrinsic reward method should be compared to support the claim.\n\n5. Experiments\n- Instead of the choice of masking velocity to make the tasks POMDP in experiments, is the proposed method still effective when masking position$^{a}$ (and let velocity visible), or input with noise$^{b}$? \n\na. Han et al., Variational Recurrent Models for Solving Partially Observable Control Tasks \n\nb. Meng et al., Memory-based Deep Reinforcement Learning for POMDP \n\n- As stated above, other explicit belief parameterization methods and other intrinsic reward for exploration are required.\n- In Tables 3 and 4, please describe the details of MLP and Trans. Does the transformer use causal encoding? Is positional encoding used? What is the number of layers and multi-heads of Trans?\n- Why the proposed method is effective in MDPs? Is that because of linear representation?\n- What is the number of negative samples $m$? Is the performance robust to  $m$ (and $\\lambda, \\alpha$)?\n\n6. Notation\n- Is $x_t$ implemented using RNN?\n- $\\gamma$ is missing in Assumption 1.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "22P-wUuSlFy", "reviewer": "Reviewer_XFFG", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes Energy-based Predictive Representation (EPR) for tractable learning, planning and exploration in POMDPs and fully observable MDPs. Concretely, the authors propose a nonlinear energy-based learnable model which does not explicitly parameterize the beliefs. Their model provides a linear sufficient representation for the state-action value function, which can be particularly beneficial to large state spaces. Subsequently, the authors introduce computationally tractable algorithms for learning, exploration/exploitation and planning, based on the principles of optimism and pessimism in the face of uncertainty for online and offline RL. The paper also discusses the connections between EPR and Predictive State Representation (PSR), and points out that the latter is not generally associated with efficient exploration and planning algorithms. The authors conduct an extensive empirical evaluation with both fully observable MDPs and POMDPs, comparing their approach to several model-based, model-free and representation-learning frameworks. The results confirm the superior empirical performance of the proposed representation learning scheme.", "summary_of_the_review": "Overall, I believe that the current work is a very interesting addition in the literature of representation learning for POMDPs with large state spaces, with strong and promising empirical results.", "strength_and_weaknesses": "Strengths\n- The proposed framework builds upon several prior works, e.g., prior work on linear/low-rank structures and spectral factorization for MDPs, random Fourier features, elliptical confidence bounds, feature space augmentation. The authors explain the various steps for building their framework in detail, and provide plenty of references in the process. This makes the paper not hard to follow.\n- Related work from several domains is covered quite thoroughly. Connections to prior works on representation learning (e.g., PSR and SPEDE) are also established.\n- The proposed framework seems able to solve a real problem for POMDPs, especially as the size of the state space becomes large. Furthermore, the 2 assumptions of \\gamma-observability and linearity are meaningful and well motivated by prior work.\n- The empirical results show strong performance. I particularly like that the EPR framework can be beneficial not just for POMDPs, but for fully observable MDPs as well.\n\nWeaknesses\n- There are quite a few typos. I suggest the authors proofread the document carefully. Examples: Th predictive -> The predictive, has has proved -> has proved, with slightly abuse -> with slight abuse, W propose -> We propose, an straightforward extension -> a straightforward extension, etc.\n- The notation was confusing at times. For instance, the authors introduced a symbol x_t^L for fixed-window histories (of size L) but then used x_t in equations (6) and (7) (and even elsewhere). I think consistency in the notation everywhere in the paper would help.\n- I think there is an (easy to fix) error in Assumption 1. \\gamma-observability needs a \\gamma multiplicative factor in the right term of the inequality. This is currently missing.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject