Distributional Policy Evaluation: a Maximum Entropy approach to Representation Learning
Authors: Riccardo Zamboni
Alberto Maria Metelli
Marcello Restelli
Venue: NeurIPS 2023 poster (2023)
Keywords: reinforcement learning
distributional reinforcement learning
maximum entropy estimation
representation learning

Abstract: The Maximum Entropy (Max-Ent) framework has been effectively employed in a variety of Reinforcement Learning (RL) tasks. In this paper, we first propose a novel Max-Ent framework for policy evaluation in a distributional RL setting, named *Distributional Maximum Entropy Policy Evaluation* (D-Max-Ent PE). We derive a generalization-error bound that depends on the complexity of the representation employed, showing that this framework can explicitly take into account the features used to represent the state space while evaluating a policy. Then, we exploit these favorable properties to drive the representation learning of the state space in a Structural Risk Minimization fashion. We employ state-aggregation functions as feature functions and we specialize the D-Max-Ent approach into an algorithm, named *D-Max-Ent Progressive Factorization*, which constructs a progressively finer-grained representation of the state space by balancing the trade-off between preserving information (bias) and reducing the effective number of states, i.e., the complexity of the representation space (variance). Finally, we report the results of some illustrative numerical simulations, showing that the proposed algorithm matches the expected theoretical behavior and highlighting the relationship between aggregations and sample regimes.

Distributional Policy Evaluation: a Maximum Entropy approach to Representation Learning
Riccardo Zamboni
Alberto Maria Metelli
DEIB, Politecnico di Milano
DEIB, Politecnico di Milano
Milan, Italy
Milan, Italy
riccardo. www.zzimmbnni.. @po limoliii.... timiiit it
albertomaria metelll@poll imi
Marcello Restelli DEIB, Politecnico Milano
Milan, Italy
carcelloorestelli@ppl/..
Abstract
The Maximum Entropy (Max-Ent) framework has been effectively employed in a variety of Reinforcement Learning (RI) task) In this paper, first propose novel Max-Ent framework for policy evaluation in a distributional RL setting, named Distributional Maximum Entropy Policy Evaluation (D-Max-Ent PE). We derive a generalization-erro bound that depends on the complexity of the representation employed, showing that this framework can explicitly take into account the features used to represent the state space while evaluating policy. Then, we exploit these favorable properties to drive the representation learning of the state space in a Structural Risk Minimization fashion. We employ aate-aggregation functions as feature functions and we specialize the D-Max-Ent approach into an algorithm, named D-Max-Ent Progressive Factorization, which constructs progressively finer finer-grained representation of the state space by balancing the trade-off between preserving information (bias) and reducing the effective number of states, i.e., the complexity of the representation space (variance). Finally, we report the results of some illustrative numerical simulations, showing that the proposed algorithm matches the expected theoretical behavior and highlighting the relationship between aggregations and sample regimes.
Introduction
In Distributional Reinforcement Learning (D-RL) [ellemare et al. 2023 an agent aims to estimate the entire distribution of the returns achievable by acting according to a specific policy. This is in contrast to and more complex than classic Reinforcement Learning (RL) [Szepesvári 2010 Sutton and Barto 2018 where the objective to predict the expected return only. In recent years, several algorithms for D-RL have been proposed, both in evaluation and control settings. The push towards distributional approaches was particularly driven by additional flavors they can bring into the discourse, such as risk-averse considerations, robust control, and many regularization techniques [Chow al. 2015] Brown et al. 2020 Keramati et al., 2022].. Most of them varied how distribution of the returns is modcled. The choice of the model on how such a distribution can be learned, how efficiently and with what guarantces, and how it can be used for the control problem. Similarly to this tradition, this paper investigates the potential of looking into the entire distribution of returns to address the representation learning of the state-action spaces, that is find good feature representation of the decision- making space as make the
shown to have cascading effect
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

overall learning problem easier, tenderly by reducing the dimensionality of such spaces, In particular, it points to answer the following research question:
Q1: What tools can return distributions provide for distributional RL?
In Section 3 we answer this methodological question, showing that it is possible to reformulate Policy Evaluation in a distributional setting so that its performance index is explicitly intertwined with the representation of the (state or action) spaces. More specifically, this work tackles Policy Evaluation in distributional setting then applies the Maximum Entropy (Max-Ent) formulation of distribution estimation [Wainwright and Jordan 2007]. In this way, it is possible to derive a novel framework for PE, which we name Distributional Max-Ent Policy Evaluation (D-Max-Ent PE), which inherits the many positives the Max-Ent framework offers. In particular, it allows the inclusion of constraints that the distribution needs to satisfy, usually called structural constraints, namely feature-basse the support of the distribution. Such constraints then appear in the generalization error bound for the Max-Ent problem, with term related to the complexity of the family of features used. Unfortunately, traditional derivations of this bound introduce quantities that are bounded but unknown. We develop the analysis further to derive a more practical bound, containing quantities that are either estimated or known. In this way, the generalization-erro bound shows a usual bias-variance trade-off that can be explicitly optimized by changing the feature functions adopted, while making the best use of the available samples. Thus, PE in a distributional setting is directly linked to representation learning.
particular case of distribution estimation problem and
constraints acting over
Now, the RL literature proved that reducing the state space size while preserving the important features of the original state space is beneficial. namely with aggregaggnnn feature functions (Singh et al.. 1994 Van Roy, 2006) Dong et al. 2020]. This is particularly true when high dimensionality can make learning slower and more unstable, as in classic RI. general, or when the learning process is almost unfeasible in small-samples regimes, as for D-RL., where learning the entire distribution of returns requires a large number of samples. Thus, motivated by these considerations, while D-Max-Ent Policy Evaluation allows for the use of any type of structural constraint, this work focuses on state-aggregation feature functions, and we exploit the first and more methodological result to answer a second more algorithmic question:
Q2: How are representation learning and policy evaluation intertwined? Do distributional methods offer new way to highlight and exploit this connection?
To answer this question, in Section 4 we show that the generalization-erve bound changes monotoni- cally when the state-aggregation constraints are changed in specific way, namely a finer-grained representation of the space. This is indeed what happens in the Structural Risk Minimization theory (SRM) [Vapnik] 1991]. Similarly, we develop a novel algorithm called D-Max- Ent Progressive Factorization, which exploits the proposed evaluation method to learn representation of the state space soundly, i.e., trying to reduce a proxy of the generalization error bound in an SRM fashion, and allowing us to answer positively to the second research question as well. Finally, in Section 5 we verify through an illustrative numerical simulation whether the proposed algorithm matches the behaviors suggested by the theoretical analysis.
2 Preliminaries
2.1 Markoy Decision Processes
discrete time finite Markov decision processes (MDP) PPteeran 1994 is a tuple M = (S, A, Ps. PR, 7, where S is a finite state space (|S]= S), A is a finite action space ([A|= A). Ps S S - (SS) is the transition kernel, PR S A (R) the reward distribution function, E A(S) is the initial-state distribution and E [0. 1) is discount factor! A policy π S + ((A) defines the behavior of an agent interacting with an environment, which goes as follows: starting from an initial state So n, an agent interacts with the environment through the policy x, generating a trajectory H = (S,, A.. Re 2o which is sequence of states, actions and rewards whose joint distribution is determined by the transition kernel, reward distribution, and the policy itself, i.c.,
A, ~ 7(-)S)), R ~ PR PR(- At), and S ~ Ps (-1S,,
(2)) denotes the simplex of space X.
2

---

2.2 Value Functions and Distributions of Returns
Given an MDP M with discount factor y, the Discounted Return is the sum of rewards received from the initial state onwards, discounted according to their time of occurence::
(*)
R s.
(I)
t=0
The Value Function of given policy is the expectation of this quantity under the policy itself:
OC ?' RR
V [ E
S =
(2)
The Return Distribution Function 17" of given policy is a collection of distributions, one for each state S, where each element is the distribution of the random variable (s):
80
1f (T
D
Σ
2 R So =
(3)
where Di, extracts the probability distribution of random variable under the joint distribution of the trajectory.
The Distributional Policy Evaluation Problem then consists of estimating the return distribution function of Eq. 3 for fixed policy
2.3 Maximum Entropy Estimation
Maximum Entropy (Max-Ent) methods (Dudik and Schapire 2006 Wainwright and Jordan 2007
Sutter et al. 2017] are density estimation methods that select the distribution that maximizes the uncertainty, i.e., the one with maximum entropy, where the entropy of a distribution p E A ( is
defined as (p) = - Ex log p(
[ Additionally, they assume that the learner has access to
feature mapping J from X' to R In the most general case, we may have M = +oc. We will denote by + the class of real-valued functions containing the component feature functions ʃj F with [M]. A distribution p is consistent with the true underlying distribution po if
E Ex~elf My
Vj E [M].
(4)
where
Bi Bx=po [X)
In this case, we say that satisfies (in expectation) the struetural constraints imposed by the features in F. In practice, Po is not available and Max-Ent methods enforce empirical consistency over N independent and i..dd observations D = (II,...,,N, ~ with support in X' by replacing the definition in Eq. (5) with
1 N (D) = N Σ f(l))
[M]
(6)
The distribution p is said be consistent with the data D if it matches the empirical expectations. The empirical Max-Ent problem consists then of the following optimization problem
H(p)
PEA(X)
(7)
s.t. Exx1 = n.. Vj E [M].
with the optimization problem in expectation differing just in the constraints (i.e., replacing constraint from Eq. (6) with the ones from Eq. (5)) It is well known that the optimal solution to the empirical
With little abuse of notation, we will use the same symbol for the probability distribution and its p.d.f., which assume exist w.r.t. reference measure.

---

Max-Ent problem in Eq. 7 is distribution PA E A(x) belonging to the class of exponential
distributions parametrized by the parameters A. namely:
px(x) = +2 exp
Σ }5Ss(a) (E(M]
(8)
where Φx = fx exp (ISEIM )xf)(2²)) dæ' is a normalization constant, which ensures that P E A(X), and its log-tue transformation takes the name of log-partition function A(A) = log Jxx exp(Eje[M) (z)) The log-partition function defines the of well-behaved dis- tributions Ω = {) E RM A(A) too). At optimality, the parameters are defined as Å and correspond to the optimal Lagrangian multipliers of the dual of the empirical Max-Ent problem in Eq. 7 Now on, we will use á to identify Py for simplicity.
3 Distributional Policy Evaluation: A Max-Ent Approach
This section aims at answering the first research question QQ The proposed ap- proach turns distributions PE into pure density estimation problem in a Max-En framework, called Distributional Max-Ent Policy Evaluation, as described in Alggo- rithm For this translation, the algorithm uses the distribution of returns as p. N- trajectory samples HN [ as data, and fixed set of features functions F be- longing a function class 4b Note that
Algorithm Distributional Max-Ent Policy Evaluation
Require: (HN, features functions n = argmax H())
N trajectory samples, set of
s.t. Ex-, Ex-,ff (x)]]=xx(44 Vj E [M]
1) E (22) return f
do this, we need to slightly change the notation concerning the D-RL framework: 7) will not be [S] -vector of distributions with support over R. but rather a joint distribution over the whole support X' S x R. Turning PE into a Max-Ent problem has many upsides. First of all, the Max-Ent principle allows deal with any kind of support X' unifying continuous and discrete cases under the same framework; secondly, it does not require specifying a family of probability distributions to choose from; moreover, it implicitly manages the uncertainty by seeking a distribution as agnostic as possible, i.e., as close to the uniform distribution as possible. Finally, Max-Ent allows to include of structural constraints over the return distribution under many different flavors, both as in the standard
approximation methods (Van Roy 2006) and as in more recent works based on statistical functionals acting over the return portion I of the support [Bellemare et al., 2023]. One of the possible limitations might be the requirement to have access batch of i.i.d. samples, but this is not necessarily restrictive: the result can be generalized for a single 3-mixing sample path by exploiting blocking techniques Yu 1994 Nachum et al.) 2019
3.1 Generalization Error Bound
As previously said, the inner properties of Max-Ent allow for translating the results from density estimation methods the distributional PE setting, and in particular, generalization-eror bounds defined as KL-divergences Unfortunately, the generalization error bounds of traditional Max-Ent theory contain conservative term that compares the solutions of the expectation and empirical Max-Ent problems, 7, A respectively, by taking the maximum between the 1-norm of the respective multipliers, namely max (X.5) ||A|[ This quantity is bounded yet unknown, making the result unpractical. In the following, we extend the previous results with more practical bound containing
||Ål] instead of the maximum, requiring some additional assumptions about the expressiveness of the feature functions. This result is of independent interest and allows us to directly use the bound from an algorithmic perspective.
Theorem 1 (Generalization Error Bound of D-Max-Ent PE). Assume that the set of features F belong to the function class 4 which it is such that sup [lf( = F < +00 and that the
The KL- divviveree between two distributions p. q defined (L(p))g) = -Er-ploog(p(=)/((t))

---

minimum singular value muin of the empirical covariance matrix of the features Cov(F) is strictly positive, namely min (C)v(F)) O. Then, given sample batch {z1..... aNn X of N i.i.d. points drawn from the true distribution 1 for any E (0, 1), it holds with probability at least 1 that the solution the sampled Max-Ent problem 9 satisfies the following:
KL t ((0) 2 (1 + L(1)) + F, ,,
(9)
N
[(i)) = N
Σ
log (((;))
(10)
O
log 1/6 2N
, = 1011A|1 Rx(()) + FV
(11)
where 3 stands for the fact that the bound comprises additional terms that decrease at a higher rate in sample complexity and were therefore neglected. (n)") and ((i). the empirical log- log-likelihood of the solution, form bias term. The remaining term B(,,F,, Nr 8) is variance term depending the multipliers characterizing the solution 2, the number of samples, the confidence level d, and the feature class complexity as the empirical Rademacher complexity of the class Rx [Mohri et al. 2018).
3.2 Proof Sketch
Here we report the main steps of the proof of Th. The interested reader can find the complete proof in Appendix A First, define the containing the solutions the expected and sampled Max-Ent problems with S {17, i the related set for the multipliers fgg {", Å}. and a quantity that will be central now on (211.... NNN == max [E,,, log n] - 1 Σ log nzx))l. Then, the building N
blocks of the error term (1) (I)),, namely L((()) and K (' (l are bounded by:
LL(ll) S 22h(-)
KT ([0) S -H())) Z(1) + 31()).
is possible to show that:
log 1/8
1(-) < sup |||l|2 Rx(()) + FV XEñs
2N
6M
sup ||А|| S |1111 XEOs
Omin (Cov(F)) 11(-)
The first inequality is obtained with standard methods as in van der Vaart and Wellner [1996], Dudley [1999], Koltchinskii and Panchenko (2002). Waan et al. [2013]. The second one obtained by exploiting the intrinsic properties of the Max-Ent solution and by noting that it is possible to link h(-) with the Bregman divergence of the log- partition function Da (X, Å). One can that the use of the second incquality introduces an additional assumption about the expressiveness of the feature functions, requiring the minimum singular value of the sampled covariance matrix mmiul (Cov(F)) to be strictly positive. As final step, setting x / h/( NN and combining the two previous inequalities yields quadratic inequality:
2 bx 0 0
b 2 amin (F))
GM
RN((1) + F
log 1/5
2N
c 2|/2112 RN(() + Fv/
log 1/8 1/5 2N
which is well-defined and solves for
log 1/8
h(21. ,NN) , 2 ||i111
RN (Rx(0)+F) FF
2N
5

---

by neglecting higher-order terms. The statement of the theorem is then just a matter of combining all these results.
4 Distributional Representation Learning with State Aggregation
This section addresses the second research question Q2 namely how to use the bound in Th. from an algorithmic perspective automatically refine the features used to represent the state space in a principled way while performing D-Max-Ent PE. In particular, the focus is on specific instance of feature functions for return distributions, namely state aggregation. More specifically, the state aggregation feature functions F (fi)ee[M split the state space into M disjoint subsets, one for each function, i.e., S = S and Sj nSy = 0, j.j' E [M], j # j', and gives back the associated return g I., namely:
5 (5 |]
(12)
These features are bounded by the maximum return Guax. while the empirical Rademacher complexity over N samples of returns {(8i, [E)N can be directly computed as in Clayton 2014]:
Rx = Gmax
Σ JE(M]
(((S))
(13)
where P(S,) = N5/N and N₁ = (( 8;) 8, E Sy. i E [N]H- The decomposition of the Rademacher term into single terms leads to rewriting B(,,F,N, in the following lemma. Lemma 1. For Distributional Max-Ent Evaluation with alate-aggregation feature class, the variance term B ((,, F,N is given by;
B(,,F,N,0) =10||||1 G max
Σ SE(M]
(P(S,) +
log 1/8
(14)
2N
4.1 Representation Refinement: Progressive Factorization
State aggregation features are of interest due to the possibility progressively refining the repre- sentation by increasing the factorization level, that is, by splitting a subset S, into further disjoint subsets. This refinement is called progressive factorization and defined as follows.
Definition 1 (Progressive Factorization). For sets of state aggregation feature func- tions, F,F,, we say that F, is a progressive factorization of F. iie, F C F,, if F (fy,FF {fi...., ffi.....f-11/ f5-1, (.............flll additional functions ]E[KK] are such that the corresponding subsets satisfy
{fi...., fy-1....f-11:ff fst1...., }
[KE]K] and the
S S n = 0. k. E [K], # kE[K]
where only non-degenerate class factorizations will be considered, meaning that the new subsets St are non-empty.
It is relevant for interests that, in the case of progressive factorizations F C. F', the respective Max-Ent solutions enjoy the following monotonicity property
Lemma 2 (Monotonicity). The multipliers of the Max- Ent solutions 7, N using F C F' are such that
|| < 1|
(15)
This result fully derived Appendix and it ensures a monotonically increasing of terms contained in the variance term of Eq. 11 since the complexity term is monotonically increasing by definition. On the other hand, the bias represented by Eq. (10) guaranteed to decrease monotonically at finer levels of factorizations.

---

4.2 D-Max-Ent Progressive Factorization Algorithm
In summary, D-Max-Ent PE shows a general- ization enror bound whose quantities are either known or estimated and change monotonically between progressive factorizations. On these results, we build an algorithm called D-Max- Ent Progressive Factorization, shown in Algorithm 2 which iteratively constructs a sequence of fea- ture sets Fo Fi ... with progressive fac- torization while performing PE. The behavior of the algorithm is similar to what is done in Struc- tural Disk-Minimization (SRM) [Vapnik] 1991] and it involves optimizing for trade-off: the bias term (i.e., empirical risk) decreases by tak- ing into account more complex features classes,
Algorithm 2 Distributional Max-Ent Progressive Factorization
Require: ,HN, 8, 3, K)
--trajectory
samples, initial feature set, confidence level, boosting factor, factorization factor
1: Done f False, i" <- while not Done do
2:
3:
F + Fir, M < |
4:
n f- D-Max-EMi Ent PE(HN..)
5:
[(i) < BL()) + B(A,F,N
6:
(F, (EMM)) Progressive Factor(F, K)
7:
for j E [M] do
8:
fi < D-Max-Ent PE(HN, Fj)
while the variance term (i.e.. the confidence in- terval) increases. The whole algorithm then based on the progressive search for the new set of feature functions which reduces a proxy of the generalization error bound of D-Max-ERi Ent
9:
[(i ) < BL((;) + B(ř F,, N,
if ((1))) ) ((i) then 1" {
11:
12:
end end for
13:
PE:
14:
if F. = then
15:
Done <- True
+ 3
(16)
16: end
and the procedure will continue until there are no further improvements in the trade-off. Due to the nature of the proxy function, the role of 0 is regulate the tendency to factorize. Higher values of 33 will increase the magnitude of the decreasing term, causing a boost in the tendency to factorize. On the other hand, lower values will further decrease the importance of this term, resulting in lower tendency factorization.
17:
end while
18: return 1):
Finally, the Progressive Factor function takes as input the list of feature functions and factor K and returns a list of progressively factored set of feature functions. More specifically, each element in (F)))E[M] corresponds to a progressive factorization of the feature fy. factoring the related subset Sj into K disjoint subsets in Definition 1 The new K subsets {S k<(K)
[K) constructed in the worst-case scenario: the complexity term in Eq. (13) is maximized with partitions of a set leading to uniform distribution of samples in each new partitioned subset, and since is not possible to know in advance which samples will be contained in which new subset, one way is then to proceed with a uniform factorization. We decided maintain the most agnostic approach over the set of possible features, but prior knowledge could be used to narrow down the partitions to consider.
5 Illustrative Numerical Simulations
This section reports the results of some illustrative numerical simulations that make use of Algorithm
Simulations Objectives The objective of the simulations is to illustrate two essential features of the proposed method that were only suggested by the theoretical results. First of all, to analyze the outcome of performing policy evaluations with aggregated states different sample regimes, by comparing the output of the proposed algorithm with some relevant baseline distributions. Secondly, the aim is to study the role of the boosting parameter / and the sampling regime N. being the main hyper- parameters of Algorithm 2 in the tendency to factor the representation at utterly different sample regimes.
MDP Instance Design The effectiveness of the proposed approach expected to be particularly evident in MDPs admitting factored representation of the return distribution, namely the ones in which many states are nearly equivalent under the evaluation a policy. This factorizability property is not uncommon in general since it is present in any environment with symmetries and Block-MDPs
7

---

A-
A-
300
2.00
A- 152
A- 152
19
4.0
30
57-
A-
30
8-
150
320
setoriationn Chops
Tastoreattenn
Figure 1: Bound Trend for different 3 (N = 50)
Figure 2: KI. Trend for different B (N = 50)
p-
04
A=
3,
32
0.550
A=
0.50
18
0.50
2,tot
12
Factorization Steps
Factorization Steps
Figure 3: Bound Trend for different | (N = 1000)
Figure 4: KL. Trend for different / (N = 1000)
[uu et 2019] as well. The MDP instance is then designed to be a Block-MDP indeed since it allows for better evaluate the simulation objectives: one would expect that operating on MDPs admitting a factored representation would allow for lower values of 3 be effective enough, while higher level of boosting would force over-factorf fverrfactorizations that are unnecessary, leading to no further improvement or even degradation of the results. The simulations are run on a rectangular GridWorld, with height of 4 and length of 8, with traps on the whole second line and goals all over the top. A visualization of the setting can be seen Appendix D The policy is selected as a uniform distribution over the set of actions A = (up, left, right).
Performance Indexes The proposed MDP instance presents many upsides in terms of the inter- pretability of the output as well. First of all, it allows us to directly compute the true underlying return distribution with Monte-Carlo estimation. Secondly, it permits to compare of the output distribution of the algorithm with the result of performing plain Distributional Max-Ent Policy Evaluation (Algorithm T) with two baseline representations: an oracle factorization that aggregates together states known to be equivalent under the policy, and in particular all the upper and lower states respectively: full factorization that employs IS -singletons of states as representations, i.e.. the most fine fine-grained representation possible. The comparison is made via two relevant quantities, the KL. divergence with respect to the true distribution (the bounded quantit), and the total bound .) (the bounding quantity). Finally, the value of the partition splitting K is set to 2, to reduce the exponential search space of all possible uniform partitions, the discount factor y is set to 0.98 and the confidence ö to 0.1, the results are averaged over 10 rounds with the respective standard deviation.
Bio = Σ(n) +
Results Discussion The results of the simulations are reported from Fig. D to Fig. with the quantity related to the oracle parametrization being in orange, while the ones related to the full parametrization being in blue. possible to notice that these two distributions have almost the same KL divergences with respect to the true return distribution (Fig. 4) yet they highly differ
8

---

in the bound B (Fig. 1 mostly due to the variance term, which is way higher in the case of full factorization. This suggests that the bound indeed able distinguish between the two. The plotting of the outputs of Algorithm 2] stops at the optimal number of factorization steps found for different values of 8, namely at Fix. The plots should be read as follows: while the bound term Bot is expected to increase at each factorization step, the KI. divergences with respect to the true return distribution should decrease as much as possible. In all cases, it is evident that the value of B pushes towards higher number of factorization steps, going from performing no factorization at all using low values (3 = 3), to performing up to 4 factorization steps even this simple scenario with higher values (B 450), both at low and high sample regimes (NV E {50, 1000)). Furthermore, at higher sample regimes, it is possible to see how the higher quality of the estimation counteracts the action of B, and increasing it generally induces still fewer factorizations compared to the low sample regimes with same values of 3, as in Fig. 3 Finally, it is apparent that minimizing for Eq. (16) successfully decreases the KI. divergence. Nonetheless, its values stop decreasing significantly after the first factorization, which splits the state space over the two rows and further factorizations might lead to performance degradation well.
6 Discussion
In this section, we briefly discuss the literature related to this work and provide some concluding remarks about the results and future research paths.
6.1 Related Works
Our work relates multiple fields. We now highlight the most relevant connections, while an exhaustive overview is beyond the scope of this paper.
Distributional Reinforcement Learning D-RL has recently received much attention. both for the richness of flavors it admits [Rowland et al. 2019, 2021], and the surprising empirical effec- tiveness [Bellemare et al., 2020). Our work tries to answer different research questions compared to traditional policy evaluation in D-RL and applies completely different techniques to derive the quantities of interest. Firstly, Bellman Operators and consequently contraction arguments cannot be applied since the estimation process is Monte-Carld based. In this way, the sequential nature of the problem is not exploited, but we show that density estimation techniques do offer interesting properties nonetheless. Additionally, the employed indexes differ from traditional D-RL results. For example, the most recent bound for Q-TD (Rowland al. 2023], , shows a sub-linear term but the bound is made over the maximum Wasserstein distance, which cannot be directly related to the KL-divergence without further assumptions. Finally, distributional considerations have been employed in the field of function approximations as well. However, to the best of our knowledge, no other D-RL works explicitly address representation learning.
Maximum Entropy and Feature Selection Max-Ent methods have long and extensive literature in the density estimation field, which mostly focused the general and algorithmic aspects of the method [Barron and Sheu 1991 Dudík et al. 22004] /Sttter et al., 2017). Among the others. (Cortes) al. [2055] proposed a Max-Ent regularized formulation for performing feature selection. Their method allocates different weights to different features to achieve an even better trade-off based on combination. Due to this, their work differs from ours in the nature of the search space, which is not built progressively but is defined a priori. Additionally, their generalization bound is of the same nature as standard Max-Ent bounds and contains supxeol |A|| term, which is bounded yet unknown. Finally, Mavridis et al. (2021] perform progressive state-aggregation through Max-Ent methods, but they try to optimize a different objective function based on -ate-dissimillrity.
6.2 Conclusions and Future Works
In our work, we presented in D-RL framework a new policy evaluation approach based on Maximum Entropy density estimation, called Distributional Max-Ent Policy Evaluation, which benefits from the learning guarantees of Max-Ent and the generality of the setting, being able to enforce even complex feature families. We extended previous results and derived a practical formulation of the generalization error bound, which contains only estimated and known quantities of the problem. We then instantiated particular class of features, namely state aggregation, and we proposed an algorithm

---

called Distributional Max- Ent Progressive Factorization to adaptively find feature representation that optimizes for a proxy of the generalization error bound in a Structural Risk Minimization fashion. In this way, we showed that performing PE can indeed drive the learning of reduced-dimentt representation in the distributional setting. We then provided illustrative simulations showing the empirical behaviors of these approaches, while clarifying the links between some hyperparameters and the sample regime. Much of our analysis and theoretical guarantees straightforwardly extend to other feature classes, and open question is to investigate other instances of features and settings that can benefit from the proposed framework. Future works will focus on the interaction between Temporal Difference (TD) distributional methods and representation learning in the proposed setting and on the existence MDP instances that enjoy some relevant properties in the bias/variance trade-off along successive factorizations, leading to high performance or better error bounds.
Acknowledgments
This paper is supported by PNRR-PE-AI FAIR project funded by the NextGeneration EU program.
References
Andrew R. Barron and Chyong-Hwa Sheu. Approximation of Density Functions by Sequences of Exponential Families. The Annals of Statistics, 19(3), September 1991. ISSN 0090-5364. doi:
Marc G. Bellemare, Salvatore Candido, Pablo Samuel Castro, Jun Gong, Marlos C. Machado, Subhodeep Moitra, Sameera Ponda, and Ziyu Wang. Autonomous navigation of stratospheric balloons using reinforcement learning. N 588(7836):77-82 December 2020. ISSN 0028- 0836, 1476-4687. doi: 10.1038/641586-000-999
Mare G. Bellemare, Will Dabney, and Mark Rowland. Distributional Reinforcement Learning. MIT Press, 2023.
Daniel Brown, Scott Nickum, and Marek Petrik. Bayesian robust optimization for imitation learning. Advances in Neural Information Processing Systems, 33:2479-2499 2020.
Yinlam Chow, M. Ghavamzadeh, Lucas Janson, and M. Pavone. Risk-cons reinforcement learning with percentile risk criteria. Journal Of Machine Learning Research, 2015.
Scott Clayton. Lecture notes (topic 10) of eecs 598: Statistical learning theory, 2014.
Corinna Cortes, Vitaly Kuznetsov, M. Mohri, and Umar Syed. Structural Maxent Models. July 2015.
Shi Dong, Benjamin Van Roy, and Zhengyuan Zhou. Provably Efficient Reinforcement Learning with Aggregated States, February 2020. arXiv:1912.06336 [cs, math, stat].
Simon S Du, Akshay Krishnamurthy, Nan Jiang, Alckh Agarwal, Miroslav Dudik, and John Langford. Provably efficient RL with Rich Observations via Latent State Decoding. 2019.
Miroslav Dudík and Robert E. Schapire. Maximum Entropy Distribution Estimation with Generalized Regularization. In Learning Theory, volume 4005, pages 123-138, Springer Berlin Heidelberg, Berlin, Heidelberg, 2006. ISBN 978-3-540-35294-5 7778---540-35296-, doi: 10.1007/11774400,
12. Series Title: Lecture Notes in Computer Science.
Miroslav Dudík, Steven J. Phillips, and Robert E. Schapire. Performance Guarantees for Regularized Maximum Entropy Density Estimation. In Learning Theory, volume 3120, pages 472-486. Springer Berlin Heidelberg, Berlin, Heidelberg, 2004. ISBN 978-3-540-22282-8 978-3-540-278199- doi:
10.1007/978- 3-540-27819.
33. Series Title: Lecture Notes in Computer Science.
R. M. Dudley. Uniform Central Limit Theorems. Cambridge University Press, edition, July 1999. ISBN 978-0-521-461022 978-0-521-05221-4 978-0-511-665523 CB09780511665662
doi: 10.1017/
Ramtin Keramati, Christoph Dann, Alex Tamkin, and Emma Brunskill. Being Optimistic to Be Conservative: Quickly Learning a CVaR Policy. Proceedings of the AAAI Conference on Artificial Intelligence, 34(04):4436-44)) April 2020. doi: 10.
10

---

V. Koltchinskii and D. Panchenko. Empirical Margin Distributions and Bounding the Generalization Error of Combined Classifiers. The Annals of Statistics, 30(1):1 50, 2002, doi: 10. 12/4/s// 1015362183.
Christos N. Mavridis, Nilesh Suriyarachchi, and John S. Baras. Maximum-Entropp Progressive State Aggregation for Reinforcement Learning. In 2021 60th IEEE Conference on Decision and Control (CDC), pages 5 1114-551400 5149, Austin, TX, USA, December 2021. IEEE. ISBN 978-1-66543-659. doi: 10.1099CDC444444222...6666
Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of machine learning. Adaptive computation and machine learning. The MIT Press, Cambridge, Massachusetts, second edition edition, 2018. ISBN 978-0-262-03940--6
Ofir Nachum, Yinlam Chow, Bo Dai, and Lihong Li. DuaIDICE: Behavior-Agnostie Estimation of Discounted Stationary Distribution Corrections, November 2019. URL http://arxii org/ abs/1906. 0t01//006607733.
04733, 10XXv:1906.0470 [es, stat].
Martin L. Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming. USA, 1st edition, 1994. ISBN 0471619779.
Mark Rowland, Robert Dadashi, Saurabh Kumar, Rémi Munos, Mare G. Bellemare, and Will Dabney. Statistics and Samples in Distributional Reinforcement Learning. 2019. doi: 10. A0.48V50/ARXIV. 1902.08102.
Mark Rowland, Shayegan Omidshafiei, Daniel Hennes, Will Dabney, Andrew Jaegle, Paul Muller, Julien Pérolat, and Karl Tuyls. Temporal Difference and Return Optimism in Cooperative Multi- Agent Reinforcement Learning. page 9, 2021.
Mark Rowland, Rémi Munos, Mohammad Gheshlaghi Azar, Yunhao Tang, Georg Ostrovski, Anna Harutyunyan, Karl Tuyls, Mare G. Bellemare, and Will Dabney. An Analysis of Quantile Temporal- Difference Learning, January 2023. arXiv:2301.0440 [cs, stat].
Satinder Singh, Tommi Jaakkola, and Michael Jordan. Reinforcement Learning with Soft State Aggregation. In Advances in Neural Information Processing Systems, volume 7. MIT Press, 1994.
Tobias Sutter, David Sutter, Peyman Mohajerin Esfahani, and John Lygeros. Generalized maximum entropy estimation. 2017. doi: 10.48500ARRXI..178.0075
Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Csaba Szepesvári. Algorithms for Reinforcement Learning. Synthesis Lectures on Artificial Intelli- gence and Machine Learning, 4(1):1-103, January 2010. ISSN 1939-4608, 1939-4616.
Aad W. van der Vaart and Jon A. Wellner. Weak Convergence and Empirical Processes. Springer Series in Statistics. Springer New York, New York, NY, 1996. ISBN 978-1-4757-2547-6699 978-1- 4757-2545-2. doi:
11177799788-477--22
Benjamin Van Roy. Performance Loss Bounds for Approximate Value Iteration with State Ag- gregation. Mathematics of Operations Research, 31(2):234-244 May 2006. ISSN 0364-765X, 1526-5471. doi: 10.1288/morr. 1060.0188.
V. Vapnik. Principles of Risk Minimization for Learning Theory. In Advances in Neural Information Processing Systems, volume 4.
1991.
Martin J. Wainwright and Michael I. Jordan. Graphical Models, Exponential Families, and Variational Inference. Foundations and Trends Machine Learning, 1:1 1:1-305, 2007. ISSN 1935-8237, 1935- 8245. doi: 10.1561/2200000001.
Shaojun Wang, Russell Greiner, and Shaomin Wang. Consistency and Generalization Bounds for Maximum Entropy Density Estimation. Entropy, 15(12):54395443, December 2013. ISSN 1099-4300.. doi: 1.339///5.115433
Bin Yu. Rates of Convergence for Empirical Processes of Stationary Mixing Sequences. The Annals of Probability, 22(1):94-116, 1994. ISSN 0091-1798. Publisher: Institute of Mathematical Statistics.


Reviews:
[{"review_id": "kWc7GJVZxe", "reviewer": "Reviewer_Tf1V", "summary": "The paper introduces a new framework called Distributional Maximum Entropy Policy Evaluation (D-Max-Ent PE) for policy evaluation in a distributional reinforcement learning setting. The framework considers the complexity of the representation used to evaluate a policy and incorporates it into a generalization-error bound. The authors then propose an algorithm called D-Max-Ent Progressive Factorization, which uses state-aggregation functions to progressively refine the state space representation while balancing information preservation and reduction in complexity. Numerical simulations demonstrate the effectiveness of the proposed algorithm and its relationship with aggregations and sample regimes.", "questions": "None", "limitations": "None", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper derived a practical formulation of the generalization error bound and proposed an algorithm, Distributional Max-Ent Progressive Factorization, that adaptively finds a feature representation to optimize the generalization error bound. Through illustrative simulations, the authors demonstrated the empirical behaviors of these approaches and explored the relationships between hyperparameters and the sampling regime. ", "weaknesses": "1. This paper lacks experiments in practical reinforcement learning tasks, such as Atari. The authors should conduct more experiments to demonstrate the efficiency of the proposed policy evaluation method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "*This paper lacks experiments in practical reinforcement learning tasks, such as Atari. The authors should conduct more experiments to demonstrate the efficiency of the proposed policy evaluation method.*\n\nWe thank the reviewer for this feedback. We remark that the main contribution of our work is of theoretical nature. Indeed, **the objective of the simulations was to illustrate two essential features of the proposed method that were only suggested by the theoretical result**, so as to close all the topics directly or indirectly introduced in the theoretical study:\n\n- whether it is possible to boost factorization by tuning the value of $\\beta$,\n- whether the resulting factorization aggregated similar states in terms of the policy's true return distribution;\n\nrather than demonstrate its applicability to general reinforcement learning tasks. The MDP instances we employed were designed to be coherent with the objectives of the simulations and to be as much interpretable as possible in terms of true return distributions, a feature which is hardly shared with more complex RL tasks (e.g., Atari). We believe that the adaptation of the proposed method to more practical contexts is out of the scope of the present paper."}, {"review_id": "vYPgrUWNu1", "reviewer": "Reviewer_W4d8", "summary": "This paper takes a maximum entropy based approach to learning the return distribution of a policy. This framework learns a distribution with maximum entropy subject to matching the expectation of a certain set of functions, referred to as the feature functions. They adapted a classical maximum entropy error bound to measure the approximation error of their learnt distribution to the true return distribution in the KL divergence. They propose a framework where the features are based on underlying state abstractions, and propose an algorithm which progressively learns a return distribution and underlying features. They then evaluate their algorithm in various small scale experiments.", "questions": "- Is there a way to connect this method to the underlying RL problem? \n- In Algorithm 1, how is $\\mu_j(\\mathcal{H}_N )$ evaluated? To my knowledge it seems we require knowledge of the true return distribution $\\eta^\\pi$ (or at the very least samples from it). ", "limitations": "The authors discuss limitations, such as their method requiring access to IID samples.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- Using a maximum entropy approach to learning the return distribution is an interesting direction, and novel from the previous parametric approaches used. \n- This paper takes a step towards studying the interaction of distributional RL and representation learning.\n- The theoretical results are cleanly proved, with very clear and organized arguments.", "weaknesses": "- The paper takes an RL agnostic approach, and learns an approximate return distribution without taking into account the underlying RL decision problem (i.e. Bellman operators aren't used, nor reward nor transition information, etc). This seems contrary to the decision-aware RL direction, which broadly states that we should focus learning capacity on aspects which are important for the decision problem. Some guarantees that learning a model in this way comes with some benefits for the RL problem would improve the contribution.\n- To my knowledge, it appears that Algorithm 1 assumes knowledge of the true return distribution (or at least access to samples from it), which is a potentially unrealistic assumption. I may be misunderstanding however, so I listed this as a question below. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the important points outlined. \n\n*The paper takes an RL agnostic approach, and learns an approximate return distribution without taking into account the underlying RL decision problem (i.e. Bellman operators aren't used, nor reward nor transition information, etc). This seems contrary to the decision-aware RL direction, which broadly states that we should focus learning capacity on aspects which are important for the decision problem. Some guarantees that learning a model in this way comes with some benefits for the RL problem would improve the contribution.\nIs there a way to connect this method to the underlying RL problem?*\n\nThe proposed method is agnostic with respect to the nature of the underlying decision problem, which has advantages and disadvantages. Indeed, as the reviewer noted, we do not make use of Bellman operators, but just of return distributions, as commonly done in Monte Carlo approaches to RL. **While we agree that this might overlook some of the MDP structure, it has the advantage of leading to (i) a simpler mathematical treatment and (ii) of being applicable even to non-Markovian environments**. As the reviewer noted, we see this work as the first contribution taking \"a step towards studying the interaction of distributional RL and representation learning\". **Future works should for sure investigate the interaction between Temporal Difference (TD) distributional methods and representation learning**.\n\n\n*To my knowledge, it appears that Algorithm 1 assumes knowledge of the true return distribution (or at least access to samples from it), which is a potentially unrealistic assumption. I may be misunderstanding however, so I listed this as a question below.\nIn Algorithm 1, how is $\\mu_j(H_n)$ evaluated? To my knowledge it seems we require knowledge of the true return distribution (or at the very least samples from it).*\n\nThe proposed algorithm does not require to evaluate $\\mu_j(H_n)$, but rather $\\hat \\mu_j(\\mathcal H_n)$, that corresponds to the empirical estimate of the value of the feature function $f_j$ as for Equation 6, computed from $N$ trajectory samples. In this sense,** it indeed requires trajectories of experience coming from the policy which is being evaluated but not the full knowledge of the environment transition models (and so of the return distribution)**.\n\n"}, {"review_id": "ktFIUNtRnU", "reviewer": "Reviewer_HRB1", "summary": "The authors proposed a novel Maximum Entropy framework for policy evaluation in a distributional RL that can explicitly take into account the features used to represent the state space while evaluating a policy. Based on the framework, the authors developed  D-Max-Ent Progressive Factorization algorithm balancing the trade-off between bias and variance of the representation space.", "questions": "I see the experiments are conducted on a Grid task, will the proposed algorithm also applicable on continuous environments? Will the algorithm help with performance, training stability, sample efficiency? \n(Again, I'm not familiar with representation learning and distributional RL so let me know if my question is not the main concern in the field.)\nMaybe add some legends and more explanations about the figures to help present the results and emphasize the contribution. ", "limitations": "Experiments designed in the paper help to support some theoretical analysis of the proposed algorithm, but I'm not sure about the benefits/advantage of the algorithm in a wider range of tasks based on my understanding of the paper. ", "rating": 5, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Overall well written and well organized paper. The authors answered two interesting research questions of representation learning. The authors analysed the questions both theoritically and experimentally by proposing a novel maximum entropy framework for policy evaluation and a new algorithm based on the framework. ", "weaknesses": "Experiment setting and results presentation are a bit hard to understand. I'm not very familiar with representation learning and distributional RL, so maybe this is standard but I find the figures hard to interpret, so as the importance of the contribution of the paper based on the experiments. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the points outlined, in particular the ones about the readability of the figures.\n\n*I see the experiments are conducted on a Grid task, will the proposed algorithm also applicable on continuous environments? Will the algorithm help with performance, training stability, sample efficiency?*\n\n\nIndeed, the **proposed method works with continuous spaces as well**, yet the factorization rule is potentially less intuitive since it needs to cut continuous spaces. As for your second question,  **reducing the input space dimension to the decision-making problem would help sample efficiency** for sure, once it assures that the new representation does not introduce bias, as the proposed method does. On the other hand,  how to combine the proposed framework with existing methods in order to boost performance was left for future work.\n\n*Maybe add some legends and more explanations about the figures to help present the results and emphasize the contribution.*\n\n\nDue to space limitations, the description of the experiments was limited to the core, as for the figures which were made to contain all the needed information, but **we intend to extend this whole section in the additional page of the Camera Ready version**.\n"}, {"review_id": "GgJSzcyyMm", "reviewer": "Reviewer_SYd5", "summary": "This paper proposes two algorithms, distributional maximum entropy policy evaluation (D-Max-Ent PE) and D-Max-Ent Progressive Factorization.\n\nThe first algorithm applies the formulations of maximum entropy RL to the problem of distributional policy evaluation, while the second extends the first by adding a progressive state space factorization step that recursively decomposes/factorizes the state space by evaluating the maxent PE at each resolution to find a decomposition that minimizes the generalization error bound.\n", "questions": "I don't see distributional policy evaluation defined explicitly anywhere (though I think it's mostly covered by equation 3 and surrounding material?), unlike maxent RL which is defined in detail. For clarity this seems like it would be nice to have.\n\nI like the use of specific grounding questions in the introduction nice, but I found the wording of the questions to be kind of leading. They seem phrased in such a way that the fact of asking them means the answer will be \"yes\" which feels not very useful rhetorically. Perhaps they can be reworded to be more thought provoking?\n\nAbsent further experiments, it'd be nice to have some analysis and discussion of what regimes these algorithms will tend to perform well in versus struggle in. For example, what happens to the test case when the number of states increases, in X versus Y dimension? Are there particular properties of an MDP that will lead to high performance/better error bounds?", "limitations": "The core of the limitations of this work are covered in my discussion of weaknesses. The authors are pretty clear about the extent of the contributions and the paper is pretty clear with its claims, however, which is always nice to see.\n\nI don't see any potential negative societal impact issues stemming from this work.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I found this paper to be generally well written and understandable. A few grammar issues are present, but nothing that seriously hurts legibility, and the math in the main paper flows intuitively.\n\nDistributional maxent policy evaluation makes sense as an approach to PE. I'm not familiar with recent literature on PE methods, but it seems like a good approach to the problem, and maxent methods have had enough empirical success in deep RL work that it seems like a reasonable approach here too.\n\nI thought combining maxent PE with an iterative state space decomposition to be a clever and intriguing idea as well. On paper it makes sense to me as a way to divide and conquer large state spaces, and I'm curious how this idea could extend to more complex tabular and continuous MDPs (as they authors note it should). \n\nOverall this seems like an interesting theoretical approach with appealing generality and the potential for extension to more complex domains.", "weaknesses": "The main issue I found with this paper is the relative simplicity of the domain studied, especially in the limited empirical evaluation.\n\nI recognize this is a theory paper, and many significant practical questions are yet to be addressed. In particular the problem of how to factorize a state space is very hard in general and a worthy topic in its own right (the authors do acknowledge this). However, I'd still like to see some more analysis of what sorts of MDPs the proposed algorithm is strong in versus weak using just the simple factorizations discussed here, ideally in the form of additional experiments on other gridworld MDPs besides the single example case. I think this idea is interesting, so as an experimentalist I'd like to see its practical behavior demonstrated a little more.\n\nI'm still inclined to recommend acceptance as the core ideas seem well motivated and interesting, but I think this would be a stronger paper with some more validation.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thorough review and for the constructive comments. \n\n*I don't see distributional policy evaluation defined explicitly anywhere (though I think it's mostly covered by equation 3 and surrounding material?), unlike maxent RL which is defined in detail. For clarity this seems like it would be nice to have.*\n\nDistributional PE was indeed left implicitly defined, but a **brief definition of the PE problem will be for sure added.**\n\n*I like the use of specific grounding questions in the introduction nice, but I found the wording of the questions to be kind of leading. They seem phrased in such a way that the fact of asking them means the answer will be \"yes\" which feels not very useful rhetorically. Perhaps they can be reworded to be more thought provoking?*\n\nThe questions were mostly used to drive the discussion toward the desired direction, but as you suggest a more provoking wording might help the rhetoric of the whole work. We would propose as a rewording something like:\n- Q1:  Does employing return distributions offer other tools out of the standard distributional RL methods?\n- Q2: How are representation learning and policy evaluation intertwined? Do distributional methods offer a new way to highlight and exploit this connection?\n\n*Absent further experiments, it'd be nice to have some analysis and discussion of what regimes these algorithms will tend to perform well in versus struggle in. For example, what happens to the test case when the number of states increases, in X versus Y dimension? Are there particular properties of an MDP that will lead to high performance/better error bounds?*\n\nThis point is an extremely interesting point indeed; we started addressing it (261-263) and left as a future contribution (337-338).  We will for sure add a section in the Appendix better explaining these considerations:\n   - In the considered MDP instances, increasing X or Y would affect the results based on which kind of change they would introduce in the return distribution landscape. In case the change would be consistent with the current factorization (for example, adding over the X dimension with the same grid layout) the results would follow; in general, **as far as the added dimensions would maintain the same factorization structure in the return distribution, the results would follow.**\n   - A factorizability structure was introduced but not directly addressed. In general, similarly to what is done in \"On the Complexity of Representation Learning in Contextual Linear Bandits\", Tirinzoni, Pirotta, Lararic (2023),  we believe that **a condition of \"realizability\" of the state-aggregation feature class for $\\eta^\\pi$ would be required**, or at least in an approximate formulation. Anyway, for state-aggregation feature classes in particular, the **high-performance MDP instances the reviewer mentioned would be the ones where acting with one policy would induce approximately the same return distribution in many different states**, which are the ones that will end up being aggregated."}, {"review_id": "aH2KmWK3HR", "reviewer": "Reviewer_wzKJ", "summary": "This paper investigates policy evaluation of distributional RL by leverage of Maximum entropy density estimations without explicitly considering Bellman operator and contraction mapping. They further derived the new generalization error bound in the maximum entropy PE process. For a practical algorithm, authors proposed progressive factorization to refine the representation with desirable properties, including the monotonicity. They did some experiments to demonstrate their claims.\n\n", "questions": "To the best of our knowledge, the representation issue in (distributional) RL is not well-defined yet. Without discussing what the representation issue you are going to study, directly incorporating an approach is very confusing to me. Also, there are three major issues I would raise here.\n\n**1.Motivation.** What is the representation issue in distributional RL? Is it just the representation of the value distribution? Note that people may also think distributional RL is also able to improve the representation of environments by additionally leverage of return distribution information. To be honest, after carefully the whole paper, I am still confused why your proposed methods are beneficial for representation. If authors refer to the monotonicity properties of progressive factorization, are there any experiments on real environments, e.g., Atari games, to show these benefits? \n\n**2.Less discussion within RL context**. The authors claim the proposed method can handle the representation issue in RL context, i.e., policy evaluation, but they did not rigorously prove the contraction properties of Bellman operator under the maximum entropy density estimate. Note that as an RL research paper, the discussion of contraction properties is the foundation of variants of policy evaluation algorithms. Otherwise, it would be technically problematic. In line 49, the authors mentioned PE in a distributional setting is directly linked to the representation issue, but as far as I can tell, they are not equivalent. Discussing the contraction mapping properties is necessary. \n\n**3.Experiments are weak**. If authors claim they are handling the representation issue of distributional RL, I do not find any direct experiment to demonstrate it on commonly used environments.  Applying the state-aggregation approach in practice is also problematic as it hinders the generalization for RL with function approximation.  In line 128, the authors claimed that Max-ent has multiple advantages. Have authors demonstrated them in their experiments? Similar issues also apply to the monotonicity property in Line 205.\n\n\nOverall, I do not think this paper really handles the representation issue in the real distributional RL setting. Some preliminary results are given without connection with Bellman equation and RL context. Experiments are also weak to support their ambitious targets mentioned in the introduction.\n", "limitations": "No", "rating": 2, "confidence": 4, "soundness": 1, "presentation": 2, "contribution": 1, "strengths": "* Introducing Maximum entropy density estimate into distributional RL has not been studied.\n* Progressive factorization seems to refine the representation of RL.\n* The writing is basically clear.\n", "weaknesses": "* The motivation is not clear. Why should we consider the maximum entropy estimate to study the representation?\n* Without discussing the RL context in the proposal, e.g., Bellman operator and contraction mapping, the technical soundness of the so-called max-ent approach is questionable.\n* Experiments are weak and not directly related to the distributional RL setting.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thorough review.  We will answer proceeding by points:\n\n*What is the representation issue in distributional RL? Is it just the representation of the value distribution? Note that people may also think distributional RL is also able to improve the representation of environments by additionally leverage of return distribution information. To be honest, after carefully the whole paper, I am still confused why your proposed methods are beneficial for representation. If authors refer to the monotonicity properties of progressive factorization, are there any experiments on real environments, e.g., Atari games, to show these benefits?* \n\nAs for the high-level representation-learning definition, we referred to it rather informally mostly because the representation issue is in general not well-defined, but we will for sure add an explicit description of what we mean by representation-learning in this context:  it consists of **finding a good feature representation of the decision-making space so as to make the overall learning problem easier**. In this sense, **a method able to reduce the dimensionality of the problem is said to address the representation-learning problem.** In this sense, the representation-learning issue is the same in distributional RL as in expected-value RL, and canonical (distributional) RL methods indirectly address it by means of generic function approximation. Nonetheless, **distributional methods further exacerbate this problem with respect to expected-value RL**, since they require learning distributions over (usually) large spaces. For this reason, we believe that investigating how to use distributions of returns to alleviate the issue is an interesting stream of research.\n\n*The authors claim the proposed method can handle the representation issue in RL context, i.e., policy evaluation, but they did not rigorously prove the contraction properties of Bellman operator under the maximum entropy density estimate. Note that as an RL research paper, the discussion of contraction properties is the foundation of variants of policy evaluation algorithms. Otherwise, it would be technically problematic. In line 49, the authors mentioned PE in a distributional setting is directly linked to the representation issue, but as far as I can tell, they are not equivalent. Discussing the contraction mapping properties is necessary.*\n\nUp to our knowledge, contraction considerations are indeed a fundamental tool in the RL methods that enforce the MDP structure (i.e. TD-based), as canonical distributional RL methods devised so far. On the other hand,  **contraction considerations are meaningless for Monte-Carlo-based methods, since no operator is actually defined**; our work is indeed in this stream of work, and **this motivates the absence of contraction considerations in the analysis**. We will for sure add one comment giving reasons for the absence of contraction arguments, but more generally we do not agree with the claim that \"Discussing the contraction mapping properties is necessary\", we rather believe that  \"**Discussing the contraction mapping properties is necessary for operator-based methods**\", which is not the case.  We point out that a large portion of sound and effective RL algorithms (including policy gradients) are based on Monte Carlo simulations and do not make use of Bellman operators. While we agree that Mote-Carlo methods might overlook some of the MDP structure, they have the advantage of leading to (i) a simpler mathematical treatment and (ii) of being applicable even to non-Markovian environments.\nAs for the final part of the question, the aim of the paper was indeed to show that policy evaluation and representation learning were linked, and in fact, **we showed that performing PE can drive the learning of a reduced-dimension representation**. Yet, **the fact that they were equivalent was not among the objectives of our work**. We can further specify this if the reviewer thinks some parts of the exposition were misleading.\n\n*If authors claim they are handling the representation issue of distributional RL, I do not find any direct experiment to demonstrate it on commonly used environments. Applying the state-aggregation approach in practice is also problematic as it hinders the generalization for RL with function approximation. In line 128, the authors claimed that Max-ent has multiple advantages. Have authors demonstrated them in their experiments? Similar issues also apply to the monotonicity property in Line 205.*\n\nThe main contribution of the paper is of theoretical nature. For this reason, the purpose of the simulations was not to demonstrate empirically the effectiveness of the method, but  to illustrate two essential features of the proposed method that were only suggested by the theoretical result, so as to close all the topics directly or indirectly introduced in the theoretical study:\n- whether it is possible to boost factorization by tuning the value of $\\beta$\n- whether the resulting factorization aggregated similar states in terms of the policy's true return distribution;\n\nThe MDP instances we employed were designed to be coherent with the objectives of the simulations and to be as much interpretable as possible in terms of true return distributions, a feature which is hardly shared with more complex RL tasks (e.g., Atari). We believe that the adaptation of the proposed method to more practical contexts is out of the scope of the present paper."}]

Author rebuttal:

Meta-review: This paper proposes a novel approach to distributional reinforcement learning (specifically policy evaluation) based on the maximum entropy principle.

The initial reviews for this paper were somewhat mixed. Reviewers generally agreed that the paper is clear, well written, and well organised. There was also general agreement that the core direction pursued by the authors (distributional policy evaluation via maximum entropy principle) is a novel, interesting direction, and that they provide rigorous theoretical analysis of the main algorithm introduced. Several reviewers also commented that this work represents an interesting route into studying the combination of distributional RL and representation learning, an area which is likely to be of interest to both RL and representation learning communities at NeurIPS.

There were also a few concerns expressed across reviews:
 - A lack of thorough experimentation in the paper. While the paper's primary contribution is purely theoretical, it is true that a wider range of experiments would strengthen the contribution. A related concern was that the experiments in the paper focus solely on tabular settings, without investigation of e.g. how applicable the core algorithm is in a setting such as the Atari suite.
 - A relatively RL-agnostic approach. The reviewers contribute a Monte Carlo learning algorithm, which does not make use of bootstrapping, and hence the theoretical analysis does not make use of the standard tools used in previous theoretical distributional RL work, such as Bellman operators.
 - The notion of representation learning used in the paper. The primary focus in the paper is on state aggregation, which is somewhat distinct from the notion of representation learning in deep learning; which is a form of representation learning commonly associated with distributional RL, since improved representation learning is a common hypothesis as to why distributional RL delivers benefits in combination with deep RL.

Different reviewers weighted these concerns differently in their reviews. After discussion, all reviewers except Reviewer wzKJ recommend acceptance for the paper, with the rebuttals also helping to resolve several minor questions raised by reviewers. Reviewer wzKJ has put forward their case clearly: that they weigh the concerns mentioned above highly, and although there has been good discussion between the authors and reviewer on these point, Reviewer wzKJ maintains their rejection recommendation. There has also been substantial discussion among reviewers on these points during the discussion period, as there is a wide range of opinion on the paper.

My overall recommendation for the paper, in light of discussion and review process described above, is for acceptance to NeurIPS. I understand Reviewer wzKJ's reservations, though I am comfortable that though not focused on temporal-difference learning, the authors' contribution represents an interesting approach to policy evaluation in distributional reinforcement learning. I would encourage the authors to incorporate the outcomes of discussion with all reviewers when preparing the camera-ready version of the paper.
Decision: Accept (poster)