【强化学习系列】Gym库使用——创建自己的强化学习环境3：矢量化环境+奖励函数设计

目录

一、概述：不同模式环境选择

二、矢量化环境1——每批次同一图片：标量离散动作空间

        1.修改reset——(1) 环境重置加载随机化

        2.修改reset——(2) 增随机选取检测框

        3.修改step——(1) 改动作与环境交互

        4.修改step——(2) 增奖励机制函数

        5.修改step——(3) 增截断：游戏结束标准

        6.代码整合与矢量化环境

一、概述：不同模式环境选择

        在之前文章中，首先简单构建了单一的 gym 环境，以帮助目标检测的强化学习训练。【强化学习系列】Gym库使用——创建自己的强化学习环境1：单一环境创建测试

        之后在此基础上继续研读官方环境例子，以规范化gym环境类中的输入输出格式并将环境通过官方支持的注册方式创建使用。

【强化学习系列】Gym库使用——创建自己的强化学习环境2：拆解官方标准模型源码/规范自定义类+打包自定义环境

        但是仅仅构建一个环境对于训练来说还不够，为了提升训练效率，本文记录使用官方文档内介绍的矢量化环境的代码实现。

        对于矢量化环境（即同时创建多个环境帮助模型训练）存在对于整体流程架构的设计和考虑，在目标检测任务中，至少存在两种可以去实验的选择方案：

        1.创建同一张图的多个相同环境，每个环境只选其中一个框为目标，做出一个动作（单值）。

        这么做的好处是环境创建简单，问题是模型一批次只看一张图，效率低且泛化性存疑。

        2.创建多种图片的不同环境，每个环境以所有框为目标，传入一串并行的动作（向量）。

        这么做的好处是和模型训练的批量对接，可以近似为一个简单的视觉分类训练任务。问题是环境创建复杂，动作空间需要对齐，存在冗余。

        下面分别对上述两种思路都进行矢量环境的搭建，以帮助后续训练测试对比。

        Gymnaium官方文档地址：https://gymnasium.farama.org

二、矢量化环境1——每批次同一图片：标量离散动作空间

        1.修改reset——(1) 环境重置加载随机化

        在之前的环境中，重置环境使用的是从外界加载好图片数据框数据，通过参数形式传入环境。现在，要实现只传给环境所需资源的地址，在环境中自动随机选择图片和对应的框信息，生成gym环境。

        首先在初始化环境__init__中，不能简单根据传入的图片地址加载图片信息了，而是生成一个文件下所有图片的地址名，方便后续在reset中加载图片和框信息。

        并且新增了一个是否是矢量化加载同一张图的参数，如果same_env填入了一个索引数，那么不管你加载多少个同步向量模型，它们都是对应于一张图的，这个是符合此处的矢量化环境1要求。如果不传入，则会在所有的图片数据集中随机选取。

        同时由于图片加载放到了 reset 中实现，因此在状态空间范围加载时，缺少了关键信息。解决方案是1.如果图片形状不统一，则一开始初始化无限大范围的宽松环境，再在每次重置环境时重新定义状态空间；2.如果形状统一，则在初始化传入图片形状shape加载状态空间。

        可视化测试代码如下。完整环境代码详见最后一部分：代码整合。

import gymnasium as gym
from gymnasium.envs.registration import register

# 资源路径
image_path = './jpg/'
box_path = './box/'

# 注册环境
register(
    id='detect_env-v2',
    entry_point='my_gym_env.vector_auto_detect_env:Detect_Env',
)

# 创建环境
envs = gym.make_vec('detect_env-v2', num_envs=10,
                    render_mode='rgb_array', image_path=image_path, box_path=box_path, same_env=None)

# 测试
state,info = envs.reset(seed=42)
print('env_name:', info['env_name'])
print('box_id:', info['box_id'])
print('reward0:',  info['init_reward'])

        2.修改reset——(2) 增随机选取检测框

        固定好每次采样同一个图片环境后same_env = id，考虑框的选择。在第一种模式的环境中，每次随机选择图片中一个检测框作为 agent 操作目标对象。因此在原有环境的基础上（初始环境代码参考上一篇博文），需要在重置环境中加入随机性——seed。

        还是按照之前先注册在gym.make创建环境的方式进行。

        gymnasium有自带的矢量环境创建函数 gym.make_vec()，这个函数用于创建多个相同环境是否方便，只需设置参数num_envs=n，n是你要创建的个数。

# 注册环境
register(
    id='detect_env-v1',
    entry_point='my_gym_env.vector_same_detect_env:Detect_Env',
)

# 创建环境
envs = gym.make_vec('detect_env-v1', num_envs=3, render_mode='rgb_array', image_path=image_path)

        seed作为随机数种子输入是为了实现实验的可重复性，保证随机性可控，如下图结果展示。可见相同随机种子的采样结果是完全相同的。

# 测试
state,info = envs.reset(seed=42, options={'real_box':real_list, 'box':box_list})
print('batch1:', info['box_id'])

state,info = envs.reset(seed=42, options={'real_box':real_list, 'box':box_list})
print('batch2:', info['box_id'])

state,info = envs.reset(seed=26, options={'real_box':real_list, 'box':box_list})
print('batch3:', info['box_id'])

         3.修改step——(1) 改动作与环境交互

        因为现在是每个时间步step接受一个动作，对当前状态中选择的那一个框做操作，因此可以很简单的修改原始代码。

        4.修改step——(2) 增奖励机制函数

        接下来是整个强化学习环境最重要的设计——奖励机制函数。在之前，对于环境返回的奖励默认为0，现在要进行真实环境的奖励返回的设计与讨论。

        奖励函数的设计需要“奖罚分明”——这样可以明确智能体的学习目标：减少坏行为，增加好行为。还要避免智能体“投机取巧”——避免智能体因为一些行为奖励陷入边缘停滞，如在生存游戏中将活着设置奖励可能会导致agent取巧一直呆在原地不去探索未知危险环境。也要兼顾长短期奖励避免“鼠目寸光”——如果短期出现过高奖励，智能体只关注短期的目标没有长期战略性。

        ①奖罚分明设计

        对于目标检测的强化学习环境，希望检测框尽可能包含实际真实框。一种朴素简单的算法就是，如果真实小框完全包含进大的检测框则给予1的奖励（可以理解为大检测框包含了一个真实框：数量为单位），如果部分包含则按照包含的部分占真实小框的比例给奖励。如果一个真实框都没有包含则给予一个较大的 -10 惩罚。

         建立初始的奖励机制函数代码如下。

        ②避免投机取巧设计

        如果仅以包含的真实框数作为奖励会存在bug，如果agent不停的向外扩展，那么其奖励值始终处于高位，这样智能体会“偷懒”的只选择扩展操作而不进行缩减操作，任务目标是精细化调整框，这样的结果显然不符合要求。

         为了避免这种情况，鼓励缩减操作，需要对检测框的空白区域进行计算，并作为奖励值的一部分考虑。因为此处动作空间只包含框的上下移动，因此将真实框最高最低点到检测框的距离视为空白区域。

        可以看到加入空白区域的减分后，会减小多余扩展的奖励，抑制错误扩展行为。并且在应该扩展时，空白得分为0，不会影响正常扩展操作。

         ③ 总结代码

        本例选择了最简单的动作操作，暂不考虑长短期策略行为对奖励函数的影响。但需要注意的是平衡包含得分（iou）和空白减分（blank）参与奖励的权值。特别注意上面没有讨论到的特殊情况——假如上方空白区域非常大，下方漏真实小框（此时需要上缩，下扩），但是两者的加减可能会抵消掉奖励绝对值，此时模型是否会表现不好，还需训练实验讨论。

        下面根据给出完整的step改进和reward奖励设计代码。

         5.修改step——(3) 增截断：游戏结束标准

        对于一个框的上下调整，如果动作数量足够，其一定存在一个终止点——框退化为线，或框顶图片上下边界，此时在做操作已经没有意义了。因此需要给环境step设置一个截断点，当出现此状况时，游戏结束，进入下一个epoch。

         只需在step代码中加入简单的判断即可，假设认为图片边缘肯定不会存在目标物体，那么设置上下都触边即退出。假设检测框的最小高度，则小于此高度结束。

        下图展示的效果和截断步数都是完全随机的，加入截断逻辑后，游戏就可以自动停止了。

        6.代码整合与矢量化环境

        上述对env环境代码的修改为了方便展示是在单一环境上测试的，下面先对其单环境测试代码进行整理，并推广到矢量化环境上。

        文件层级结构：

        Env环境类主程序代码：

import pygame
import gymnasium as gym
from gymnasium import spaces
import numpy as np
from typing import Optional
import os
import json
from my_gym_env.utils import create_database, get_reward

class Detect_Env(gym.Env):
    metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 1}
    def __init__(self, render_mode: Optional[str] = None, image_path: Optional[str] = None, box_path: Optional[str] = None,
                 img_shape: Optional[tuple] = None, same_env: Optional[int] = None):
        super(Detect_Env, self).__init__()
        # 资源信息
        self.img_path = image_path
        self.box_path = box_path
        # 所有文件列表
        self.filename = self._load_img_paths()

        # 如果是同一环境，则选出这个特定的环境索引
        if same_env != None:
            self.id_choice = same_env
        else:
            self.id_choice = None

        # 过程信息
        self.real = None  # 真实框数据
        self.real_datasbase, self.real_id = None, None
        self.box_num = None  # 初始预测框数量
        self.box_id = None  # 当前环境选择移动的是第几个框
        # 设置最大框数
        self.box_max_num = 3

        # 创建动作范围
        self.action_space = spaces.Discrete(5)
        # 创建状态范围，如果图片尺度统一则初始化init生成，如果不统一则reset处生成
        if img_shape is not None:
            self.width, self.height = img_shape
            self.observation_space = spaces.Box(low=np.zeros((self.box_max_num,4),dtype=np.float32),
                                                high=np.tile([self.width, self.height,self.width, self.height],(self.box_max_num,1)),
                                                shape=(self.box_max_num,4), dtype=np.float32)
        else:
            self.width, self.height = None, None
            # 使用宽松的无限值定义——np.inf
            self.observation_space = spaces.Box(low=np.zeros((self.box_max_num,4),dtype=np.float32),
                                                high=np.tile([np.float32(np.inf), np.float32(np.inf),
                                                              np.float32(np.inf), np.float32(np.inf)],(self.box_max_num,1)),
                                                shape=(self.box_max_num,4), dtype=np.float32)

        # 定义可视化模式：人类可视化or机器人训练
        assert render_mode is None or render_mode in self.metadata["render_modes"]
        self.render_mode = render_mode
        self.window = None  # 可视化窗口
        self.clock = None   # 可视化时钟
        self.window_size = (600,600)  # 窗口大小
        self.background = None  # 背景图
        self.scale_x = None     # x横轴缩放比
        self.scale_y = None     # y竖轴缩放比

    def reset(
            self,
            *,
            seed: Optional[int] = None,
            options: Optional[dict] = None) :
        super().reset(seed=seed)
        # 重置时随机选择当前环境
        if self.id_choice == None:
            indices = np.arange(len(self.filename), dtype=np.int8)  # 生成索引数组
            cur_id_choice = self.np_random.choice(indices)
        else:
            cur_id_choice = self.id_choice
        basename = self.filename[cur_id_choice]
        # 加载资源
        imgdir = os.path.join(self.img_path, basename+'.jpg')
        img = pygame.image.load(imgdir)
        if self.width or self.height == None:
            self.width, self.height = img.get_width(), img.get_height()
            self.observation_space = spaces.Box(low=np.zeros((self.box_max_num, 4), dtype=np.float32),
                                                high=np.tile([self.width, self.height, self.width, self.height],
                                                             (self.box_max_num, 1)),
                                                shape=(self.box_max_num, 4), dtype=np.float32)
        boxdir = os.path.join(self.box_path, 'detect/'+basename+'.json')
        realdir = os.path.join(self.box_path, 'real/'+basename+'.json')
        state, self.real = self._load_box_jsons(boxdir, realdir)
        # 生成真实框rtree数据库
        self.real_datasbase, self.real_id = create_database(self.real)
        # 记录框数
        self.box_num = state.shape[0]
        # 选择当前图片第几个框作为目标
        self.box_id = self.np_random.integers(0,self.box_num)
        # 填充零使得状态达到指定维度
        self.state = np.pad(state,((0,self.box_max_num-state.shape[0]),(0,0)), mode='constant', constant_values=0)

        # 设置初始奖励
        # 观测目标框
        obv_box = self.state[self.box_id]
        reward = get_reward(obv_box, self.real_datasbase, self.real_id)

        if self.render_mode == 'human':
            self.background = pygame.transform.scale(img, self.window_size) # 设置背景图
            # 计算x和y方向的缩放比例
            self.scale_x = self.window_size[0]/ self.width
            self.scale_y = self.window_size[1] / self.height
            self.render()

        info = {'box_num':self.box_num, 'env_name':basename, 'box_id':self.box_id, 'init_reward':reward}
        return self.state, info

    def step(self, action):  # 动作序列{0：不操作，1：上扩， 2：上缩， 3：下扩， 4：下缩}
        # 根据action生成移动数组
        movement_np = np.zeros_like(self.state,dtype=np.float32)
        i = self.box_id
        if action == 1:
            movement_np[i,1] = 10
        elif action == 2:
            movement_np[i,1] = -10
        elif action == 3:
            movement_np[i,3] = 10
        elif action == 4:
            movement_np[i,3] = -10
        # 移动当前状态框
        self.state += movement_np
        self.state = np.clip(self.state, self.observation_space.low, self.observation_space.high)

        # 观测目标框
        obv_box = self.state[self.box_id]
        # 奖励获取
        reward = get_reward(obv_box, self.real_datasbase, self.real_id)
        # 截断游戏结束判断
        terminated = False
        if obv_box[3]-obv_box[1] < 1e-3:
            terminated = True
        if obv_box[1] ==0 and obv_box[3] == self.height:
            terminated = True

        return np.array(self.state, dtype=np.float32), reward, terminated, False, {}

    def render(self):
        # 初始化窗口和时钟
        if self.window is None and self.render_mode == 'human':
            pygame.init()
            pygame.display.init()
            self.window = pygame.display.set_mode(self.window_size)
        if self.clock is None and self.render_mode == 'human':
            self.clock = pygame.time.Clock()
        # 重新绘制背景，以清除上一层
        if self.background is not None:
            self.window.blit(self.background,(0,0))
            for real in list(self.real):
                rect_real = [real[0], real[1], abs(real[2]-real[0]), abs(real[3]-real[1])]
                # 根据x和y方向的缩放比例计算每个矩形框的新位置
                stretched_rectangle = [
                    rect_real[0] * self.scale_x,  # x 坐标
                    rect_real[1] * self.scale_y,  # y 坐标
                    rect_real[2] * self.scale_x,  # 宽度
                    rect_real[3] * self.scale_y  # 高度
                ]
                pygame.draw.rect(self.window, (255, 0, 0), stretched_rectangle, 2)  # 绘制矩形框，线宽为2
        # 绘制框
        for box in list(self.state):
            rect = [box[0], box[1], abs(box[2]-box[0]), abs(box[3]-box[1])]
            # 根据x和y方向的缩放比例计算每个矩形框的新位置
            stretched_rectangle = [
                rect[0] * self.scale_x,  # x 坐标
                rect[1] * self.scale_y,  # y 坐标
                rect[2] * self.scale_x,  # 宽度
                rect[3] * self.scale_y  # 高度
            ]
            pygame.draw.rect(self.window, (0, 255, 0), stretched_rectangle, 3)  # 绘制矩形框，线宽为3
        # 更新显示
        pygame.display.flip()
        self.clock.tick(self.metadata["render_fps"])

    def close(self):
        pygame.quit()

    def _load_img_paths(self):
        return np.array([fname.split('.')[0] for fname in os.listdir(self.img_path) if fname.endswith('jpg')])

    def _load_box_jsons(self, boxfile, realfile):
        with open(boxfile, 'r') as f:
            box_dict_list = json.load(f)
        box_list = []
        for box_dict in box_dict_list:
            box = box_dict['box']
            box_list.append(box)
        with open(realfile, 'r') as f:
            real_list = json.load(f)

        return np.array(box_list,dtype=np.float32).reshape(-1,4), np.array(real_list, dtype=np.float32).reshape(-1, 4)

        奖励函数独立py文件代码：

from rtree import index
import numpy as np

# 创建r数数据库
def create_database(box_list):
    tuple_id_box = []
    for i, box in enumerate(box_list):
        if box[0] > box[2]:
            box[0], box[2] = box[2], box[0]
        if box[1] > box[3]:
            box[1], box[3] = box[3], box[1]
        tuple_id_box.append([i, [box[0], box[1], box[2], box[3]]])
    # 实例化
    rtree = index.Index()
    # 插入框——元组形式
    for i_box in tuple_id_box:
        i, insert_box = i_box
        rtree.insert(i, (insert_box[0], insert_box[1], insert_box[2], insert_box[3]))

    return rtree, tuple_id_box

# r树查询返回框坐标
def rtree_retrieve_interbox(rtree, rtree_axis_id, search_area):
    id = list(rtree.intersection(search_area))
    near_list = [rtree_axis_id[i][1] for i in id]
    return near_list


def cal_batch_iou(boxes1, boxes2):
    # 计算交集坐标
    x1_inter = np.maximum(boxes1[:, 0], boxes2[:, 0])
    y1_inter = np.maximum(boxes1[:, 1], boxes2[:, 1])
    x2_inter = np.minimum(boxes1[:, 2], boxes2[:, 2])
    y2_inter = np.minimum(boxes1[:, 3], boxes2[:, 3])
    # 计算交集面积
    inter_area = np.maximum(0, x2_inter - x1_inter) * np.maximum(0, y2_inter - y1_inter)
    # 计算真实框的面积
    area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])
    # 计算IoU
    iou = inter_area / np.maximum(area1, 1e-6)  # 防止除以零
    return iou

# 设计detect_env的奖励函数
def get_reward(obv_box, real_box_database, real_id):
    obv_score = 0  # 当前框的得分
    # 根据包含真实框的iou均值来加分
    search_area = (min(obv_box[0],obv_box[2]), min(obv_box[1],obv_box[3]), max(obv_box[0],obv_box[2]),max(obv_box[1],obv_box[3]))
    overlap_list = rtree_retrieve_interbox(real_box_database, real_id, search_area)

    if overlap_list == []:
        obv_score -=10   # 没有包含任何真实框则减分
    else:
        real_np = np.array(overlap_list).reshape(-1, 4)
        obv_np = np.tile(list(obv_box), (real_np.shape[0], 4))
        # 计算包含得分
        iou = np.sum(cal_batch_iou(real_np, obv_np))
        obv_score += iou
        # 计算空白减分
        obv_hgt = np.maximum(obv_np[0][3]-obv_np[0][1],0)
        comb_hgt = np.max(real_np[:,3])-np.min(real_np[:,1])
        blank = np.maximum(obv_hgt-comb_hgt,0)/np.maximum(obv_hgt, 1e-6)*10
        obv_score -= blank
        print('blank:',blank,'   ','iou:',iou)
    return obv_score

        矢量化环境测试：

        在矢量化环境类中，如果其中一个环境截断，包装器会自动重置环境（在矢量化环境1中是重新选择框），因此这里设定，假如创建的n个环境，每个都经历过一次截断，则退出当前轮次。可以在下图看到，直到1341步时，所有环境都经历过一次截断（terminated=True）。   

import gymnasium as gym
from gymnasium.envs.registration import register
import numpy as np
from gymnasium.wrappers import TimeLimit

# 资源路径
image_path = './jpg/'
box_path = './box/'

# 注册环境
register(
    id='detect_env-v2',
    entry_point='my_gym_env.vector_auto_detect_env:Detect_Env',
)
# 创建环境个数
num_envs = 3
# 创建环境
envs = gym.make_vec('detect_env-v2', num_envs=num_envs,
                    render_mode='rgb_array', image_path=image_path, box_path=box_path, same_env=0)

# 测试
state,info = envs.reset(seed=42)
print('env_name:', info['env_name'])
print('box_id:', info['box_id'])
print('reward0:',  info['init_reward'])

break_point = np.array([0]*num_envs)  # 统计环境退出次数
epid = 0
while not np.all(break_point) > 0:
    action = np.random.choice([1,2,3,4],size=[num_envs,])
    state, reward, terminated, _, _ = envs.step(action)
    if np.any(terminated) != False:
        print('terminated'+str(epid+1), '  ',terminated)
    break_point[terminated] +=1
    #print('action'+str(epid+1)+':', action,'   ', 'reward'+str(epid+1)+':', np.round(reward,5))
    epid +=1

         至此矢量化第一种环境——“同图异框”情况的强化学习矢量环境创建完毕，在下一篇文章将先尝试在此环境基础上进行模型训练积累经验。之后再在所有的基础上，对第二种矢量环境——“异图向量框”进行建模训练。