We sincerely thank you for the thoughtful and constructive feedback. Please find our detailed responses to the comments below.

1. Supplemental experiments of MPGE

Thank you for raising this important point. To address your concern, we have conducted additional walltime experiments to assess the computational cost of MPGE. The updated results are provided below:

The results show that MPGE incurs moderate overhead compared to standard off-policy methods due to the added diffusion and world model components. However, this overhead is effectively controlled through two mechanisms:

1. We perform MPGE training only once every 10 environment steps, significantly reducing overall compute.

2. Unlike PGR-style methods that require retraining the generative model to convergence before use, MPGE supports online training with incremental updates, which simplifies training and improves efficiency.

We will include this discussion and the corresponding results in the final version of the paper.

Additionally, we appreciate your suggestion regarding baseline selection. We agree that including recent off-policy exploration methods such as [1–4] would provide a more comprehensive evaluation. We conduct the experiments in DMC Humanoid tasks, which are one of the hardest tasks within 3 seeds, and the results are illustrated as follows:

As shown in this task, except for the simplest case (Humanoid-Stand), MPGE achieves the best performance across all more challenging tasks.

2. Discussion of the Encoder-decoder structure in MPGE

Thank you for this insightful question. We would like to clarify that MPGE is developed under the MDP setting. And why we choose to model dynamics in a latent space instead of directly learning the mapping $f(s,a)→s′$ is motivated by several practical considerations:

1. Learning a compact latent representation helps capture abstract, task-relevant features of the environment, improving generalization and training efficiency even under MDPs. Besides, different dimensions of the state may contribute unevenly to dynamic modeling. Using an encoder to transform the raw state allows the model to extract the most relevant components, which facilitates more accurate transition prediction.

2. Latent dynamics are often smoother and easier to predict compared to raw state transitions, especially in high-dimensional environments. Fitting dynamics in a latent space typically improves the accuracy of transition modeling.

3. Decoupling representation learning from dynamics prediction allows better reuse and transfer of components. In fact, in MPGE, our policy network is built on top of the same encoder, enabling a unified state representation that facilitates more effective policy learning.

To further support our claim, we conducted a simple experiment across three environments (each with 3 random seeds), using the same diffusion and baseline algorithm setup, to compare the TAR between vanilla dynamics (Using Transformer as predictor of next states and rewards) and latent dynamics. The results are as follows:

The results show that, compared to directly mapping the current state and action to the next state, introducing a latent space leads to better learning of both environment transitions and state representations in the policy network.

3. Comparison with different exploration baselines

Thank you very much for the valuable suggestion. In response, we have conducted additional experiments to study the effect of different exploration strategies on the DSAC algorithm via ablation comparisons. Due to the development and implementation effort involved, and time constraints, we focus on three representative environments—the Humanoid tasks in DMC—for this evaluation.

Before presenting the results, we provide the following clarifications regarding the methods considered in this comparison:

1. Plan2Explore: Due to the limited compatibility of the original codebase, we re-implemented this method within our own framework. Specifically, we replaced the world model with Dreamer, and trained five independent predictors to construct a latent disagreement signal. An exploration policy was then trained using this signal to collect additional data for the main task policy.
2. EIPO: All experiments of EIPO are conducted on on-policy algorithms, specifically based on PPO, which poses challenges when attempting to transfer the method to off-policy algorithms such as DSAC. Although EIPO is theoretically algorithm-agnostic, applying it to off-policy methods (e.g., DSAC) poses practical challenges. Its core mechanism relies on alternating optimization between two policies and dynamically adjusting intrinsic reward via Lagrangian updates. However, off-policy training uses replay buffers, making it difficult to accurately evaluate both policies’ extrinsic returns for updating the multiplier. Additionally, the assumption of policy closeness often fails in off-policy settings, reducing surrogate objective reliability. Due to these challenges and time constraints, we were unable to include EIPO in our experiments. We appreciate your understanding.
3. MaxInfoRL: The method is highly transferable, and we directly integrated its core components into our own codebase. We also attempted to build DSAC on top of the original implementation, but the performance was not as good as our reimplementation. Therefore, we report results based on the version implemented in our framework.
4. OMBRL: Since the original implementation is in JAX, we reimplemented the method in our own codebase by replacing the SAC component in MBPO with DSAC. We use Deep Ensembles to compute $σ_n​(x,a)$, with each network sharing the same architecture as the predictor used in Plan2Explore.

The final results are as follows:

The results show that MPGE still offers a significant advantage compared to these methods. Intrinsic exploration methods encourage visiting novel or unpredictable states, but their signals (e.g., prediction error, uncertainty) are often task-agnostic and may lead to uninformative or misaligned exploration. In contrast, MPGE leverages task-aware utility functions (e.g., TD-error, entropy) to generate states that are directly aligned with policy improvement objectives. Besides, unlike intrinsic reward methods that require careful reward balancing and affect the policy's optimization objective, MPGE decouples exploration from reward design, enabling more stable training without interfering with the task-specific learning signal.

4. Clarification of Theorem 1

We fully understand your concern regarding this statement.

In fact, the purpose of this theorem is to ensure that the generated states, which are learned based on the state distribution in the replay buffer, can gradually follow a steady-state distribution. This is motivated by existing studies[4-5] showing that unstable or non-convergent training distributions in generative models often lead to degraded or uncontrolled generation quality. Our theoretical analysis aims to provide a principled guarantee that the generated states remain within the environment’s reachable state space, thereby ensuring that they can effectively support policy exploration and improvement.

As stated in line 532 of the paper, we also provide conditions under which the conclusion still holds even if the policy does not satisfy finite-step convergence. And the actual assumption made here is that the policy changes diminish over time as optimization proceeds. Importantly, the final stationary policy distribution does not need to correspond to an optimal policy (our derivation does not rely on any properties of the optimal policy, such as its return being the highest). In our paper, Lastly, we would like to reiterate that this is merely an idealized assumption. Our theoretical result is not tied to the convergence of the optimal policy, but only to its convergence rate, at which policy updates diminish over time.

We acknowledge that our previous wording may have caused confusion, and we will clarify it accordingly in the final version. We are sincerely grateful for your thoughtful question and suggestion.

Conclusion

We sincerely thank you for your thoughtful and constructive feedback. In response, we have conducted substantial additional experiments and analyses to directly address your concerns. We hope these efforts help clarify the points in question, and we would be truly grateful if you would consider them in your consideration for a higher score.

Dear Reviewer DNyW,

Please provide your feedback to the authors as early as possible after reading the rebuttal and other reviews, so there is enough time for back-and-forth discussion. Thank you!

Best, Area Chair

We sincerely appreciate your detailed and constructive feedback. Below are our responses to your concerns.

1. Explanation of Equation 5

Equation (5) follows the standard gradient decomposition for classifier-guidance diffusion models [1]. It is used to guide the generation of states toward those favored by the target function. The guidance scale coefficient $\omega$ controls the strength of this guidance, balancing between following the unconditional diffusion prior and aligning with the target utility function.

2. Description of the two utility functions

We appreciate the reviewer’s suggestion and agree that these concepts are better introduced earlier. In the revised version, we will move the descriptions of policy entropy and TD error to the Preliminaries section for clarity. We also clarify that these are example utility functions within our framework, and other choices are equally applicable to support the generation of diverse and informative critical states.

3. Swapped label

Thank you very much for pointing this out! We confirm the labels are indeed swapped and will correct this in the final version.

4. About Equation 10

Thank you for the insightful question. We follow the design principle introduced in Dreamer series[2-3], where stop-gradient (sg) is applied to avoid representational collapse and to stabilize joint training of the encoder $g_\phi$ and the dynamics model $d_\phi$. Specifically, the sg operator ensures that the prediction target $g_{s_{t+1}}$ is treated as fixed when training the dynamics model. At the same time, the target is treated as fixed when training the encoder, so that the encoder is not optimized to merely match the dynamics model output, but instead learns representations that are useful for downstream tasks such as value prediction or policy learning.

In summary, this design is implemented to prevent trivial solutions where both modules co-adapt to minimize the loss without learning meaningful structure. Moreover, Dreamer demonstrates that using such asymmetric losses with sg leads to more stable and scalable representation learning in model-based RL. Empirically, we observe that this design promotes stable convergence across tasks, consistent with prior works.

5. Extension

Yes, our method can be extended to discrete action spaces. As long as a differentiable mapping from given utility functions to states can be established, it remains feasible to compute utility functions under a state, resulting in critical generation. While our current experiments focus on continuous control tasks, we plan to explore more discrete-action environments in future work.

Q&A1: Clarification of Theorem 1

We fully understand your concern regarding this statement.

In fact, the purpose of this theorem is to ensure that the generated states, which are learned based on the state distribution in the replay buffer, can gradually follow a steady-state distribution. This is motivated by existing studies[4-5] showing that unstable or non-convergent training distributions in generative models often lead to degraded or uncontrolled generation quality. Our theoretical analysis aims to provide a principled guarantee that the generated states remain within the environment’s reachable state space, thereby ensuring that they can effectively support policy exploration and improvement.

As stated in line 532 of the paper, we also provide conditions under which the conclusion still holds even if the policy does not satisfy finite-step convergence. And the actual assumption made here is that the policy changes diminish over time as optimization proceeds. Importantly, the final stationary policy distribution does not need to correspond to an optimal policy (our derivation does not rely on any properties of the optimal policy, such as its return being the highest). In our paper, Lastly, we would like to reiterate that this is merely an idealized assumption. Our theoretical result is not tied to the convergence of the optimal policy, but only to its convergence rate, at which policy updates diminish over time.

We acknowledge that our previous wording may have caused confusion, and we will clarify it accordingly in the final version. We are sincerely grateful for your thoughtful question and suggestion.

Q&A2: Comparison of the on-policy methods

We have reported comparisons with two representative on-policy methods, PPO and TRPO, in Appendix C.2. Please refer to that section for detailed results.

Q&A3: Discrete action spaces generalization

Yes, the approach can be generalized to discrete action spaces. Please refer to our earlier response above for a detailed explanation.

Conclusion

We truly appreciate your thoughtful and constructive feedback. We hope that our detailed responses and additional efforts have addressed your concerns, and we kindly ask that you take them into consideration during your re-evaluation.

[1] Dhariwal P, Nichol A. Diffusion models beat GANs on image synthesis[J]. Advances in neural information processing systems, 2021, 34: 8780-8794.

[2] D. Hafner et al. Learning latent dynamics for planning from pixels. ICML, 2019.

[3] D. Hafner et al. Mastering diverse control tasks through world models. Nature, 2025.

[4]Argenson A, Dulac-Arnold G. Model-based offline planning[J]. arXiv preprint arXiv:2008.05556, 2020.

[5]Luo C. Understanding diffusion models: A unified perspective[J]. arXiv preprint arXiv:2208.11970, 2022.

Dear Reviewer ZXF4,

Please provide your feedback to the authors as early as possible after reading the rebuttal and other reviews, so there is enough time for back-and-forth discussion. Thank you!

Best, Area Chair

We sincerely thank you for the thoughtful and constructive feedback. Please find our detailed responses to the comments below.

1. Qualitative Analysis of MPGE

We will address your concerns through theoretical analysis from two perspectives: novelty and dynamic compliance. Specifically, we aim to clarify the following questions:

1. Why do the actions or transitions generated by MPGE exhibit novelty, and how does this benefit policy optimization and evaluation?

2. How does MPGE ensure compliance with environmental dynamics when generating transitions?

3. Why are the transitions generated by MPGE more effective than those produced by methods like PGR, ultimately leading to improved policy performance?

1.1 The novelty of transitions

The joint distribution of transitions generated by MPGE can be derived as follows:

$p(s,a,s^{\prime},r)=p(s)p(a|s)p(s^{\prime},r|s,a)\quad(1)$

Based on the structure of MPGE, we can further decompose the source of its novelty into three main components:

(i) State-Level Generation Instead of Transition Resampling

MPGE generates high-utility states directly and then uses the world model to simulate plausible transitions. This contrasts with prior methods such as PGR, which often resample or perturb transitions from the buffer. By targeting utility-defined regions, MPGE explores more structurally diverse transitions.

(ii) Dynamically Shifting Utility Landscape During Training

The utility signal in MPGE evolves as training progresses, continuously shifting the generation target. This dynamic conditioning promotes continual novelty, as the generator focuses on states that are currently difficult for the policy, rather than statically defined regions.

(iii) Policy-Centric Guidance Instead of Buffer Imitation

Crucially, MPGE does not generate data based on the static behavior policy distribution in the buffer. Instead, the generation is guided by utility functions evaluated under the current policy, such as current-policy TD error or entropy. This ensures that the generated transitions align with the evolving learning needs of the agent, rather than being confined to outdated or suboptimal trajectories. The world model further enables rollouts from generated states, allowing MPGE to synthesize transitions that the current policy is likely to visit but has not yet explored, thus supporting meaningful off-policy generalization.

Summary:

MPGE leverages utility functions conditioned on the current policy to generate reachable but under-visited states with long-tail characteristics. Rather than producing out-of-distribution samples, it rebalances in-distribution critical states to align with the training objective.

As a result, transitions generated by MPGE are more novel and task-relevant compared to methods like PGR, which imitate the buffer and remain constrained by the behavior policy distribution.

1.2 The dynamical feasibility of transitions

(i) analysis of the dynamical feasibility in MPGE

In our setting, the world model is pretrained using data collected from a random policy, ensuring broad coverage of the state-action space and capturing generalizable environment dynamics rather than overfitting to narrow trajectories. This is a standard practice in offline model-based RL to ensure that rollouts remain dynamically plausible.

Importantly, MPGE generates only states, and transitions are simulated through the learned dynamics model. As such, all transitions remain on-manifold with respect to the model, provided it is well-trained. This follows the same principle as in Dreamer [1] and related work, where stochastic data and latent dynamics promote reliable predictions.

(ii) Comparison with methods like PGR

In fact, compared to MPGE—which explicitly assembles transitions using the environment dynamics and ensures that the resulting buffer complies with the true transition function—methods that imitate and resample transitions directly from the replay buffer are more likely to violate dynamics consistency.

Take PGR as an example: PGR models the joint distribution $p(s,a,r,s^{{\prime}})$ by learning directly from tuples in the buffer. However, this paradigm does not enforce any structural constraints that the transitions must satisfy the environment dynamics. This lack of constraint makes the generative model prone to overfitting, potentially capturing spurious correlations between variables rather than learning the true mapping from $(s, a)$ to $(s^{{\prime}}, r)$. For example, although both approaches model joint distributions $p(s, a,s^{\prime},r)$, PGR may end up modeling a distribution of the form such as $p(s)p(a)p(s^{{\prime}})p(r)$, which does not reflect the causal or functional dependencies present in the environment and may results in poor policy evaluation.

This observation is further supported by our experiments below, which demonstrate that MPGE is more capable of generating transitions that are dynamically consistent with the environment, in contrast to methods that rely solely on joint distribution modeling without enforcing transition constraints.

2. Quantitative Analysis of MPGE (Supplementary Experiments)

To validate our analysis, we conducted an experiment to assess the dynamic compliance of transitions generated by MPGE and PGR. We trained a dynamics discriminator that classifies whether a given tuple $(s,a,s ′)$ aligns with the true environment dynamics.

The discriminator was trained on positive samples from the replay buffer and a random policy—both guaranteed to be valid. To ensure fairness, MPGE and PGR were trained independently, each generating an equal number of negative samples, while their learning was not influenced by the discriminator’s feedback.

Structure of the discriminator:

To mitigate potential overfitting as much as possible, we introduced an information bottleneck in the design of the dynamics discriminator. Specifically, we adopt the architecture where the input pair $(s, a)$ is first processed through a multi-layer MLP, which is then directly connected to the encoded representations of $s^{{\prime}}$. This combined representation is subsequently passed through a lower-dimensional MLP before producing a softmax-based binary classification output.

   def forward(self, s, a, s_prime):
       sa_feat = self.mlp1(torch.cat([s, a], dim=-1))
       sp_feat = self.mlp2(s_prime)
       fused = torch.cat([sa_feat, sp_feat], dim=-1)
       out = self.fusion_discriminator(fused)
       return out 

Results:

We conducted a joint training experiment on the Humanoid-Run with 3 seeds, where the dynamics discriminator, PGR, and MPGE were trained simultaneously. The resulting figure summarizing the method's accuracy (the ratio of the number of samples classified as real to the total number of test samples) over time is shown below:

The results show that as training progresses, MPGE consistently generates a larger proportion of transitions that do not violate the environment dynamics constraints, thereby demonstrating the higher quality and plausibility of its generated data.

Q&A1: Discussion on guidance scale

The guidance scale plays a critical role in balancing sample diversity and target alignment. In our current implementation, we select $\omega$ through cross-validation on a held-out set (or via grid search across a small range, e.g., $[0.2,0.5,1.0,2.0]$) and apply the same value across all environments for consistency.

In our ablations, $ω=1$ performs best under standard training steps. However, with extended training (e.g., 3M steps), smaller scales (e.g., $ω<1$) also achieve strong performance, differing mainly in convergence speed. Larger scales tend to distort the generated data. For simple tasks, a small $ω$ is sufficient, while complex tasks may benefit from adaptive schemes such as annealing.

Q&A2: Discussion on hyperparameter

We fully understand the reviewer’s concern regarding generalization and hyperparameter sensitivity. To address this:

1. We used the same set of hyperparameters across all tasks on both benchmarks (DMC and Gym MuJoCo), without any task-specific tuning for MPGE.

2. The world model parameters are directly adopted from existing implementations of DreamerV3[1] and STORM[2], and the diffusion model settings follow standard DDPM[3] configurations—no further tuning was applied.

3. For our method-specific hyperparameters $λ$ and $ω$, we conducted ablation studies, and then applied the same values across all tasks.

These results demonstrate that MPGE maintains strong performance without extensive tuning, and we believe this reflects its robustness across diverse environments.

Conclusion

We sincerely appreciate your detailed and constructive feedback. We hope our responses address your concerns and kindly request your consideration for a higher score.

[1] D. Hafner et al. Mastering diverse control tasks through world models. Nature, 2025.

[2] Zhang W, Wang G, Sun J, et al. Storm: Efficient stochastic transformer-based world models for reinforcement learning[J]. Advances in Neural Information Processing Systems, 2023, 36: 27147-27166.

[3] Ho J, Jain A, Abbeel P. Denoising diffusion probabilistic models[J]. Advances in neural information processing systems, 2020, 33: 6840-6851.

Dear Reviewer W4cH,

Please provide your feedback to the authors as early as possible after reading the rebuttal and other reviews, so there is enough time for back-and-forth discussion. Thank you!

Best, Area Chair

We sincerely appreciate your detailed and constructive feedback. Below are our responses to your concerns.

1. Discussion of the utility function with Policy entropy

We sincerely appreciate your insightful comment. Indeed, policy entropy is only one possible choice of utility function within the MPGE framework. In tasks with inherently multi-modal optimal policies, such as those observed in robotics, we approximate entropy using a Gaussian Mixture Model (GMM). Prior works (e.g., [1]) have demonstrated the feasibility of this approach in the context of diffusion-based policies.

More generally, since entropy merely quantifies stochasticity, it can be replaced by other metrics to capture utility, such as the distance between a given state and the empirical state distribution in the buffer, or other differentiable surrogate signals. The key requirement is that the utility function must be differentiable with respect to the state, so that it can guide generation. MPGE is fully compatible with a broad family of such utility functions, and we view entropy as one heuristic example rather than a universal choice.

2.Discussion of the hyperparameters

We fully understand your concern regarding hyperparameter sensitivity. We will explain the meaning and selection strategy for each of the following components: the utility function, the guidance scale $ω$, and the trade-off parameter $λ$. I will also describe how these parameters are chosen in practice.

1. First, regarding the choice of utility functions: as mentioned in the paper, there exists a wide variety of utility functions in reinforcement learning tasks. Our goal is to provide a general architecture that can accommodate and benefit from better-designed or more informative utility signals in the future. Depending on the specific task, different utility functions (e.g., Q-values or intrinsic curiosity as used in PGR) can be adopted to guide the generation of task-specific critical states. Our current choiceslike policy entropy and temporal-difference (TD) error are meant to serve as illustrative examples: entropy naturally appears in the policy optimization objective, and TD error is central to value estimation. These choices are heuristic rather than exhaustive, and we encourage practitioners to adapt and extend this framework using other utility formulations suitable for their own downstream tasks.

2. In our ablations, $ω=1$ performs best under standard training steps. However, with extended training (e.g., 3M steps), smaller scales (e.g., $ω<1$) also achieve strong performance, differing mainly in convergence speed. Larger scales tend to distort the generated data. For simple tasks, a small $ω$ is sufficient, while complex tasks may benefit from adaptive schemes such as annealing:

3. As for the value of $λ$, we have discussed its setting in Appendix A. A large $λ$ can lead to a significant distributional shift, which reduces the effectiveness of prioritized imitation mechanisms (PIM). Our ablation studies also confirm this observation. Therefore, actually $λ$ is not a freely tunable parameter—it must remain within a small range to be effective, similar in spirit to the clipping threshold in PPO.

3.Supplemental experiments

We fully understand your concern. To address this, we have re-run the experiments on the Gym tasks using five random seeds based on the original setup. The updated results are presented below:

The results show that our method consistently achieves state-of-the-art performance across all five seeds compared to the baseline algorithms. As further discussed in the supplementary experiments (Appendix C), we note that the baseline algorithm DSAC already achieves near-optimal performance on these Gym tasks, leaving limited room for improvement. The fact that both PGR and MPGE still manage to outperform DSAC highlights the importance of passive, utility-guided exploration, which contributes to further policy enhancement even in highly optimized settings.

4. Quality analysis of transitions generated by MPGE

MPGE generates imaginary transitions by first sampling critical states under utility guidance, then assembling transitions via a learned world model. These generated transitions exhibit two important properties like novelty and dynamical feasibility, which together contribute to improved policy performance.

4.1 Novelty of generated states

MPGE generates states that are reachable but under-visited, conditioned on utility signals (e.g., TD error or entropy) that reflect the current policy’s learning needs. This design introduces the following advantages:

(i) State-level generation allows MPGE to focus on critical regions beyond the static buffer coverage, rather than simply perturbing stored transitions as done in PGR.

(ii) Dynamic utility evolution enables continual novelty as the utility function shifts with policy updates, steering generation toward states where the policy is uncertain or under-trained.

(iii) Policy-centric guidance means MPGE is not constrained by the historical behavior policy; it rebalances the frequency of critical states under the current policy distribution, leading to better task relevance.

Together, these results in generated transitions that are more diverse and better aligned with the evolving policy, which supports exploration and evaluation in otherwise poorly covered areas of the environment.

4.2 Dynamical feasibility of generated states

MPGE ensures that all transitions respect environmental dynamics by simulating them through a pretrained world model, following standard practice in model-based RL. The world model is trained using diverse data (e.g., from random policy), covering broad regions of the state-action space and avoiding overfitting to narrow behavioral trajectories. In contrast, methods like PGR directly learn joint distributions over $(s,a,r,s ′)$ from buffer tuples. This paradigm does not enforce the structural constraints of real dynamics (i.e., the mapping from $(s,a)→(r,s′))$. As a result, PGR may capture spurious statistical correlations rather than true functional dependencies, potentially modeling unstructured factorizations. Such dynamics-inconsistent samples may degrade the value estimates and hinder policy improvement.

4.3 Summary

In summary, MPGE generates imaginary transitions that are:

Novel: they target critical, under-visited states tailored to the current policy;

Valid: they are assembled via a learned model trained to respect environment dynamics.

These properties enable MPGE to produce high-quality synthetic data that complements the replay buffer and improves learning efficiency. Our additional experiments further confirm that transitions generated by MPGE are more dynamically consistent and task-relevant, contributing to its superior performance over PGR in several benchmark tasks. To find more quantitative analysis, please refer to the rebuttal to reviewer W4cH due to the limitation of space.

5. Compute costs in MPGE

Thank you for the suggestion. We have added the relevant comparison. The walltime results(s) are now included as follows:

The results show that MPGE incurs slightly higher walltime compared to baseline algorithms. Notably, to reduce computational overhead, we perform MPGE training only once every 10 environment interaction steps(please refer to codebase), which significantly controls the overall compute cost.

Moreover, unlike PGR, which requires retraining the generative model on the buffer until convergence before use, MPGE continuously updates its generator during training. This makes MPGE’s diffusion-based training more efficient and lightweight in practice.

Minor

We sincerely appreciate and acknowledge your valuable suggestion. We will carefully consider this point and incorporate the corresponding clarification or discussion in the revised version.

Conclusion

We sincerely appreciate your thoughtful and constructive feedback. We hope our responses have effectively addressed your concerns and respectfully ask that you consider a more favorable evaluation.

[1] Wang Y, Wang L, Jiang Y, et al. Diffusion actor-critic with entropy regulator[J]. Advances in Neural Information Processing Systems, 2024, 37: 54183-54204.

  def get_reward(self, physics):
    """Returns a reward to the agent."""
    standing = rewards.tolerance(physics.head_height(),
                                 bounds=(_STAND_HEIGHT, float('inf')),
                                 margin=_STAND_HEIGHT/4)
    upright = rewards.tolerance(physics.torso_upright(),
                                bounds=(0.9, float('inf')), sigmoid='linear',
                                margin=1.9, value_at_margin=0)
    stand_reward = standing * upright
    small_control = rewards.tolerance(physics.control(), margin=1,
                                      value_at_margin=0,
                                      sigmoid='quadratic').mean()
    small_control = (4 + small_control) / 5
    if self._move_speed == 0:
      horizontal_velocity = physics.center_of_mass_velocity()[[0, 1]]
      dont_move = rewards.tolerance(horizontal_velocity, margin=2).mean()
      return small_control * stand_reward * dont_move
    else:
      com_velocity = np.linalg.norm(physics.center_of_mass_velocity()[[0, 1]])
      move = rewards.tolerance(com_velocity,
                               bounds=(self._move_speed, float('inf')),
                               margin=self._move_speed, value_at_margin=0,
                               sigmoid='linear')
      move = (5*move + 1) / 6
      return small_control * stand_reward * move