Thank you for taking the time to thoroughly review our work and provide valuable insights and constructive feedback.

Here we try to address your questions and concerns.

Experiments

• Figures clarification
  • Figures 3 and 4: we clarified (in red) the meaning of each curve. Sobolev (blue) indicates we use gradient information to learn the function or distribution, No Sobolev (red) only uses the scalar information from the dataset. More specifically about the dashed lines: it is a setup where a similar architecture to the distributional setting is used except it is simply L2 regression.
  • Figure 1: the blue dots are indeed hidden behind the other curves. Blue are samples from the true 5 modes of the distribution over functions. You can find clearer examples in Figure 7 and 8 (in the Appendix). We will make sure to change the Figure to something clearer if needed.
• Fairness we acknowledge we only showcase our method on $0.25 \times 10^6$ training steps. However, as we wrote in the experiment section, we use 512 actors to speed up convergence similarly to [1]. Furthermore, as you can see most methods reached convergence before hitting the maximum number of steps.
• Pre-trained transition information we do not pre-train our world model but learn it concurrently to the policy and critic.
• Incorporating gradient information into distributional modelling we argue this is the idea scenario as policy improvement is carried out from the critic's action gradient. Learning a more accurate gradient as in [2] and then a distribution over that gradient is done to help policy improvement by more accurate policy evaluation.

Random state-action Sobolev return

• It is unclear whether Sobolev returns are actually used for decision making They are modeled during policy evaluation. The modeled Sobolev return distribution is then sampled from during the policy improvement step, where the action-gradient part is used to provide an improvement direction to the policy. We further clarify our method is value-based. This process is described in equation 8 except in practice we draw samples from a replay buffer.
• Eq. 15 We clarified equation 15 and put a small derivation in the appendix.
• Handling of complex state-spaces indeed learning world-model for complex state spaces might not be trivial and the estimated state-gradient might not be accurate. It is a limitation that we do not tackle in this work but we point to methods that use variational inference and cVAEs to do transition modelling in a latent space instead of the state-space directly [4]. This could be the topic of further work.

Discussion with other methods

• The works you mentioned are cited in the Section 3.3 after the Distributional Sobolev critic paragraph. As we wrote, neither QR-DQN nor C51 were applicable to learning a distribution over gradient. Our framework is similar to learning a distribution over multi-variate returns [5]. Learning quantiles or bins isn't tractable over many dimensions thus we favored a method based on generative modelling and design our distributional Sobolev critic as a geneative model that minimizes a distributional discrepancy (MMD) that can be tractacbly estimated via samples.
• Kernel tuning: Indeed, tuning the kernel and its hyperparameters is critical. We used a value we found work well on most environments which is also the same as in [3]

Questions

• Computational cost Indeed, modeling the return distribution using generative modeling is more computationally expensive than via pseudo-samples or quantiles for example. This is a property we acknowledge in the conclusion. One path towards reducing that issue would be to share the forward pass between the K samples. We investigated that idea but did not succeed. Furthermore, as far as we are aware, there is no other obvious way to model a random gradient (using Sobolev inductive bias we mentionned in the submission) than via generative modelling and the reparametrization trick.
• Diffusion Diffusion is a possible avenue for the world model. However, the computational cost will grow dramatically. As we need to differentiate the compute graph of the generative model with respect to the conditioning variable. Diffusion models implicitly build very deep compute graphs. Furthemore, our method for gradient estimation (of the environment) relies on inference. Switching to diffusion would probably mean we would rather use the world-model to imagine new samples (that would be differentiable) instead of reconstructing/inferring.
• Have the authors considered implementing Distributional Sobolev RL on top of value-based RL? We are confused as we did implement Distributional Sobolev training on a value-based baseline (DDPG).

We hope the revisions address your concerns and provide further clarity on the novelty and contributions of our method. Thank you once again for your time and constructive comments.

Thank you for taking the time to thoroughly review our work and provide valuable insights and constructive feedback.

Here we try to address your questions and concerns.

• Motivation we clarified our theoretical motivation by adding Proposition 3.1 in the approach section as well as some more explanations. We do not use the gradient information in the policy optimization but in the policy evaluation in order to learn a more accurate distributional critic. We also provide a more general motivation in the general reply.
• Extension of D4PG [1] D4PG did not study the distribution over the gradient. Our work can be seen as a Sobolev extension of this work. Extending it required a new framework for multi-variate distribution RL akin to [2].
• Methodological contribution Indeed Distributional RL and Sobolev training are known methods: but this work is the first to propose a combination of the two which required a novel design for the distributional critic and the distributional discrepancy measure (we use MMD on this new application of (Distributional Sobolev)).
• Limitation of deterministic policy gradient Indeed, this is a limitation of our work but studying stochastic policy gradient would be out of scope for this paper. Our baselines are D4PG which is a distributional extension of DDPG and MAGE which is a Sobolev extension of DDPG. Furthermore, we work on value-based methods which exclude methods like PPO/A2C.
• Limitation of cVAE We are quite confused by this sentence as a great part of the literature on model-based reinforcement involves variational inference and cVAEs [3] [4]
• Likelihood estimation issue As we explained in this paragraph, modelling a distribution on a gradient-space isn't trivial as we want to preserve the Sobolev inductive bias. Since we want to model a gradient using the gradient of a neural network, we cannot easily put a probability distribution that would make density estimation possible. This rules out many methods for generative modelling such as normalizing flows or simply mixture of Gaussians as we explained. The reason is straightforward: even if one can model the likelihood of some value of a random variable, there is no obvious way to estimate the likelihood of some value of its gradient (with respect to a conditioning variable). Thus we must rely on a generative modelling scheme that learns only on the base of true samples and generated samples with not densities available.
• Meaning of the derivative of a random variable in Eq. 8 We did not invent this equation. This is the standard distributional deterministic policy gradient from D4PG paper [1]. We added a few words below the confusing equation and added an appendix section.
• Incomplete references we corrected all the references.
• Supervised learning experiments we argue theses experiments are useful to motivate the primary novely of this paper: distributional Sobolev training. We show that our implementation of this framework allows to effectively use gradient information to learn a more accurate distribution. We also showcases that a distributional estimation can have benefitial properties in the low data regime.

We hope the revisions address your concerns and provide further clarity on the novelty and contributions of our method. Thank you once again for your time and constructive comments.

[1] Distributed Distributional Deterministic Policy Gradients https://arxiv.org/abs/1804.08617

[2] Distributional Reinforcement Learning for Multi-Dimensional Reward Functions https://openreview.net/pdf?id=u7oKU1iXTa9

[3] MASTERING ATARI WITH DISCRETE WORLD MODELS https://arxiv.org/pdf/2010.02193

[4] World Models https://arxiv.org/pdf/1803.10122

Thank you for taking the time to thoroughly review our work and provide valuable insights and constructive feedback.

Here we try to address your questions and concerns.

• Equation 11 vs equation 16 we removed equation 16 as learning the state-action gradient is an additional but non-necessary improvement to the framework as it is and it brings confusion. In the removed paragraph, we wrote this to argue that as the state space is sometimes much larger in dimensionality than the action space, learning the state part of the gradient may bring more information from the world-model. Also, environments may exhibit different types of distribution whether on the state part or the action part of the return's gradient which could or could not be benefitial to model.
• Theoretical justification we added a part in the approach section that motivates our framework (Section 3.1).
• Lack of rigor in the definitions We added a section in the appendix for the derivation of the new operator defined in Eq 18. May we ask what parts you would like to see improved/clarified ?
• Existence of the gradient(s) we were not clear enough about this. Our work only tackles continuous state-action spaces. It can be seen as a Sobolev extension of D4PG [2]. Furthermore, we assume this gradient exist and that we can calculate its value. It is an assumption that we then relax using the world-model similarly to [1]. We added a few words in the approach section (in red) to make this clearer.
• Meaning of bootstrapped target we substantially modified this paragraph. It is hopefully clearer now. The bootstrapped target was related to re-using the current estimator for the distributional critic in the next state-action pair as in temporal-difference. We removed this phrasing as it brought more confusion. We also added a small derivation for the new operator in the appendix. We indeed assume it is differentiable. We made it clearer we start from this assumption and relax it later using the world-model similarly to [1].
• $\nabla_s$ and $\nabla_a$ the gradient is taken with respect to the application $a_0=a$. We added a sentence to clarify this. \item Why learning the return distribution along the gradient ? First we clarify that we learn the joint distribution over the returns and their gradients. For the motivations we refer to the general reply above.
• Related works We merge the related works with the introduction as was suggested by another reviewer. Furthermore, we mention several related works in the Section 3.3.

We hope the revisions address your concerns and provide further clarity on the novelty and contributions of our method. Thank you once again for your time and constructive comments.

[1] How to Learn a Useful Critic? Model-based Action-Gradient-Estimator Policy Optimization

[2] Distributed Distributional Deterministic Policy Gradients https://arxiv.org/abs/1804.08617

[1] Distributed Distributional Deterministic Policy Gradients https://arxiv.org/abs/1804.08617

[2] Dyna, an integrated architecture for learning, planning, and reacting

[3] Learning Continuous Control Policies by Stochastic Value Gradients https://arxiv.org/abs/1510.09142

[4] How to Learn a Useful Critic? Model-based Action-Gradient-Estimator Policy Optimization

Thank you for taking the time to thoroughly review our work and provide valuable insights and constructive feedback.

We will try to address your concerns and then answer your questions.

Concerns

• Performance: We acknowledge the limited empirical performance of our method and are actively trying to improve it. We suspect the issue lies in the conditional VAE. As we rely on variational inference to infer gradients from true observations as explained in Appendix A.2 and A.3, the performance of the cVAE is critical. It must provide good reconstruction to query in the correct next-state and provide a accurate enough reward. However, just as critical is new sample generation. Indeed, since the gradients are inferred, the posterior (from the encoder) must be close to the prior while providing good reconstruction (which is related to good samples generation). This is a usual balance to find for Variational Autoencoder that has been particularly difficult to strike in this work.
• Environment stochasticity: Regarding the stochasticity of the environments, we used N-step return (with N = 5) which makes the dynamics highly stochastic from the perspective of policy evaluation. Indeed, as explained in Section 2.2, N-step return uses the exploration policy for N steps which makes the mapping $(s, a) \rightarrow (s', r)$ non-deterministic even if it originally was. We further motivate this from D4PG [1] which showcased empirical improvements once making the environment stochastic this way.
• Baselines About the baselines, we argue that only few are needed to motivate the empirical improvement. Indeed, DSDPG is not a model-based algorithm like the Dyna family [2]. We do not use the world-model for imagination but to provide a differentiable surrogate of true observations via inference akin to [3]. The primary reason is that we wanted our implementation of Distributional Sobolev Deterministic Policy Gradient to be as minimal as possible. Hence, we claim that only the following baselines are needed

• A scalar-based distributional method (MMD)
• A scalar-based deterministic method (DDPG)
• A gradient-based deterministic method (MAGE) [4]

Related works

We merged this section with the introduction. Several related works are also discussed in Section 3.3.

Clarity

Novelty We added a paragraph about the contributions in the introduction and in the general reply above. Indeed, our method can be seen as

• A distributional generalization of MAGE [4] that used Sobolev training in value-based RL and/or
• A Sobolev generalization of D4PG [1]

It required developping new tools to model gradient as a random variable along random returns.

Mathematical section We modified Section 3.2 heavily and tried to make it clearler what is new.

Questions

• What are the paper's contributions ? We refer to the general reply and the added paragraph in the introduction.
• Is it possible to extent your method to support other actor-critic methods ? probably but it would certainly not provide the same benefits as on value based methods such as DDPG/TD3/SAC. Indeed, we improve the critic by learning a distribution over gradient as the critic's action gradient is the one quantity that is used for policy improvement in this family of methods. The usual policy gradient and its PPO variant only use the critic through its scalar output (or the conditional mean in the case of distributional RL) which could be made more accurate by learning the gradient. Furthermore, as it seems other reviewers had similar question: our method is primarly designed for continuous action spaces. It would then rule out the discrete applications of PPO/A2C.
• In your experiments, it seems that you evaluate the learning curve, while eventually most methods perform on par after enough interaction. Have you tried limiting the data (to show better sample complexity) ? we are not sure exactly what you propose. We implemented the usual training loop where the more training steps the more samples from the observations are collected and buffered. Would you mean reducing the number of actors (for exploration) or doing more learning steps per exploration step ? In which case, we didn't investigate in that direction.
• One of the advantages of using a model-based method, given accurate enough model, is the ability to train without direct interaction (only from the world model) have you tried it ? we have not explored this possibility as our world-model is only used to infer gradient from true observations and not imagination as in Dyna [2]. Although this could be possible, this would be an orthogonal modification to the current framework. We stress again that our implementation was minimal and a proof of concept that DSDPG can work.

We hope the revisions address your concerns and provide further clarity on the novelty and contributions of our method. Thank you once again for your time and constructive comments.