Title: Distributional Policy Evaluation: a Maximum Entropy approach to Representation Learning
Venue: NeurIPS 2023
Accepted: yes
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

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.

1 Introduction

In Distributional Reinforcement Learning (D-RL) [Bellemare 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 is 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 et al., 2015, Brown et al., 2020, Keramati et al., 2020]. Most of them varied in how the distribution of the returns is modeled. The choice of the model was shown to have a cascading effect on how such a distribution can be learned, how efficiently and with what guarantees, 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 to find a good feature representation of the decision-making space so as to make the

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 a distributional setting as a particular case of a distribution estimation problem and 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-based constraints acting over the support of the distribution. Such constraints then appear in the generalization-error bound for the Max-Ent problem, with a 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-error 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.

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 state-aggregation 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 RL in 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 a new way to highlight and exploit this connection?

To answer this question, in Section 4, we show that the generalization-error bound changes monotonically when the state-aggregation constraints are changed in a 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 a 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 Markov Decision Processes

A discrete-time finite Markov decision processes (MDP) [Puterman, 1994] is a tuple M := ( S , A , P S , P R , µ, γ ) , where S is a finite state space ( |S| = S ), A is a finite action space ( |A| = A ), P S : S×A → ∆( S ) is the transition kernel, P R : S×A → ∆( R ) is the reward distribution function, µ ∈ ∆( S ) is the initial-state distribution and γ ∈ [0 , 1) is discount factor. 1 A policy π : S → ∆( A ) defines the behavior of an agent interacting with an environment, which goes as follows: starting from an initial state S 0 ∼ µ , an agent interacts with the environment through the policy π , generating a trajectory H = ( S t , A t , R t ) ∞ t =0 , which is a sequence of states, actions and rewards whose joint distribution is determined by the transition kernel, reward distribution, and the policy itself, i.e., A t ∼ π ( ·| S t ) , R t ∼ P R ( ·| S t , A t ) , and S t +1 ∼ P S ( ·| S t , A t ) .

2.2 Value Functions and Distributions of Returns

Given an MDP M with discount factor γ , the Discounted Return is the sum of rewards received from the initial state onwards, discounted according to their time of occurrence:

The Value Function of a given policy π is the expectation of this quantity under the policy itself:

The Return Distribution Function η π of a given policy π is a collection of distributions, one for each state s ∈ S , where each element is the distribution of the random variable G π ( s ) :

where D π ( s ) extracts the probability distribution of a 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 a fixed policy π .

2.3 Maximum Entropy Estimation

Maximum Entropy (Max-Ent) methods [Dudík 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 ∈ ∆( X ) is defined as H ( p ) := -E X ∼ p [log p ( X )] . 2 Additionally, they assume that the learner has access to a feature mapping F from X to R M . In the most general case, we may have M = + ∞ . We will denote by Φ the class of real-valued functions containing the component feature functions f j ∈ F with j ∈ [ M ] . A distribution p is consistent with the true underlying distribution p 0 if

In this case, we say that p satisfies (in expectation) the structural constraints imposed by the features in F . In practice, p 0 is not available and Max-Ent methods enforce empirical consistency over N independent and i.i.d. observations D = { x 1 , . . . , x N } ∼ p 0 with support in X by replacing the definition in Eq. (5) with

The distribution p is said to be consistent with the data D if it matches the empirical expectations. The empirical Max-Ent problem consists then of the following optimization problem

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

where

Max-Ent problem in Eq. (7) is a distribution p λ ∈ ∆( X ) belonging to the class of exponential distributions parametrized by the parameters λ , namely:

where Φ λ := ∫ X exp ( ∑ j ∈ [ M ] λ j f j ( x ′ ) ) dx ′ is a normalization constant, which ensures that p ∈ ∆( X ) , and its log-transformation takes the name of log-partition function A ( λ ) := log ∫ X exp( ∑ j ∈ [ M ] λ j f j ( x )) dx . The log-partition function defines the set of well-behaved distributions Ω = { λ ∈ R M : A ( λ ) < + ∞} . 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 ˆ p to identify p ˆ λ for simplicity.

3 Distributional Policy Evaluation: A Max-Ent Approach

This section aims at answering the first research question Q1. The proposed approach turns distributional PE into a pure density estimation problem in a Max-Ent framework, called Distributional Max-Ent Policy Evaluation , as described in Algorithm 1. For this translation, the algorithm uses the distribution of returns η as p , N -trajectory samples H N = {H} N n =0 as data, and a fixed set of features functions F belonging to a function class Φ . Note that to

Algorithm 1 Distributional Max-Ent Policy Evaluation Require: ( H N , F ) ▷ N trajectory samples, set of features functions ˆ η =argmax η H ( η ) s.t. E X ∼ η [ f j ( X )] = ˆ µ j ( H N ) ∀ j ∈ [ M ] η ∈ ∆( X ) return ˆ η

do this, we need to slightly change the notation concerning the D-RL framework: η will not be a |S| -vector of distributions with support over R , but rather a joint distribution over the whole support X = S × R . Turning PE into a Max-Ent problem has many upsides. First of all, the Max-Ent principle allows to 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 value-function approximation methods [Van Roy, 2006] and as in more recent works based on statistical functionals acting over the return portion R of the support [Bellemare et al., 2023]. One of the possible limitations might be the requirement to have access to a batch of i.i.d. samples, but this is not necessarily restrictive: the result can be generalized for a single β -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 to the distributional PE setting, and in particular, generalization-error bounds defined as KL-divergences. 3 Unfortunately, the generalization error bounds of traditional Max-Ent theory contain a conservative term that compares the solutions of the expectation and empirical Max-Ent problems, ¯ η, ˆ η respectively, by taking the maximum between the 1-norm of the respective multipliers, namely max λ ∈{ ¯ λ, ˆ λ } || λ || 1 . This quantity is bounded yet unknown, making the result unpractical. In the following, we extend the previous results with a more practical bound containing || ˆ λ || 1 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 Φ , which it is such that sup x ∈X ,f ∈F || f ( x ) || ∞ = F < + ∞ and that the

minimum singular value σ min of the empirical covariance matrix of the features ˆ Cov( F ) is strictly positive, namely σ min ( ˆ Cov( F )) > 0 . Then, given a sample batch { x 1 , . . . , x N } ∈ X N of N i.i.d. points drawn from the true distribution η π , for any δ ∈ (0 , 1) , it holds with probability at least 1 -δ that the solution to the sampled Max-Ent problem ˆ η satisfies the following:

where ≾ stands for the fact that the bound comprises additional terms that decrease at a higher rate in sample complexity and were therefore neglected. H ( η π ) and ˜ L (ˆ η ) , the empirical log-likelihood of the solution, form a bias term. The remaining term B ( ˆ λ, F , N, δ ) is a variance term depending on the multipliers characterizing the solution ˆ λ , the number of samples, the confidence level δ , and the feature class complexity as the empirical Rademacher complexity of the class R N (Φ) [Mohri et al., 2018].

3.2 Proof Sketch

Here we report the main steps of the proof of Th. 1. The interested reader can find the complete proof in Appendix A. First, define the set containing the solutions to the expected and sampled Max-Ent problems with S := { ¯ η, ˆ η } , the related set for the multipliers Ω S := { ¯ λ, ˆ λ } , and a quantity that will be central now on h ( x 1 , · · · , x N ) := max η ∈S | E η π [log η ] -1 N ∑ N i log η ( x i ) | . Then, the building blocks of the error term KL ( η π || ˆ η ) , namely KL (¯ η || ˆ η ) and KL ( η π || ¯ η ) are bounded by:

It is possible to show that:

The first inequality is obtained with standard methods as in van der Vaart and Wellner [1996], Dudley [1999], Koltchinskii and Panchenko [2002], Wang et al. [2013]. The second one is 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 D A ( ¯ λ, ˆ λ ) . One can see that the use of the second inequality introduces an additional assumption about the expressiveness of the feature functions, requiring the minimum singular value of the sampled covariance matrix σ min ( ˆ Cov( F )) to be strictly positive. As a final step, setting x = √ h ( x 1 , · · · , x N ) and combining the two previous inequalities yields a quadratic inequality:

which is well-defined and solves for

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. 1 from an algorithmic perspective to 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 a specific instance of feature functions for return distributions, namely state aggregation. More specifically, the state aggregation feature functions F = { f j } j ∈ [ M ] split the state space into M disjoint subsets, one for each function, i.e., S = ∪ j ∈ [ M ] S j and S j ∩ S j ′ = ∅ , j, j ′ ∈ [ M ] , j = j ′ , and gives back the associated return g ∈ R , namely:

✶ These features are bounded by the maximum return G max, while the empirical Rademacher complexity over N samples of returns { ( s i , g i ) } i ∈ [ N ] can be directly computed as in Clayton [2014]:

where ˆ P ( S j ) = N j /N and N j = |{ ( g i , s i ) : s i ∈ S j , i ∈ [ N ] }| . The decomposition of the Rademacher term into single terms leads to rewriting B ( ˆ λ, F , N, δ ) as in the following lemma.

Lemma1. For Distributional Max-Ent Evaluation with a state-aggregation feature class, the variance term B ( ˆ λ, F , N, δ ) is given by;

4.1 Representation Refinement: Progressive Factorization

State aggregation features are of interest due to the possibility of progressively refining the representation by increasing the factorization level, that is, by splitting a subset S j into further disjoint subsets. This refinement is called progressive factorization and is defined as follows.

Definition 1 (Progressive Factorization) . For two sets of state aggregation feature functions, F , F j , we say that F j is a progressive factorization of F , i.e., F ⊂ F j , if F = { f 1 , . . . , f j -1 , f j +1 , . . . , f M } ∪ { f j } , F j = { f 1 , . . . , f j -1 , f j +1 , . . . , f M } ∪ { f k j } k ∈ [ K ] and the additional functions { f k j } k ∈ [ K ] are such that the corresponding subsets satisfy

̸

where only non-degenerate class factorizations will be considered, meaning that the new subsets S k j are non-empty.

It is relevant for our interests that, in the case of progressive factorizations F ⊂ F ′ , the respective Max-Ent solutions enjoy the following monotonicity property

Lemma 2 (Monotonicity) . The multipliers of the Max-Ent solutions ˆ λ, ˆ λ ′ using F ⊂ F ′ are such that

This result is fully derived in Appendix C, and it ensures a monotonically increasing of all 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) is 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 generalization error 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 feature sets F 0 ⊂ F 1 ⊂ . . . with progressive factorization while performing PE. The behavior of the algorithm is similar to what is done in Structural Risk-Minimization (SRM) [Vapnik, 1991], and it involves optimizing for a trade-off: the bias term (i.e., empirical risk) decreases by taking into account more complex features classes, while the variance term (i.e., the confidence interval) increases. The whole algorithm is 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-Ent PE:

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

Algorithm 2 Distributional Max-Ent Progressive Factorization

β > 0 is to regulate the tendency to factorize. Higher values of β 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 a lower tendency to factorization.

Finally, the Progressive Factor function takes as input the list of feature functions and a factor K and returns a list of progressively factored set of feature functions. More specifically, each element in {F j } j ∈ [ M ] corresponds to a progressive factorization of the feature f j , factoring the related subset S j into K disjoint subsets as in Definition 1. The new K subsets {S j k } k ∈ [ K ] are constructed in the worst-case scenario: the complexity term in Eq. (13) is maximized with partitions of a set leading to a uniform distribution of samples in each new partitioned subset, and since it 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 to 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 2.

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 at 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 is expected to be particularly evident in MDPs admitting a factored representation of the return distribution, namely the ones in which many states are nearly equivalent under the evaluation of a policy. This factorizability property is not uncommon in general since it is present in any environment with symmetries and Block-MDPs

β=3

B

100

β= 152

tot

100

150

B

50

β=301

150

100

B

50

β= 450

100

150

B

50

0

2

3

4

Factorization Steps

β=3

tot

40

B

20

β=152

tot

40

B

20

β=301

tot

40

20

B

β=450

tot

40

120

B

2

FactorizationSteps

β=3

2.00

1.75

β= 152

1.8

1.9

1.7

β=301

1.8

1.6

β=450

1.75

2.00

0

1

2

3

4

Factorization Steps

β=3

0.8

0.6

β=152

0.550

0.525

β=301

0.55

0.50

β=450

0.550

0.525

0

2

FactorizationSteps

[Du et al., 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 β to be effective enough, while a higher level of boosting would force over-factorizations that are unnecessary, leading to no further improvement or even degradation of the results. The simulations are run on a rectangular GridWorld, with a 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 in 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 interpretability 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 1) 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; a full factorization that employs |S| -singletons of states as representations, i.e., the most fine-grained representation possible. The comparison is made via two relevant quantities, the KL divergence with respect to the true distribution (the bounded quantity ), and the total bound Btot = ˜ L (ˆ η ) + B ( ˆ λ, F , N, δ ) (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 γ is set to 0 . 98 and the confidence δ to 0 . 1 , the results are averaged over 10 rounds with the respective standard deviation.

Results Discussion The results of the simulations are reported from Fig. 1 to Fig. 4, with the quantity related to the oracle parametrization being in orange , while the ones related to the full parametrization being in blue . It is possible to notice that these two distributions have almost the same KL divergences with respect to the true return distribution (Fig. 2, 4), yet they highly differ

in the bound Btot (Fig. 1, 3) mostly due to the variance term, which is way higher in the case of full factorization. This suggests that the bound is indeed able to 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 β , namely at F i ⋆ . The plots should be read as follows: while the bound term Btot is expected to increase at each factorization step, the KL divergences with respect to the true return distribution should decrease as much as possible. In all cases, it is evident that the value of β pushes towards a higher number of factorization steps, going from performing no factorization at all using low values ( β = 3) , to performing up to 4 factorization steps even in this simple scenario with higher values ( β = 450) , both at low and high sample regimes ( N ∈ { 50 , 1000 } ) . Furthermore, at higher sample regimes, it is possible to see how the higher quality of the estimation counteracts the action of β , and increasing it generally induces still fewer factorizations compared to the low sample regimes with same values of β , as in Fig. 3, 4. Finally, it is apparent that minimizing for Eq. (16) successfully decreases the KL 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 as 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 to 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 effectiveness [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-Carlo 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 et 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 a long and extensive literature in the density estimation field, which mostly focused on the general and algorithmic aspects of the method [Barron and Sheu, 1991, Dudík et al., 2004, Sutter et al., 2017]. Among the others, Cortes et al. [2015] 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 a 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 a sup λ ∈ Ω || λ || 1 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 state-dissimilarity.

6.2 Conclusions and Future Works

In our work, we presented in a 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 a particular class of features, namely state aggregation, and we proposed an algorithm

called Distributional Max-Ent Progressive Factorization to adaptively find a 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 a reduced-dimension 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 an 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 of 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: 10.1214/aos/1176348252.

Marc G. Bellemare, Salvatore Candido, Pablo Samuel Castro, Jun Gong, Marlos C. Machado, Subhodeep Moitra, Sameera S. Ponda, and Ziyu Wang. Autonomous navigation of stratospheric balloons using reinforcement learning. Nature , 588(7836):77-82, December 2020. ISSN 00280836, 1476-4687. doi: 10.1038/s41586-020-2939-8.

Marc G. Bellemare, Will Dabney, and Mark Rowland. Distributional Reinforcement Learning . MIT Press, 2023.

Daniel Brown, Scott Niekum, and Marek Petrik. Bayesian robust optimization for imitation learning. Advances in Neural Information Processing Systems , 33:2479-2491, 2020.

Yinlam Chow, M. Ghavamzadeh, Lucas Janson, and M. Pavone. Risk-constrained 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.06366 [cs, math, stat].

Simon S Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudík, 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 978-3-540-35296-9. doi: 10.1007/11776420_ 12. Series Title: Lecture Notes in Computer Science.

Miroslav Dudík, Steven J. Phillips, and Robert E. Schapire. Performance Guarantees for Regularized MaximumEntropy 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-27819-1. doi: 10.1007/978-3-540-27819-1_33. Series Title: Lecture Notes in Computer Science.

M. Dudley. Uniform Central Limit Theorems . Cambridge University Press, 1 edition, July 1999. ISBN 978-0-521-46102-3 978-0-521-05221-4 978-0-511-66562-2. doi: 10.1017/ CBO9780511665622.

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-4443, April 2020. doi: 10.1609/aaai.v34i04.5870.

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.1214/aos/ 1015362183.

Christos N. Mavridis, Nilesh Suriyarachchi, and John S. Baras. Maximum-Entropy Progressive State Aggregation for Reinforcement Learning. In 2021 60th IEEE Conference on Decision and Control (CDC) , pages 5144-5149, Austin, TX, USA, December 2021. IEEE. ISBN 978-1-66543-659-5. doi: 10.1109/CDC45484.2021.9682927.

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. DualDICE: Behavior-Agnostic Estimation of Discounted Stationary Distribution Corrections, November 2019. URL http://arxiv.org/ abs/1906.04733 . arXiv:1906.04733 [cs, 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, Marc G. Bellemare, and Will Dabney. Statistics and Samples in Distributional Reinforcement Learning. 2019. doi: 10.48550/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 MultiAgent Reinforcement Learning. page 9, 2021.

Mark Rowland, Rémi Munos, Mohammad Gheshlaghi Azar, Yunhao Tang, Georg Ostrovski, Anna Harutyunyan, Karl Tuyls, Marc G. Bellemare, and Will Dabney. An Analysis of Quantile TemporalDifference Learning, January 2023. arXiv:2301.04462 [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.48550/ARXIV.1708.07311.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

Csaba Szepesvári. Algorithms for Reinforcement Learning. Synthesis Lectures on Artificial Intelligence 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-6 978-14757-2545-2. doi: 10.1007/978-1-4757-2545-2.

Benjamin Van Roy. Performance Loss Bounds for Approximate Value Iteration with State Aggregation. Mathematics of Operations Research , 31(2):234-244, May 2006. ISSN 0364-765X, 1526-5471. doi: 10.1287/moor.1060.0188.

Vapnik. Principles of Risk Minimization for Learning Theory. In Advances in Neural Information Processing Systems , volume 4. Morgan-Kaufmann, 1991.

Martin J. Wainwright and Michael I. Jordan. Graphical Models, Exponential Families, and Variational Inference. Foundations and Trends in Machine Learning , 1:1-305, 2007. ISSN 1935-8237, 19358245. doi: 10.1561/2200000001.

Shaojun Wang, Russell Greiner, and Shaomin Wang. Consistency and Generalization Bounds for Maximum Entropy Density Estimation. Entropy , 15(12):5439-5463, December 2013. ISSN 1099-4300. doi: 10.3390/e15125439.

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n* 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\nWEAKNESSES:\n* 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\nQUESTIONS:\nTo 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\nLIMITATIONS:\nNo", "TITLE": "Official Review", "RECOMMENDATION": 2, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 1, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\nI 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.\n\nWEAKNESSES:\nThe 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\nQUESTIONS:\nI 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

comments: 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.

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.

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? 
(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.)
Maybe 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\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.\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.\n\nQUESTIONS:\n- 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).\n\nLIMITATIONS:\nThe authors discuss limitations, such as their method requiring access to IID samples.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-08"}

comments: 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.

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.

QUESTIONS:
None

LIMITATIONS:
None
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-27