如何在Keras-RL的DQNAgent中实现动作掩码功能？

问题背景 #

我正在做一个训练智能体在道路网络（图结构）里找最优路径的项目，已经基于OpenAI Gym搭了自定义环境，用Keras建模型、Keras-RL训练智能体。

核心问题：

• 网上的Keras DQN示例都是固定动作集，但我的场景里，可行动作数量随节点变化——起始节点可能只有2个可选邻接节点，后续节点可能有4个可行动作。
• 试过给不可行动作加负奖励，但效率不高；测试时用space.Discrete().sample(mask)（掩码np.array([1,1,0,0,0,0,0,0,0])，1代表可行、0代表不可行）过滤动作有效，但不知道怎么在训练流程里实现。
• 当前DQNAgent()设了10个可选动作，训练时智能体总会从这10个里选，导致可行邻接节点不足时抛出IndexError: list index out of range。

现有代码 #

模型构建 #

def build_model(env):
    model = Sequential()
    input_shape = (1, env.observation_space.shape[0]) # 输出为(1,8)
    model.add(Flatten(input_shape=input_shape))
    model.add(Dense(24, activation='relu'))
    model.add(Dense(24, activation='relu'))
    n_output_nodes = env.action_space.n
    model.add(Dense(n_output_nodes, activation='linear'))
    return model

智能体构建与训练 #

def build_agent(model, actions):
    policy = BoltzmannQPolicy()
    memory = SequentialMemory(limit=50000, window_length=1)
    dqn = DQNAgent(
        model=model, 
        memory=memory, 
        policy=policy,
        nb_actions=actions, 
        nb_steps_warmup=10, 
        target_model_update=1e-2,
    )
    return dqn

# 训练代码
model = build_model(env)
dqn = build_agent(model, env.action_space.n)
dqn.compile(Adam(learning_rate=1e-3), metrics=['mae'])
dqn.fit(env, nb_steps=50000, visualize=False, verbose=1)

解决步骤 #

1. 自定义带掩码支持的策略 #

Keras-RL默认策略不支持动作掩码，需要继承原有策略重写动作选择逻辑，只保留可行动作的Q值：

from rl.policy import BoltzmannQPolicy
import numpy as np

class MaskedBoltzmannQPolicy(BoltzmannQPolicy):
    def select_action(self, q_values, mask=None):
        if mask is not None:
            # 把不可行动作的Q值设为极小值，确保不会被选中
            q_values = q_values.copy()
            q_values[mask == 0] = -np.inf
        return super().select_action(q_values)

2. 修改智能体的动作选择逻辑 #

重写DQNAgent的动作获取方法，传入当前状态对应的掩码：

def build_agent(model, actions):
    policy = MaskedBoltzmannQPolicy()
    memory = SequentialMemory(limit=50000, window_length=1)
    dqn = DQNAgent(
        model=model, 
        memory=memory, 
        policy=policy,
        nb_actions=actions, 
        nb_steps_warmup=10, 
        target_model_update=1e-2,
    )
    # 替换原有get_action方法，加入掩码逻辑
    def masked_get_action(state):
        # 从自定义环境获取当前节点的动作掩码
        mask = env.get_current_mask()
        q_values = dqn.compute_q_values(state)
        return dqn.policy.select_action(q_values, mask)
    dqn.get_action = masked_get_action
    return dqn

3. 自定义环境添加掩码获取方法 #

在你的OpenAI Gym自定义环境中，实现get_current_mask方法，根据当前节点的邻接节点生成掩码：

import gym
from gym import spaces

class RoadNetworkEnv(gym.Env):
    # 原有环境初始化、step、reset等方法...
    
    def get_current_mask(self):
        # 假设self.current_node是当前所在节点，self.adj_list是道路网络的邻接表
        mask = np.zeros(self.action_space.n, dtype=int)
        # 把当前节点的邻接节点对应的动作索引设为1（需根据你的动作-节点映射逻辑调整）
        for action_idx in self.adj_list[self.current_node]:
            mask[action_idx] = 1
        return mask

4. 保险措施：环境step方法容错 #

在环境的step方法中，增加不可行动作的判断，避免意外报错：

def step(self, action):
    # 先判断动作是否可行
    mask = self.get_current_mask()
    if mask[action] == 0:
        # 选中不可行动作，返回惩罚，保持当前状态
        return self._get_observation(), -10, False, {}
    # 原有正常动作处理逻辑...

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