使用sb3-contrib的RecurrentPPO策略梯度损失趋近于0的问题排查

我基于sb3-contrib训练RecurrentPPO（RPPO）智能体用于温控系统控制，已经实现了自定义Gym环境（代码如下）。训练时发现策略梯度损失始终趋近于0，完全没参与训练，但价值损失持续下降，训练监控数据如下，请问这一异常的原因是什么？

自定义Gym环境代码 #

class TempControlSeqEnv(gym.Env):
    def __init__(self, curriculum_phase, time_steps=5):
        super(TempControlSeqEnv, self).__init__()
        self.time_steps = time_steps
        self.num_sensors = 6
        self.stage = 0
        self.stage_3_init = []
        self.stage_1_counter = 0
        self.stage_2_counter = 0
        self.stage_3_counter = 0
        self.momentum = np.array([0.0, 0.0], dtype=np.float64)
        self.curriculum_phase = curriculum_phase

        self.X_data = []
        self.X_actions = []

        #step
        self.aux_action = [0.0,0.0,0.0,0.0,0.0,0.0]
        
        #Observation is a sequence: shape (time_steps, 6 TCs + stage)
        self.observation_space = self.observation_space = spaces.Dict({
            "obs":spaces.Box(
                low=1200, 
                high=1300, 
                shape=(time_steps, self.num_sensors + 1), 
                dtype=np.float32
            )
        })
        #Actions are +0, +1, +2, +3 for Z2 and Z5Z6
        self.action_space = spaces.MultiDiscrete([4, 4])

        self.max_steps = 24
        self.current_step = 0
        self.history = deque(maxlen=self.time_steps)
        self.state = None

    def reset(self):
        self.current_step = 0
        self.stage = 0
        self.stage_3_init = []
        self.stage_1_counter = 0
        self.stage_2_counter = 0
        self.stage_3_counter = 0
        self.history.clear()

        #Initialize history with random readings
        reading = self._generate_temp_reading()
        self.X_data = [reading.tolist()]
        self.X_actions = []
        for _ in range(self.time_steps):
            self.history.append(np.concatenate([reading, [self.stage]], axis=0))

        self.state = np.array(self.history)
        return {"obs": self.state}

    def get_stage(self):
        return self.stage
    
    def step(self, action):
        #force no action during phase change
        if self.stage == 1:
            action = np.array([0, 0], dtype=np.int32)

        obs, reward, done, info = self._true_step(action)
        return obs, reward, done, info
    
    def _true_step(self, action):
        self.current_step += 1
        self.momentum *= 0.4

        #State_finder
        action = np.asarray(action)
        self.aux_action[:] = [0.0, float(action[0]), 0.0, 0.0, float(action[1]), float(action[1])]
        if len(self.X_actions) > 0:
            self.X_actions.append(self.aux_action)
        else:
            self.X_actions = [self.aux_action]

        timestep_list = list(range(1, self.current_step+1))
        nested_timestep_list = [[x] for x in timestep_list]
        pred = PredNextState(self.X_data, action, nested_timestep_list, self.momentum)
        new_x = np.add(self.X_data[-1], np.round(pred[-1], 1))
        new_x = np.clip(new_x, 1200, 1300)
        print("new x = ", new_x)
        logging.info(f"Predicted new x(next stage) = {new_x}")
        print("action = ", self.aux_action)
        logging.info(f"Current action = {self.aux_action}")
        self.X_data.append(new_x)
        self.history.append(np.concatenate([new_x, [self.stage]], axis=0))
        del timestep_list, nested_timestep_list, pred
        self.state = np.array(self.history)

        print("stage = ", self.stage)
        logging.info(f"Current stage = {self.stage}")
        print("momentum = ", self.momentum)
        logging.info(f"Current momentum = {self.momentum}")
        #Placeholder for reward function
        #Stage 0, pre phase change
        if self.curriculum_phase == "pretrain":
            if self.stage == 0:
                if self.stage_1_counter > 6: #assume phase change 6 hours in, configurable
                    self.stage = 1
                else:
                    self.stage_1_counter += 1
                #Reward Z2
                init_Z2 = np.round(self.X_data[0][1])
                target_Z2_stage_1 = 1228
                target_Z2 = init_Z2 + self.current_step
                if target_Z2 > target_Z2_stage_1:
                    target_Z2 = target_Z2_stage_1
                reward_Z2 = ComputeZoneReward(new_x[1] + self.momentum[0], target_Z2, action[0])
        
                #Reward Z6
                init_Z6 = np.round(self.X_data[0][5])
                target_Z6_stage_1 = 1246
                target_Z6 = init_Z6 + self.current_step
                if target_Z6 > target_Z6_stage_1:
                    target_Z6 = target_Z6_stage_1
                reward_Z6 = ComputeZoneReward(new_x[5] + self.momentum[1], target_Z6, action[1])
    
                reward = reward_Z2 + reward_Z6
    
                print("Z2 reward: ", reward_Z2)
                print("Z6 reward: ", reward_Z6)
                logging.info(f"Z2 reward = {reward_Z2}")
                logging.info(f"Z6 reward = {reward_Z6}")
 
            #Stage 1, during phase change
            elif self.stage == 1:
                if self.stage_2_counter > 2: #assumes exit phase change after 2 hours and temperature change < 0.5, configurable
                    if (self.X_data[-1][1] - self.X_data[-2][1]) < 0.5:
                        self.stage = 2
                        self.stage_3_init = self.X_data[-1]
                else:
                    self.stage_2_counter += 1
                #No change
                reward = -np.abs(action[0] + action[1])
    
                print("stage 1 no reward")
                logging.info("stage 1 no reward")

            elif self.stage == 2:
                reward = 0.0
                print("stage 2 during pretrain no reward")
                logging.info("stage 2 during pretrain no reward")

        elif self.curriculum_phase == "stage2_only": 
            if self.stage != 2:
                reward = 0.0
                print("stage 0 or 1 after pretrain no reward")
                logging.info("stage 0 or 1 after pretrain no reward")

                #Stage 3, post phase change
            else:
                self.stage_3_counter += 1
                
                #Reward Z2
                init_Z2 = np.round(self.stage_3_init[1])
                target_Z2_stage_3 = 1239
                target_Z2 = init_Z2 + self.stage_3_counter
                if target_Z2 > target_Z2_stage_3:
                    target_Z2 = target_Z2_stage_3
                reward_Z2 = ComputeZoneReward(new_x[1] + self.momentum[0], target_Z2, action[0])
    
                #Reward Z6
                init_Z6 = np.round(self.stage_3_init[5])
                target_Z6_stage_3 = 1253
                target_Z6 = init_Z6 + self.stage_3_counter
                if target_Z6 > target_Z6_stage_3:
                    target_Z6 = target_Z6_stage_3
                reward_Z6 = ComputeZoneReward(new_x[5] + self.momentum[1], target_Z6, action[1])
    
                reward = reward_Z2 + reward_Z6
    
                print("Z2 reward: ", reward_Z2)
                print("Z6 reward: ", reward_Z6)
                logging.info(f"Z2 reward = {reward_Z2}")
                logging.info(f"Z6 reward = {reward_Z6}")
        

        print("reward = ", reward)
        logging.info(f"Total reward = {reward}")
        print(psutil.Process(os.getpid()).memory_info().rss / 1024 ** 2, "MB")
        del new_x
        gc.collect()
        done = self.current_step >= self.max_steps
        info = {}
        self.momentum += np.round(np.array(action)/3, 2)

        return {"obs": self.state.astype(np.float32)}, reward, done, info

    #Placeholder for init temperature
    def _generate_temp_reading(self):
        init_temp = [[1216.3, 1218, 1234.6, 1241.8, 1239, 1245.4], 
                     [1222.6, 1224.8, 1237.2, 1242.4, 1239.9, 1241.4], 
                     [1216.2, 1224, 1240.2, 1242.5, 1239.3, 1243.7],
                     [1215.9, 1219.2, 1241.2, 1241.7, 1241, 1241.9],
                     [1219.5, 1222.4, 1236.5, 1241.1, 1240.4, 1242.8],
                     [1217.9, 1221.1, 1240, 1243.1, 1237.8, 1240.7],
                     [1222.6, 1224.8, 1237.2, 1242.4, 1239.9, 1241.4],
                     [1224.5, 1226.5, 1241.3, 1243.7, 1240.5, 1243.2],
                     [1221.5, 1223.5, 1242.4, 1244.4, 1239.6, 1241.1],
                     [1222.6, 1225.3, 1240.7, 1243.9, 1241.4, 1244.4],
                     [1216.2, 1224, 1240.2, 1242.5, 1239.3, 1243.7],
                     [1219.1, 1226.2, 1239.1, 1241.8, 1240.1, 1244.3],
                     [1218, 1224.7, 1239.4, 1242.6, 1239.3, 1243.5],
                     [1220.8, 1225.7, 1240.6, 1242.7, 1239, 1241.8],
                     [1221.4, 1226.5, 1237.7, 1241.7, 1240.5, 1244.3]]
        return np.array(random.choice(init_temp)).astype(np.float32)

    def set_curriculum_phase(self, phase: str):
        assert phase in ["pretrain", "stage2_only"]
        self.curriculum_phase = phase
        if phase == "pretrain":
            self.max_steps = 10
        else:
            self.max_steps = 24
    
    def seed(self, seed=None):
        self.np_random, seed = gym.utils.seeding.np_random(seed)
        return [seed]
    
    def render(self, mode='human'):
        print(f"Step {self.current_step}: {self.history[-1]}")


#Create vectorized environment
env = TempControlSeqEnv(curriculum_phase="pretrain")
vec_env = DummyVecEnv([lambda: env])
vec_env = VecNormalize(vec_env, norm_obs=True, norm_reward=True)
vec_env.envs[0].set_curriculum_phase("pretrain")

#Instantiate Recurrent PPO agents with LSTM policy
model_A = RecurrentPPO("MultiInputLstmPolicy", vec_env, verbose=1, tensorboard_log="./ppo_rnn_tb/")
for _ in range(10):
    model_A.learn(2000)
    model_A.save("TC_PPO_v0.10_pretrain.zip")

vec_env.envs[0].set_curriculum_phase("stage2_only")

model_B = RecurrentPPO.load("TC_PPO_v0.10_pretrain.zip", vec_env, verbose=2)
for _ in range(10):
    model_B.learn(2000)
    model_B.save("TC_PPO_v0.10_stage2.zip")

训练监控数据 #

------------------------------------------
| time/                   |              |
|    fps                  | 0            |
|    iterations           | 2            |
|    time_elapsed         | 562          |
|    total_timesteps      | 256          |
| train/                  |              |
|    approx_kl            | 0.0015819944 |
|    clip_fraction        | 0            |
|    clip_range           | 0.2          |
|    entropy_loss         | -2.77        |
|    explained_variance   | 6.14e-06     |
|    learning_rate        | 0.0003       |
|    loss                 | 470          |
|    n_updates            | 10           |
|    policy_gradient_loss | -0.00765     |
|    value_loss           | 965          |
------------------------------------------

异常原因分析 #

• 奖励信号无效或缺失
  多个阶段的奖励设计导致策略无法感知动作价值：pretrain阶段进入stage2后奖励恒为0；stage2_only阶段非stage2状态奖励为0；stage1中动作被强制设为[0,0]，此时奖励-np.abs(action[0]+action[1])也恒为0。大量无意义的0奖励让策略无法建立动作与奖励的关联，策略梯度自然趋近于0。另外，ComputeZoneReward函数未给出实现，如果其返回值过小或范围被压缩，也会导致梯度信号微弱。

• 动作被强制覆盖，策略失去控制权
  在step方法中，当stage为1时，智能体输出的动作会被强制替换为[0,0]，此时策略的输出完全无法影响环境和奖励，策略梯度更新失去依据，损失自然趋近于0。

• 状态与动作的因果关系薄弱
  状态转移依赖未实现的PredNextState函数，如果该函数中状态更新与动作的关联极弱，或者状态变化随机性过强，智能体无法学习到动作对环境的影响，策略梯度会因信号不足而趋近于0。

• 奖励归一化压缩有效信号
  使用VecNormalize对奖励做归一化，而训练中存在大量0值奖励，归一化后有效奖励信号会被进一步压缩，导致策略梯度更新的驱动力不足。

• 训练阶段与步长不匹配
  pretrain阶段max_steps=10，但stage1在stage_1_counter>6时触发，意味着每个episode中大部分时间处于stage1或stage2，这两个阶段要么动作被强制，要么奖励为0，有效训练样本极少，策略无法获得足够的梯度信号来更新。

内容的提问来源于stack exchange，提问作者Anthony