Paper: Model-based Reinforcement Learning for Parameterized Action Spaces (ICLR-2024)

Reviewer comment (Writing):
Quite a few grammar errors:  
- Sec 1, first paragraph: “… at each time step the agent must choose a discrete action type (move, ribble or shoot) and also continuous parameters (related to) that chosen discrete action, … each discrete action is parameterized by some continuous parameters, …”.  
- Sec 1, second paragraph: “By contrast, in continuous/discrete-only action spaces, …”.  
- Sec 1, third paragraph: ”It then performs Model Predictive Control …”.  
- Sec 4.2, third paragraph: “The discrete actions and continuous parameters …”.  
- Sec 4.2, fourth paragraph: “This is because each time, the agent needs to first sample from the discrete actions then sample from the corresponding continuous parameters. And each discrete action has a independent continuous parameter space. This becomes …”.  
- Sec 4.2, fifth paragraph: “… the agent executes E steps of forward planning while updating the distribution parameters over discrete actions and continuous parameters. Then it uses the first action …”.  
- Sec 5, first paragraph: “Goal: The agent needs …”.  
- Sec 5, first paragraph: “The number of parameterized actions is …”.  
- Sec 5, second paragraph: “HyAR learns an embedding of the parameterized action space …”.  
- Sec 5.1, first paragraph: “We find that DLPA achieves significantly higher sample efficiency …, … which is consistent with the results …, … DLPA on average achieves 30× higher sample efficiency compared to the best model-free method …”.

Author response:
Thank you for the insightful feedback on our submission. We believe we can address your concerns. We have fixed the grammar errors accordingly in our updated manuscript.

Reviewer comment (Novelty):
The author's description of the key innovations in this paper, such as the calculation of the H-step loss in model learning or the introduction of the CEM method, does not place enough emphasis on what makes these approaches unique when applied to the specific problem scenario PAMDP. For instance, what sets apart the H-step loss in DLPA compared to the traditional model-based approaches, or how does this paper's use of CEM bring innovation in the context of PAMDP, especially regarding handling the relationship between discrete and continuous actions? These are the questions that are more central to the paper and that I would like to have a better understanding of.

Author response:
We believe the primary contribution of this paper is the very first model-based RL algorithm for the PAMDP setting. It outperforms all the previous model-free RL methods - many of which have been published at ICLR, including one just last year - for PAMDP on 8 standard PAMDP benchmarks. The PAMDP setting is very common in real-world tasks, and has drawn a lot of attention in recent papers, but all of them are model-free methods.  We hope our paper can draw other researchers' attention to the model-based aspect of this problem setting since DLPA performs better than even the most sophisticated model-free methods. Calculation of the H-step loss in model learning is an important technique that we found to be especially useful in practice. Additionally, for the CEM method in PAMDPs, previous works have only shown its efficiency on either discrete action space or continuous action space, while we’re the first to apply such a method to the parameterized action space. Note that **one major modification we make when applying MPPI to PAMDPs is that we keep a separate distribution over the continuous parameters for each discrete action and update them at each iteration instead of keeping just one independent distribution for the discrete actions and one independent distribution for the continuous parameters.** In other words, we let the distribution of the continuous parameters condition on the chosen discrete action during the sampling process. We believe this is an important change to make when using MPPI for PAMDPs as we don’t want to throw away the established dependency between the continuous parameter and corresponding discrete action in PAMDPs. In the updated manuscript Section (6.3) we have included a new ablation study result for this change which shows that this change is indispensable for DLPA to do proper planning in PAMDP.  

In terms of novelty, **we have also added an entirely new theoretical analysis section in the updated manuscript where we derive some performance guarantees / bounds for the proposed algorithm (Section 5).** In general, we provide three main theorems that describe the performance guarantees of DLPA from three aspects: 1. A theoretical bound for the cumulative return difference between the rollout trajectories of DLPA and the optimal trajectories. 2. A theoretical bound for the multi-step prediction error of DLPA. 3. A theoretical bound for how the estimation error will change with respect to the number of samples and the cardinality of the continuous parameter space.

Reviewer comment (Experiments):
During the experimental process, the number of training steps for DLPA differs from the other baseline methods. This discrepancy makes it challenging to adequately assess the stability of DLPA's training, including whether it exhibits oscillations or other issues. Additionally, it appears that HyAR may not achieve full convergence in certain environments. Therefore, it would be advisable to provide additional experimental results with a consistent number of training steps for all algorithms to ensure a fair comparison.

Author response:
Thanks for your suggestion, before we only show the result when DLPA has already converged in Figure 3. We aim to use Figure 3 to show the superior sample efficiency achieved by DLPA. Table 1 shows the asymptotic performance where all the other algorithms have fully converged. Even after giving all the other baselines a significantly larger number of samples, DLPA still achieves better performance in 6 out of the 8 tested domains and comparable performance in the other 2. **We have also updated the plots for the comparison of all methods with full steps in Figure 6 in the appendix**, where the results for the baselines are very close to what Li et al. 2022 reported in their paper. Note that all the experiments were run on a Nvidia GeForce RTX 3090 GPU. We directly use the source code provided by the authors for those baselines and so the resources are kept the same for all benchmarks. We also run DLPA for an additional amount of steps after convergence for each environment, and the performance is still quite stable as shown in Figure 6.

Reviewer comment (Experiments):
For the experiments assessing the impact of CEM-based MPC on DLPA, it would be beneficial to include ablation studies with other traditional model-based non-MPC methods, such as MBPO, applied in the context of PAMDP.

Author response:
Thanks for your suggestion. We add the comparison to the Dyna-like model-based RL method in Section 6.3, which is a baseline very close to MBPO that does not use MPC. We use a model-free RL framework for PAMDP—PDQN—as the policy learning framework given the data generated from the learned model. We have also tuned the proportion between the generated data and the real-world rollouts to make sure the baseline reaches its best performance. As we can see from the training results on two benchmarks,  this non-MPC baseline generally performs better than the model-free baselines but is not as good as DLPA.

Reviewer comment (Reproducibility):
What specific form do the parameters of CEM take in this context? Are they similar to the parameters of a policy neural network, and how do they relate to the state input?

Author response:
The parameters of CEM in DLPA directly control the distribution over the discrete actions and continuous parameters and are different from the parameters for a policy neural network. Specifically, for discrete actions, the parameters of CEM are the parameters for a multinomial distribution. We sample from this distribution to generate discrete action during planning. For continuous parameters, the parameters of CEM are the parameters for a Gaussian distribution. We directly sample from it to generate the corresponding continuous parameters during planning. After randomly initializing the action distributions, we use the learned model and the **initial state input** to rollout imagination trajectories and estimate their cumulative return. And we pick the actions to execute based on the estimated returns. Note that for each planning step, we will update the CEM parameters for several iterations (described by Eqn 5, 6, 7) given the estimated returns.

Reviewer comment (Theory):
How does the computational complexity of CEM change as the action space expands?

Author response:
We have included a new Table in appendix 8.6 that compares the clock time of planning and the number of timesteps needed to converge as the action space expands. We tested this on the Hard Move domain, where the number of discrete actions changes from $2^4$ to $2^{10}$ (one continuous parameter for each of them). As shown in the table, while the number of samples increases as the action space expands, it’s still within an acceptable range even when it’s extremely large. We also quantify this with theoretical analysis which is included as a new section in the updated manuscript.

Reviewer comment (Novelty):
Novelty and contribution is limited.

Author response:
Thank you for the review. The research question of this paper centers on Reinforcement learning problems where actions are of a mixed type and interactions are readily described by Parameterized Action MDPs (PAMDPs). Here we would like to emphasize again that the primary contribution of this paper is the first model-based RL algorithm for the PAMDP setting, which outperforms all the previous model-free RL methods for PAMDP on 8 standard PAMDP benchmarks. PAMDP is a very common problem setting in real-world tasks, and has gained a lot of attention in recent papers despite all of them being model-free methods.  We hope our paper can draw other researchers' attention to the model-based aspect of this problem setting since DLPA already performs better than the model-free methods without very fancy designs. 

In terms of the specific design when applying model-based RL on PAMDP tasks, besides the multi-step prediction technique which is also mentioned by the other reviews, **one major modification we make when applying MPPI to PAMDPs is that we keep a separate distribution over the continuous parameters for each discrete action and update them at each iteration instead of keeping just one independent distribution for the discrete actions and one independent distribution for the continuous parameters.** In other words, we let the distribution of the continuous parameters condition on the chosen discrete action during the sampling process. This is an important change to make when using MPPI for PAMDPs as we don’t want to throw away the established dependency between the continuous parameter and corresponding the discrete action in PAMDPs.  In the updated manuscript we have included a new ablation study result for this change which shows that this change is indispensable for DLPA to do proper planning in PAMDP (Section 6.3).

In terms of novelty, we have also added an entirely new theoretical analysis section in the updated manuscript where we derive some performance guarantees / bounds for the proposed algorithm (Section 5). In general, we provide three main theorems that describe the performance guarantees of DLPA from three aspects: 1. A theoretical bound for the cumulative return difference between the rollout trajectories of DLPA and the optimal trajectories. 2. A theoretical bound for the multi-step prediction error of DLPA. 3. A theoretical bound for how the estimation error will change with respect to the number of samples and the cardinality of the continuous parameter space.

Reviewer comment (Writing):
What is the research question you study? In other words, what is the specific difficulty or particular problem when applying existing model-based RL methods on PAMDP tasks?

Author response:
Thank you for the review. The research question of this paper centers on Reinforcement learning problems where actions are of a mixed type and interactions are readily described by Parameterized Action MDPs (PAMDPs). Here we would like to emphasize again that the primary contribution of this paper is the first model-based RL algorithm for the PAMDP setting, which outperforms all the previous model-free RL methods for PAMDP on 8 standard PAMDP benchmarks. PAMDP is a very common problem setting in real-world tasks, and has gained a lot of attention in recent papers despite all of them being model-free methods.  We hope our paper can draw other researchers' attention to the model-based aspect of this problem setting since DLPA already performs better than the model-free methods without very fancy designs. 

In terms of the specific design when applying model-based RL on PAMDP tasks, besides the multi-step prediction technique which is also mentioned by the other reviews, **one major modification we make when applying MPPI to PAMDPs is that we keep a separate distribution over the continuous parameters for each discrete action and update them at each iteration instead of keeping just one independent distribution for the discrete actions and one independent distribution for the continuous parameters.** In other words, we let the distribution of the continuous parameters condition on the chosen discrete action during the sampling process. This is an important change to make when using MPPI for PAMDPs as we don’t want to throw away the established dependency between the continuous parameter and corresponding the discrete action in PAMDPs.  In the updated manuscript we have included a new ablation study result for this change which shows that this change is indispensable for DLPA to do proper planning in PAMDP (Section 6.3).

In terms of novelty, we have also added an entirely new theoretical analysis section in the updated manuscript where we derive some performance guarantees / bounds for the proposed algorithm (Section 5). In general, we provide three main theorems that describe the performance guarantees of DLPA from three aspects: 1. A theoretical bound for the cumulative return difference between the rollout trajectories of DLPA and the optimal trajectories. 2. A theoretical bound for the multi-step prediction error of DLPA. 3. A theoretical bound for how the estimation error will change with respect to the number of samples and the cardinality of the continuous parameter space.

Reviewer comment (Writing):
Section 2.2: "model to take in ~~this~~ actions".

Author response:
Thanks for pointing all these out! We have updated the manuscript accordingly.

Reviewer comment (Reproducibility):
Section 2.2: From what state are the action sequences sampled?

Author response:
Action sequences are sampled starting from the current timestep t (so the initial state s_t) until timestep t+H from randomly initialized distributions like Gaussians.

Reviewer comment (Writing):
Section 2.2: It would be helpful to include a reference to introductory material on Model Predictive Control.

Author response:
Thanks for pointing all these out! We have updated the manuscript accordingly.

Reviewer comment (Presentation):
Section 3: This dichotomy of model-based methods confused me, because the two categories (random shooting and data augmentation) don't entirely split the class of methods. The issue is with random shooting. The suggestion is to use a more descriptive and potentially conventional dichotomy which splits the class based on how the model is used in planning, such as Dyna-based methods and decision-time planning methods, where one uses planning for credit assignment and the other for policy evaluation.

Author response:
Thanks for pointing all these out! We have updated the manuscript accordingly.

Reviewer comment (Writing):
Section 4: Section 4.1 and Figure 1 point to notation that isn't defined until Section 4.2. Consider introducing notation earlier.

Author response:
Thanks for pointing all these out! We have updated the manuscript accordingly.

Reviewer comment (Presentation):
Figure 1: The illustrations are nicely done. However, the caption could use more supporting text so readers understand the semantics of the green dots and the trajectories inside the states.

Author response:
Thanks for pointing all these out! We have updated the manuscript accordingly.

Reviewer comment (Writing):
Section 4.1: Why do you call \(s_t\) an observation when it was previously defined as a state in Section 2 (assuming it comes from \(S\))? If you want to distinguish it from outputs of the dynamics model, then you can always use "environment state" and "model state."

Author response:
Thanks for pointing all these out! We have updated the manuscript accordingly.

Reviewer comment (Writing):
Section 4.1: Should \(r_{t+1}\) have a hat, since it is an estimated quantity?

Author response:
Thanks for pointing all these out! We have updated the manuscript accordingly.

Reviewer comment (Writing):
Section 4.1: "by learning to inferring" → "by learning to infer".

Author response:
Thanks for pointing all these out! We have updated the manuscript accordingly.

Reviewer comment (Writing):
Section 4.2: Defining the parameters before they are used in \(\mathcal{C}^0\) would add clarity to this section.

Author response:
Thanks for pointing all these out! We have updated the manuscript accordingly.

Reviewer comment (Writing):
Section 4.2: Equation 4 should explicitly state that \(\hat{s}_{t_0} = s_{t_0}\).

Author response:
Thanks for pointing all these out! We have updated the manuscript accordingly.

Reviewer comment (Writing):
Section 4.2: You are overloading notation for \(d\). Consider using \(\text{d}\) or some other notation to avoid overloading.

Author response:
Thanks for pointing all these out! We have updated the manuscript accordingly.

Reviewer comment (Presentation):
Algorithm 1: "Execute the first action ..." — It would add clarity to write out the action using your notation.

Author response:
Thanks for pointing all these out! We have updated the manuscript accordingly.

Reviewer comment (Reproducibility):
Algorithm 1: The "training" variable is never defined, so it is unclear why the conditional statement is included. There is an initial period of data gathering with no updates. The pseudocode should either reflect this more accurately or exclude this logic, considering it as an implementation detail.

Author response:
Thanks for pointing all these out! We have updated the manuscript accordingly.

Reviewer comment (Presentation):
Figure 3: It would be more accurate to label the vertical axis "Average Episodic Return", as the reward is just a momentary quantity.

Author response:
Thanks for pointing all these out! We have updated the manuscript accordingly.

Reviewer comment (Writing):
Section 5.2: "action-spercialized"?

Author response:
Thanks for pointing all these out! We have updated the manuscript accordingly.

Reviewer comment (Writing):
Section 5.3: "that just do random shooting" — fix grammar.

Author response:
Thanks for pointing all these out! We have updated the manuscript accordingly.

Reviewer comment (Writing):
Section 5.3: The last paragraph needs a grammatical revision.

Author response:
Thanks for pointing all these out! We have updated the manuscript accordingly.

Reviewer comment (Writing):
Section 5.3: The last paragraph could do a better job explaining why Random Shooting fails.

Author response:
Thanks for pointing all these out! We have updated the manuscript accordingly.

Reviewer comment (Theory):
This algorithm seems to use a heavy amount of computation. Can you comment on the algorithm's complexity?

Author response:
We have included a new Table in appendix 8.6 that compares the clock time of planning and the number of timesteps needed to converge as the action space expands. We tested this on the Hard Move domain, where the number of discrete actions changes from $2^4$ to $2^{10}$ (one continuous parameter for each of them). As shown in the table, while the number of samples increases as the action space expands, it’s still within an acceptable range even when it’s extremely large. We also quantify this with theoretical analysis which is included as a new section in the updated manuscript.

Reviewer comment (Novelty):
The novelty seems rather limited as the main framework of DLPA follows the standard data-driven MPPI, with the contextualization from MDP to PAMDPs. For example, an information theoretic MPC framework was proposed in [G. Williams et al, ICRA 20217] to incorporate the model learning in the MPPI-based planning procedure.

Author response:
The primary contribution is the first model-based RL algorithm for the PAMDP setting, which outperforms all the previous model-free RL methods for PAMDP on 8 standard PAMDP benchmarks. The PAMDP model is appropriate to many very common real-world tasks, and has gained a lot of attention in recent papers - however, all existing methods are model-free.  We hope our paper can draw other researchers' attention to the model-based aspect of this problem setting since DLPA performs better than even the most sophisticated model-free methods. In terms of novelty, inspired by the reviewer’s question, **we have added a new theoretical analysis section in the updated manuscript where we derive some performance guarantees / bounds for the proposed algorithm.** Thank you for the suggestion!

For using MPPI for PAMDPs, previous works have only shown its efficiency on either discrete action space or continuous action space, while we’re the first to apply such a method to the parameterized action space. **The main modification we make when applying MPPI to PAMDPs is that we keep a separate distribution over the continuous parameters for each discrete action and update them at each iteration instead of keeping just one independent distribution for the discrete actions and one independent distribution for the continuous parameters.** In other words, we let the distribution of the continuous parameters condition on the chosen discrete action during the sampling process. This is an important change to make when using MPPI for PAMDPs as we don’t want to throw away the established dependency between the continuous parameter and corresponding the discrete action in PAMDPs.  In the updated manuscript we have included a new ablation study result for this change which shows that this change is indispensable for DLPA to do proper planning in PAMDP (Section 6.3).

Reviewer comment (Reproducibility):
Some key information is not provided, e.g. which neural network algorithms are used to optimize Eq. 1?

Author response:
To solve Eqn. (1), we simply learn three neural networks for the reward predictor, transition predictor and termination predictor respectively using supervised learning. And we use the mean square error as the loss of the state and reward prediction, and use cross entropy loss as the objective function for updating the termination signal. We train our model based on multi-step prediction (only the ground truth state at step 1 is needed) to encourage the model to do better H-step long-term prediction. This design for the training process is consistent with the planning process, as H is also the planning horizon in MPPI where we do the same long-term prediction to select actions. As also pointed out by reviewer WRCa, this choice allows gradients to flow back through time and assign credit more effectively than the alternative, which is one-step training. We found this design choice to be especially useful during empirical evaluation as shown in Section 6.3.  We have also included one additional experimental results plot on Platform to show this in the updated manuscript.

Reviewer comment (Theory):
No theoretical analysis is given to justify the performance applied to PAMDPs, which seems necessary to help readers understand the significance of the DLPA beyond simple contextualization of data-driven MPPI to PAMDPs. As PAMDPs concerned about finding the right action and the associated parameters to use, is there any guarantee of convergence to the selected parameterized actions during training from DLPA?

Author response:
We thank the reviewer for pointing this out. We have added a theoretical analysis section (also in the appendix) in the updated manuscript (section 5). In general, we provide three main theorems that describe the performance guarantees of DLPA from three aspects: 1. A theoretical bound for the cumulative return difference between the rollout trajectories of DLPA and the optimal trajectories. 2. A theoretical bound for the multi-step prediction error of DLPA. 3. A theoretical bound for how the estimation error will change with respect to the number of samples and the cardinality of the continuous parameter space.

Reviewer comment (Novelty):
What is the main contribution behind DLPA, or could authors comment on specific challenges solved when using MPPI for PAMDPs?

Author response:
The primary contribution is the first model-based RL algorithm for the PAMDP setting, which outperforms all the previous model-free RL methods for PAMDP on 8 standard PAMDP benchmarks. The PAMDP model is appropriate to many very common real-world tasks, and has gained a lot of attention in recent papers - however, all existing methods are model-free.  We hope our paper can draw other researchers' attention to the model-based aspect of this problem setting since DLPA performs better than even the most sophisticated model-free methods. In terms of novelty, inspired by the reviewer’s question, **we have added a new theoretical analysis section in the updated manuscript where we derive some performance guarantees / bounds for the proposed algorithm.** Thank you for the suggestion!

For using MPPI for PAMDPs, previous works have only shown its efficiency on either discrete action space or continuous action space, while we’re the first to apply such a method to the parameterized action space. **The main modification we make when applying MPPI to PAMDPs is that we keep a separate distribution over the continuous parameters for each discrete action and update them at each iteration instead of keeping just one independent distribution for the discrete actions and one independent distribution for the continuous parameters.** In other words, we let the distribution of the continuous parameters condition on the chosen discrete action during the sampling process. This is an important change to make when using MPPI for PAMDPs as we don’t want to throw away the established dependency between the continuous parameter and corresponding the discrete action in PAMDPs.  In the updated manuscript we have included a new ablation study result for this change which shows that this change is indispensable for DLPA to do proper planning in PAMDP (Section 6.3).

Reviewer comment (Reproducibility):
What algorithms are used to solve Eq. 1 in terms of model learning? Also, it is unclear why the authors imply that the ground truth state was not used as input to train their model in DLPA. For example, the ground-truth state s_{t+1} is needed to compute the loss L_joint defined in Eq. 1.

Author response:
To solve Eqn. (1), we simply learn three neural networks for the reward predictor, transition predictor and termination predictor respectively using supervised learning. And we use the mean square error as the loss of the state and reward prediction, and use cross entropy loss as the objective function for updating the termination signal. We train our model based on multi-step prediction (only the ground truth state at step 1 is needed) to encourage the model to do better H-step long-term prediction. This design for the training process is consistent with the planning process, as H is also the planning horizon in MPPI where we do the same long-term prediction to select actions. As also pointed out by reviewer WRCa, this choice allows gradients to flow back through time and assign credit more effectively than the alternative, which is one-step training. We found this design choice to be especially useful during empirical evaluation as shown in Section 6.3.  We have also included one additional experimental results plot on Platform to show this in the updated manuscript.

Reviewer comment (Theory):
Does DLPA deliver any theoretical properties to justify its performance on PAMDPs? Authors are encouraged to discuss the convergence of parameterized actions from the DLPA framework, if any.

Author response:
We thank the reviewer for pointing this out. We have added a theoretical analysis section (also in the appendix) in the updated manuscript (section 5). In general, we provide three main theorems that describe the performance guarantees of DLPA from three aspects: 1. A theoretical bound for the cumulative return difference between the rollout trajectories of DLPA and the optimal trajectories. 2. A theoretical bound for the multi-step prediction error of DLPA. 3. A theoretical bound for how the estimation error will change with respect to the number of samples and the cardinality of the continuous parameter space.