SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\mathcal{F}_0 \subset \mathcal{F}_1 \subset \ldots$  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:
$$
\mathcal {J} (\hat {\eta}) = \beta \mathcal {L} (\hat {\eta}) + B (\hat {\lambda}, \mathcal {F}, N, \delta), \tag {16}
$$
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
$\beta > 0$  is to regulate the tendency to factorize. Higher values of  $\beta$  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  $\{\mathcal{F}_j\}_{j \in [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_k^j\}_{k \in [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  $\beta$  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
[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  $\beta$  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  $\mathcal{A} = (\mathrm{up},\mathrm{left},\mathrm{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  $\mathrm{B_{tot}} = \tilde{\mathcal{L}} (\hat{\eta}) + B(\hat{\lambda},\mathcal{F},N,\delta)$  (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  $\gamma$  is set to 0.98 and the confidence  $\delta$  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  $\mathrm{B}_{\mathrm{tot}}$  (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  $\beta$ , namely at  $\mathcal{F}_{i^{\star}}$ . The plots should be read as follows: while the bound term  $\mathrm{B}_{\mathrm{tot}}$  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  $\beta$  pushes towards a higher number of factorization steps, going from performing no factorization at all using low values  $(\beta = 3)$ , to performing up to 4 factorization steps even in this simple scenario with higher values  $(\beta = 450)$ , both at low and high sample regimes  $(N \in \{50, 1000\})$ . Furthermore, at higher sample regimes, it is possible to see how the higher quality of the estimation counteracts the action of  $\beta$ , and increasing it generally induces still fewer factorizations compared to the low sample regimes with same values of  $\beta$ , 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.
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_{\lambda \in \Omega}||\lambda ||_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.
This paper is supported by PNRR-PE-AI FAIR project funded by the NextGeneration EU program.
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 0028-0836, 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 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-27819-1. doi: 10.1007/978-3-540-27819-1_33. Series Title: Lecture Notes in Computer Science.
R. 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.