Near-Minimax-Optimal Distributional Reinforcement Learning with a Generative Model

Summary

This paper proposes a novel model-based algorithm for policy evaluation in distributional RL and shows its near optimality within the generative model framework. Additionally, they conduct an experimental study comparing their algorithm with quantile DP in a synthetic environment.

Strengths

Discretizing the distribution's support and then rewriting the update as a linear mapping is an important algorithmic contribution of this paper and is particularly interesting. It reveals a connection between categorical algorithms and linear systems---and, as they did, we can reduce the former to the latter. The paper is well-written. The algorithm and its main idea are clearly presented. The theory looks sound.

Weaknesses

Why do you assume N iid samples for *each state* (line 56)? The common wisdom in policy evaluation under generative model is that more samples are needed for states frequently visited by the policy and fewer samples are needed for states that are less visited. As an extreme case, consider a state that is entirely unreachable---in this case no samples would be needed for that state, and the theoretical analysis should still hold. Thus, I wonder if the results in this paper can be achieved under that assumption. According to the experimental results, the proposed algorithm does not outperform quantile DP in terms of sample efficiency. However, it has significantly better running time.

Questions

Is the cubic dependence on the state space for computational efficiency good?

Rating

6

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors discussed the limitations, and I found no potential negative societal impact.

Summary

The paper presents a novel algorithm for model-based distributional RL and establishes that it achieves near-minimax-optimal performance for approximating return distributions in the generative model regime. This is a significant contribution, showing that the distributional RL is sample-efficient with a generative model, and being the first work providing practical implementation guidance to achieve this.

Strengths

The paper introduces the new algorithm DCFP for distributional RL, which directly approximates the cumulative probability. The development of the DCFP algorithm is a key contribution. The DCFP algorithm effectively leverages a generative model to achieve efficient performance. The authors show that DCFP achieves minimax lower bounds for sample complexity in Wasserstein distance, thus addressing a question raised by Zhang et al. (2023). Besides, the paper's theoretical contributions, such as the introduction of the stochastic categorical CDF Bellman equation and the detailed derivations of the minimax-optimal sample complexity, are significant but may be challenging for readers without a strong background. I believe this work provides a deeper understanding of the field of distributional RL.

Weaknesses

I have some concerns about the experiments. The authors claim that DCFP is generally faster than QDP. I feel this conclusion is primarily due to computationally costly sort operations in QDP. I am also curious about whether both DCFP and QDP are generally lower than CDP, as DCFP also involves costly matrix inverse operations. It would be better to show some results of CDP in the main paper. Additionally, the authors mentioned $T_P$ has some sparse structure allowing for potential speed-ups and making use of sparse linear solvers from the scipy package in implementation. If solving this linear system using these solvers is stable and whether the results remain reliable as the state space scales up. It would be helpful to discuss these.

Questions

Q1: The algorithm operates under the generative model assumption, where it can sample from any state-action pair. Since I am not particularly familiar with this field, could you elaborate on the justifiability of this assumption? In my limited knowledge, this assumption seems quite far away from reality, and discussing its limitations and applicability would be beneficial. Q2 The proposed DCFP algorithm aims to align theoretical analysis with practical implementation and can be seen as a modified category-based distributional RL algorithm. However, However, is the sacrifice in computational efficiency deserved? It would be better to discuss the scalability of the proposed DCFP algorithm in real-world, large-scale environments. Q3: The current algorithm is model-based and solving this algorithm relies on constructing the empirical transition probabilities. Does this suggest that DCFP is not easily applicable in the model-free setting? In other words, does the idea of directly learning the cumulative probability mass rather than discrete probability mass face additional challenges in a model-free setting? Furthermore, if sample-efficient distributional RL can only be achieved under a generative model assumption?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Please refer to the question and weakness part.

Summary

The authors propose a new algorithm for distributional reinforcement learning under the generative model setting with categorical representation. New upper bound for sample complexity is presented for the proposed algorithm. Some empirical results comparing the method with other alternatives are presented, showcasing its benefits over the alternatives.

Strengths

The main algorithm as well as the sample complexity results are solid contribution to the reinforcement learning community. Although I did not check the proofs, the results seem believable.

Weaknesses

The generative model setting is rather limited. The results do not directly contribute towards learning a better policy, but I assume some of the tools presented will find their uses in many scenarios.

Questions

How about the all-gaussian setting where every reward distribution is a gaussian?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Summary

This paper proposes a min-max optimal algorithm for model-based distributional RL, which is used to approximate the return distributions under the assumption of having access to generative model. New theoretical analysis is provided for categorical approaches in distributional RL, with the introduction of a new distributional Bellman equation and the stochastic categorical CDF Bellman equation. Empirical studies are performed for benchmark comparison with other model-based distributional RL algorithms.

Strengths

The main contribution of the paper is the theoretical results for distributional RL by assuming the access of a generative model, where the goal is to estimate the full distribution of returns at each state instead of just the expected returns. In particular, this work exhibits several promising results: 1. Authors establish theoretical guarantees for the proposed distributional RL algorithm, i.e. the direct categorical fixed-point algorithm (DCFP), which directly computes the fixed point of CDP (an existing distributional RL algorithm). It achieves a sample complexity that matches with the lower bound up to logarithmic factors when estimating return distribution in Wasserstein distance. 2. Theoretically, authors introduces a new distributional Bellman equation and the stochastic categorical CDF Bellman equation for categorical distributional RL.

Weaknesses

While this paper provides insights into the distributional RL, there are still potential improvements that authors can further consider: 1. It is unclear about the motivation of considering categorical approaches to distributional RL and how the proposed algorithm excels existing distributional RL approaches throughout the text. While it makes sense to choose a tractable representation of probability distributions for approximation, the representation considered as in line 81 - 88 appears to be simple and may have limited expressiveness in practice. 2. The convergence result provided in Proposition 2 does not categorize the convergence rate, which can be essential for computational efficiency. 3. While the theoretical results are interesting, the practical applicability is not quite clear. Numerical results are presented on a 5-state simple tabular environment. It will be beneficial to include a practical example to better illustrate how the proposed algorithm can be particularly useful.

Questions

1. In the studied setting, how does "policy" and "actions" come into play when evaluating return / value functions? Do you consider fixed policy? 2. In Section 2, it seems that the randomness being considered comes from transition, reward function and initial state distribution. How about the randomness in policy?

Rating

5

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

This is a theoretical work, no potential negative social impact