Constrained Latent Action Policies for Model-Based Offline Reinforcement Learning

Summary

This paper proposes C-LAP, a model-based offline RL method that mitigates the distributional shift problem without using uncertainty penalties or modifying the Bellman update. It learns a joint distribution of latent states and latent actions and constrains the latent action policy to the dataset distribution by using a linear transformation. Experiments on D4RL datasets and V-D4RL datasets report performance comparable to other strong offline RL baselines.

Strengths

1. This paper addresses the issue of distributional shift in offline RL by incorporating latent actions to enhance policy learning. The empirical results demonstrate that the proposed method outperforms other offline RL baselines by requiring fewer gradient steps to learn a policy. 2. The idea of combining a latent action prior and a bounded policy to obtain the constrained policy is novel and intuitive.

Weaknesses

1. Figure 1 requires a clearer presentation of the model diagram. It would be better to simultaneously show the prior and posterior in Figure 1(a), and the latent action prior $p_\theta(u_t|s_t)$ and the policy $\pi_\psi(u_t|s_t)$ in Figure 1(b) to improve understanding. During the policy training phase, are the imagined trajectories rolled out from the latent state prior $p_\theta(s_t|s_{t−1}, u_{t−1})$ or the latent state posterior $q_\phi(s_t|s_{t−1}, a_{t−1}, o_t)$? If it is the latent state prior, why is the output of the action decoder used? If it is the latent state posterior, $o_t$ is unavailable. 2. Could the authors provide some distribution results to support the claim that the latent actions generated by the action prior are close to the dataset's action distribution? Does the predictive power of using latent actions degrade in performance compared to models using real actions? Some quantitative or qualitative results can be provided. 3. It would be better to highlight the differences between the proposed method and Dreamer in terms of model designs and algorithms. 4. The experiments mainly focus on Mujoco tasks. How does the proposed method perform on navigation tasks such as AntMaze? 5. Why do some experiments in Figure 4, 5 show high returns at the beginning of policy learning (zero gradient step)? 6. In table 5 and table 6, the performance of C-LAP is inferior to some baseline models in medium or medium-relay datasets. This results seem to indicate that the effectiveness of the proposed method depends on the quality of the offline dataset.

Questions

Please see the weaknesses section.

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

3

Limitations

This paper discusses the limitations, but that is not enough.

Summary

The paper approaches the problem of model-based policy learning from static datasets. It firstly identifies that this offline MBRL setup inherits two major concerns from its components: the problem of value overestimation from out-of-distribution actions (common in offline RL) and the bias originated by learning policies from model predictions. In order to address these issues, it proposes C-LAP, a MBRL method that besides learning a latent state space model - as common in the literature - also learns a latent action space from the static dataset. This generative model of states AND actions is further leveraged by the policy optimization objective to constrain the policy actions to share support with the latent action prior, in order to prevent out-of-distribution samples. The paper provides experiments in two offline RL benchmarks, with feature and pixel-based observations, and showing good improvements in the pixel-based observation setting.

Strengths

- The paper addresses the open problem of offline model-based reinforcement learning, which is very relevant to the RL community nowadays. The paper is well-written, very clear and organized: the methodology is carefully and formally described; the baselines are clearly and concisely described (and justified). - Despite the incremental nature of the paper (while mixing up different building blocks from previous literature), the work is really clever on identifying that a latent modeling of actions would implicitly constrain the action space to the data distribution, which would naturally provide a regularization effect on learning policies with static datasets. Therefore, the general methodology is very well motivated and applied to the problem. - The asymptotic results in the V-D4RL are really strong and the ablations in Section 4.2 presents very clearly the effect of the proposed method in the problem of value overestimation. These results give solid evidence that the method is indeed addressing the previously raised problems.

Weaknesses

- The claim in L200 regarding a “significant speed up in policy learning” is questionable, since comparing the number of gradient steps to achieve a certain performance is not fair. Some methods can leverage less computation in the updates than others. Also, one can carefully tune the hyperparameters (such as the learning rate, batch size) to minimize the number of gradient steps required, at the cost of more computation. Therefore, showing the gradient steps in the x-axis is not ideal and does not seem to provide useful information. It would be interesting to bring the computational cost for claiming speed up. Optionally, the work could remove such a claim. - There is an unclear trend in the results of the proposed method: in many cases the starting policy already performs close to the best, while other methods start from performance 0. For instance, refer to the expert cases in D4RL. Why does this happen? - The most interesting case is the walker 2d-expert-v2 in Figure 6, where one ablation starts optimal and then goes down. It would be interesting to justify or at least hypothesize what is happening. - The bound of Eq. 11 is not justified in the paper and looks arbitrary. From my understanding no prior is placed in the policy so this looks like an assumption. It would be interesting to elaborate better about this. - The last concern is perhaps the most general and crucial to be cleared out during rebuttal. It is clear that the method worked much better for visual domains than feature domains (the results in the D4RL dataset do not show any asymptotic improvement – it is actually worse than MOBILE as per Table 5). In both cases, the action space is the same. It is also low-dimensional. This raises the question: is the model really implementing latent variables for actions or is it actually implementing a hierarchical latent model for the state space? The work explicitly defines “u_t” as latent actions, but it could also be interpreted as a higher-level latent state variable. Given that the gradients from the action decoder backpropagates through both s_t and u_t, I couldn’t find an explanation on why this is not a valid explanation of the method, and this would better justify the discrepancy in the results. **Minor Concerns/Further Suggestions** - It would be great to show the value overestimation plots (Figure 6) for the considered baselines, in order to compare the difference among them. It is unclear if they perform worse than the proposed method because of value overestimation or other reason. - It would also be nice to show a sensitivity analysis of the policy constraint hyperparameter epsilon across a few environments, to help understanding how the proposed method behaves and how much it is environment specific. - In L105, the paper contrasts the proposed C-LAP to previous works stating that they “rely on uncertainty estimates to generate trajectories within the data distribution”, but I would consider rephrasing that since C-LAP can also be understood as leveraging the uncertainty of the latent action prior to generate trajectories. Perhaps the distinction is that previous methods explicitly use uncertainty estimates, while C-LAP relies on a more implicit constraint. **Typo:** - A missing “)” in Eq. 12.

Questions

Please refer to the weaknesses section. Each Major Concern contains questions to be addressed during the rebuttal.

Rating

7

Confidence

4

Soundness

2

Presentation

4

Contribution

3

Limitations

- The main limitation in my perspective is that it is not clear why the method presents much stronger results in the visual benchmark (which is supposed to be harder) than in the feature-based benchmark. This opens questions about whether the introduced latent modeling is indeed doing what is hypothesized in the paper. - There are some limitations in how the results are presented and discussed in the work (please refer to first two bullet points in the Weaknesses section). **Summary of the Review**: Overall, the paper is well-written, clearly describes the proposed methodology and offers strong results in the visual benchmarks. The work is based upon a very interesting insight on what latent variables could potentially offer to restrict out-of-distribution actions in offline RL settings. The paper does present some concerns, most related to better explaining the presented results and adjusting claims accordingly. There are also other minor concerns/further suggestions that do not prevent acceptance but would also enrich the paper (which I would appreciate with a higher score). That being said, I believe the paper is already above the acceptance threshold and I am willing to increase my score if my concerns (mostly questions) are properly discussed in the rebuttal phase. # Post-Rebuttal Thank you, authors, for properly addressing my questions and concerns during the rebuttal. After rebuttal, I am increasing my score.

Summary

The work introduce C-LAP (Constrained Latent Action Policies), a novel approach to model-based offline reinforcement learning in POMDPs. Notably, C-LAP does not employ explicit reward penalty terms regarding the action space. The paper proposes a methodology for learning latent variables for both states and actions based on the state-space models (SSMs). It then frames policy optimization as a constrained optimization problem, aiming to maximize N-step returns while ensuring that latent actions remain within the support of the latent action prior. This approach allows the policy to maximize returns predicted by the model without incorporating reward penalty terms, such as uncertainty penalization. Empirically, C-LAP demonstrates superior performance compared to baseline methods that incorporate uncertainty penalties in their reward structure.

Strengths

The manuscript presents a novel and effective model-based method for offline reinforcement learning in POMDPs. Prior approaches in the domain of model-based offline RL often incorporate penalty terms to mitigate OOD actions and prevent overestimation of the value function. This penalization, however, can restrict the maximization of value functions related to possibly promising actions. In this work, the authors extend a standard state-space model (SSM), exemplified by model-based approaches such as Dreamer, to include latent action variables. By constraining the latent action space to adhere to an action prior, the generation of OOD actions is effectively precluded. This modification is both simple and effective, addressing a critical challenge in the field. The effectiveness of the proposed method is validated across both MDP and POMDP settings using standard benchmarks such as D4RL and V-D4RL. Moreover, the authors provide a thorough ablation study to identify the critical components contributing to performance improvements. As described in Figure 6, omitting the constraint on the latent action prior is the most crucial to performance, underscoring its significance as a key contribution of this research.

Weaknesses

Despite the novelty of this work, the motivation for adopting latent action spaces is insufficiently articulated. The paper does not clearly address why constraining with respect to the latent action space is advantageous over the actual action space. Specifically, it remains unclear whether C-LAP offers any benefits compared to SSMs that include a support constraint term directly on the actual action space, as expressed by the constraint $\mathbb{E}_{s \sim p, a\sim \mu(s)}[\pi (a|s) ] \ge \epsilon$ (where $\mu$ denotes the behavior policy) for an instance. I believe that a more thorough explanation for the motivation could significantly strengthen the contributions and arguments of this work. (Minor comment) In Eq. (7), a constraint term would be a typo. It would be more accurate to revise this as $\mathbb{E}_{p,\pi}[p(h_t|s_t)]\ge \epsilon$, given that $h_t$ is sampled from $\pi(s_t)$.

Questions

Q1. The manuscript suggests potential applicability of the C-LAP framework beyond the traditional models explored. Have the authors considered leveraging foundation models, such as transformers or diffusion models, within the latent space formulation of C-LAP? If so, could you elaborate on how C-LAP might integrate with these advanced model architectures? Q2. The proposed C-LAP method appears to be promising for environments characterized by high-dimensional action spaces, including those analogous to language action spaces. Could the authors provide insights into the adaptability of C-LAP to such high-dimensional scenarios? What modifications, if any, would be necessary to optimize its performance in these contexts?

Rating

6

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors included their limitation in Section 4.2. They address societal impacts of their work in the checklist #10. Broader Impacts.

Please respond to author response if you haven't

Dear Reviewers, Before the author-reviewer discussion period ends, please make sure that you have read the author responses, acknowledge that your review reflects them, and engage with the authors if necessary. Thank you!