From 509b3d635cc301c4f373f1e932b8a7a830b893da Mon Sep 17 00:00:00 2001 From: 8x54zj-m <8x54zj-m@motrixlab.local> Date: Wed, 22 Jul 2026 19:58:19 +0800 Subject: [PATCH] Cap DreamWaQ adaptive learning rate --- motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py | 5 ++++- motrix_rl/tests/test_rslrl_np_env_reward_flow.py | 5 +++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py b/motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py index b41c663..570ab65 100644 --- a/motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py +++ b/motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py @@ -23,6 +23,7 @@ class DreamWaQPPO(PPO): # VAE 损失权重（与上游 beta=1.0 一致） vae_beta: float = 1.0 + max_learning_rate: float = 1e-3 @staticmethod def construct_algorithm(obs: TensorDict, env, cfg: dict, device: str) -> "DreamWaQPPO": @@ -133,7 +134,9 @@ class DreamWaQPPO(PPO): if kl_mean > self.desired_kl * 2.0: self.learning_rate = max(1e-5, self.learning_rate / 1.5) elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: - self.learning_rate = min(1e-2, self.learning_rate * 1.5) + self.learning_rate = min( + self.max_learning_rate, self.learning_rate * 1.5 + ) for param_group in self.optimizer.param_groups: param_group['lr'] = self.learning_rate diff --git a/motrix_rl/tests/test_rslrl_np_env_reward_flow.py b/motrix_rl/tests/test_rslrl_np_env_reward_flow.py index 88a07cb..a652cd5 100644 --- a/motrix_rl/tests/test_rslrl_np_env_reward_flow.py +++ b/motrix_rl/tests/test_rslrl_np_env_reward_flow.py @@ -3,9 +3,14 @@ from types import SimpleNamespace import numpy as np import torch +from motrix_rl.rslrl.torch.train.dreamwaq_ppo import DreamWaQPPO from motrix_rl.rslrl.torch.wrap_vec_env import RslrlNpEnvWrap +def test_dreamwaq_adaptive_learning_rate_is_capped_at_initial_rate(): + assert DreamWaQPPO.max_learning_rate == 1e-3 + + def test_step_preserves_rewards_and_marks_only_truncations_as_timeouts(): state = SimpleNamespace( obs=np.zeros((2, 3), dtype=np.float32),