1. 为什么我要折腾自定义gym环境可视化可视化是强化学习调试的刚需先交代一下背景。我做强化学习实验也有几年了从最早的表格型Q-learning到后来的PPO、SAC环境大部分时候用的都是现成的。但一到自己设计任务——比如让智能体学会控制一个机械臂末端去抓取目标点或者在一个迷宫里自主探索——就一定会碰到一个硬需求我得“看见”智能体在做什么。标题里的“自定义gym环境可视化绘制”其实就是这么一件事自己按照Gym接口写一个环境并且用一种可靠的方式把环境状态画出来。很多人写自定义环境第一步是搭状态转移逻辑也就是step()函数里怎么更新状态、怎么给奖励第二步才开始犯难这个环境怎么渲染出来有些人干脆就不渲染了只打印reward、step数这些标量数据最后训练出来的策略效果不好根本分不清是因为奖励设计不合理还是因为环境本身就写错了。我个人强烈建议只要你打算认真调一个自定义环境可视化就绝对不能省。哪怕只是一个简陋的矩形方块在画布上移动你都能一眼看出很多问题——比如智能体是不是穿墙了、目标点坐标是不是生成错了、碰撞检测是不是写反了。这种“眼见为实”的调试方式比单纯看数值高效得多。这篇博文就把我从“环境逻辑”到“可视化渲染”这一整套实践流程写清楚包括Gym渲染机制的底层逻辑、pygame绘制环境状态的核心技巧、以及各种我踩过的坑。适合刚接触自定义环境、或者写了环境但不知道怎么渲染的读者参考。2. 摸清Gym渲染机制的底细三条可视化路线怎么选2.1 Gym/gymnasium的render_mode是怎么运作的聊可视化之前必须先搞清楚Gym环境里渲染这件事的“协议”。社区里现在主流用的是gymnasiumOpenAI Gym已经停止维护了当然很多老项目还在用gym 0.2x。在新版API里render_mode是一个设计得很巧妙的参数它直接影响环境对象创建时的行为。我最早入坑的时候用的是老版gym那会儿渲染靠的是env.render()直接调用后来升级到gymnasium才发现接口变了。新版思路是在环境类的__init__里接收render_mode然后根据这个参数决定初始化哪些渲染相关的资源比如要不要创建pygame窗口、要不要把像素缓存保留下来。render()方法本身只负责“画一帧”至于画出来的东西是显示到屏幕上、还是存成数组、还是根本不要画面都由render_mode决定。对比一下三种常用的render_moderender_mode取值用途数据形式典型场景human弹出一个窗口实时显示直接显示到屏幕调试、演示rgb_array返回一帧画面的RGB数组(H, W, 3)的ndarray训练时记录视频、传给神经网络ansi返回字符串形式的画面文本终端环境、简易调试这三种模式里我最推荐的做法是环境同时支持human和rgb_array。因为训练的时候你不会一直盯着窗口看大概率是要把帧存下来做轨迹回放的但调试的时候你又确实需要一个窗口来看实时表现。两个都支持切换成本就只在一个参数上。2.2 为什么社区里大多数人选了pygameGym官方文档里其实给出了一条内置的渲染路径就是每个环境要实现render()并且可以使用底层的像素缓冲。但实际动手写的时候你会发现完全没有辅助工具的情况下靠numpy直接手动画图是极其痛苦的——你要自己处理抗锯齿、画圆、画矩形、写文字基本等于从零写一个图形库。所以社区里的自定义环境十个有九个用的是pygame。原因很简单pygame就是为2D游戏设计的Gym环境绝大多数也是2D离散或连续控制任务什么画板、画圆、碰撞检测、键盘监听pygame全都给你包好了。它的语法很直观pygame.draw.rect()画矩形pygame.draw.circle()画圆坐标系统跟屏幕像素一致理解成本极低。一条简单的经验法则如果你的环境是2D场景——小车、网格迷宫、机械臂侧视图、智能体移动——直接无脑选pygame如果是3D场景那就换mujoco或者gymnasium的RenderEngine接口去做深度渲染。别试图在pygame里勉强模拟3D那是给自己找罪受。2.3 渲染API设计思路把绘制与仿真解耦这一步很多人会忽略。我一开始写环境的时候是把绘制逻辑直接写在render()方法里的结果环境逻辑一变绘制代码也跟着乱掉。现在我的推荐做法是把环境拆成两层仿真逻辑层负责状态转移、奖励计算、动作更新渲染层只负责根据当前状态画出画面。这么做的好处是环境逻辑和画面显示互不干扰step()里改了状态render()里从上到下重新画一遍就行不会出现“状态变了但画面没更新”的狗血问题。可以很方便地支持多种可视化模式同一个状态用pygame画2D图也可以用matplotlib画曲线甚至可以输出成gif互不影响。方便做headless训练——服务器上没有显示器的时候不初始化pygame窗口只用rgb_array模式录像减少不必要的性能开销。这个设计思路贯穿整篇博文后面的代码也都按这个原则来写。3. 核心实现从零写一个可渲染的gymnasium自定义环境3.1 环境类骨架先搭仿真逻辑我以一个“二维平面上的点到点导航”任务为例来展开说明。这个任务足够简单同时又覆盖了自定义环境的大部分关键点状态定义、动作空间、奖励设计、单步推进、回合重置。智能体是一个小圆点目标是一个小方块智能体每步可以改变速度向量目标是到达目标点。直接上代码。这是一个最简化但五脏俱全的环境骨架import gymnasium as gym import numpy as np from gymnasium import spaces class PointNavEnv(gym.Env): def __init__(self, render_modergb_array, max_steps200): super().__init__() self.render_mode render_mode self.max_steps max_steps self.step_count 0 # 观察空间智能体x坐标、y坐标、目标x坐标、目标y坐标 self.observation_space spaces.Box( low0.0, high1.0, shape(4,), dtypenp.float32 ) # 动作空间x方向速度增量 y方向速度增量范围[-0.1, 0.1] self.action_space spaces.Box( lownp.array([-0.1, -0.1]), highnp.array([0.1, 0.1]), dtypenp.float32, ) self.agent_pos None self.target_pos None def reset(self, seedNone, optionsNone): super().reset(seedseed) self.agent_pos np.array([0.1, 0.1], dtypenp.float32) self.target_pos np.array([0.9, 0.9], dtypenp.float32) self.step_count 0 obs np.concatenate([self.agent_pos, self.target_pos]) return obs, {} def step(self, action): # 更新位置 self.agent_pos self.agent_pos np.clip(action, -0.1, 0.1) self.agent_pos np.clip(self.agent_pos, 0.0, 1.0) self.step_count 1 # 判断是否到达目标 dist np.linalg.norm(self.agent_pos - self.target_pos) terminated dist 0.05 truncated self.step_count self.max_steps # 奖励距离越近奖励越大到达终点给大额奖励 reward -dist if terminated: reward 10.0 obs np.concatenate([self.agent_pos, self.target_pos]) return obs, reward, terminated, truncated, {}有几个关键点需要展开说说。第一spaces.Box的边界要仔细设计。这个环境里观察空间是四个0到1之间的坐标值因为我把整个场景归一化到1x1的坐标空间里。这里有个常见的坑Box的low和high如果是标量表示所有维度同一个边界如果是数组就必须与维度数严格对应而且数组形状必须匹配。写错shape的话gymnasium的check_env工具会直接报错。第二动作空间要限制步长。我这里是速度增量每一次最多移动0.1对应一秒钟只能走十分之一的空间200步最多从头走到尾再折返一次。这个步长设计决定了任务难度——如果步长只有0.01那200步最多移动2看起来够用但实际训练起来会非常稀疏因为智能体很容易迷失方向。这种参数设定没有标准答案完全取决于你的任务但作为环境设计者你必须先自己“人脑模拟”一下这个值是否合理。第三奖励函数决定了学习难度。我用了-dist 10的稠密奖励设计让智能体每步都能获得“越近越好”的反馈。如果你的任务是稀疏奖励——到达终点才给1分——那训练会难很多。在调试阶段我建议先上稠密奖励等策略学好了再换稀疏奖励对比效果这样你能更快判断问题出在环境还是算法。3.2 render()方法实现按render_mode分流接下来是重头戏render()方法。按我前面说的解耦原则render()只做一件事根据当前状态画一帧画面然后根据render_mode决定怎么输出。def render(self): if self.render_mode rgb_array: return self._render_frame_rgb() elif self.render_mode human: self._render_frame_human() def _render_frame_rgb(self): # 返回一个HxWx3的RGB数组 import pygame # 用一个离屏surface画图避免弹窗 surface pygame.Surface((400, 400)) self._draw_scene(surface) # pygame surface转numpy数组并转成RGB顺序 frame pygame.surfarray.array3d(surface) frame np.transpose(frame, (1, 0, 2)) return frame def _render_frame_human(self): import pygame # 如果窗口还没创建先初始化 if not hasattr(self, screen): pygame.init() self.screen pygame.display.set_mode((400, 400)) pygame.display.set_caption(Point Nav) self._draw_scene(self.screen) pygame.event.pump() pygame.display.flip() def _draw_scene(self, surface): # 绘制背景、目标点、智能体 pass这里有个非常重要的细节rgb_array模式和human模式的初始化逻辑必须区分开来。如果你用rgb_array模式却在render()里调用了pygame.display.set_mode()在无显示器的服务器上运行会直接抛错。所以我只在human模式才创建窗口而rgb_array模式一律在离屏Surface上绘制最后转成数组。这个区分是很多人第一个踩的坑——在服务器上训练时一调用render()就报“video system not initialized”。根源就是没搞明白渲染模式之间的差异。3.3 坐标转换仿真坐标与像素坐标必须分开现在到可视化最核心的技术点坐标换算。仿真逻辑里的坐标是归一化的数值范围在0到1之间对应1x1的仿真空间。而pygame画布是400x400的像素空间如果直接把agent_pos当像素坐标画画面就全挤在左上角了几乎看不见。所以必须有一个“仿真坐标 → 像素坐标”的映射。def _sim_to_pixel(self, pos): # 将仿真坐标 [0,1] 映射到像素坐标 [0, 400] return (int(pos[0] * 400), int(pos[1] * 400))这个映射本质就是一次线性变换pixel sim * scale。如果你的仿真空间不是从0开始的——比如智能体坐标在[-1, 1]之间——那就要先平移再缩放pixel (sim 1) / 2 * scale。上面这个公式看着简单但很多人画图出错都是栽在这个计算上。我再给一个复杂一点的例子假设仿真空间是矩形x范围[-5, 5]y范围[0, 10]画布大小是800x600那映射函数应该这么写def _sim_to_pixel(self, sim_pos): px (sim_pos[0] - (-5)) / (5 - (-5)) * 800 py (sim_pos[1] - 0) / (10 - 0) * 600 return int(px), int(py)这个公式的通用形式是pixel (sim - sim_min) / (sim_max - sim_min) * canvas_size所有2D可视化都是这么干的。我建议务必单独抽一个函数来写不要每个绘制点手算一遍不然改画布尺寸的时候你会想骂人。4. 把状态画出来pygame绘制核心技巧4.1 基础图形绘制圆形区域、矩形目标点接下来补全_draw_scene()的具体实现。这个例子里要画三样东西背景、目标点矩形、智能体圆形。def _draw_scene(self, surface): # 背景白色 surface.fill((255, 255, 255)) # 目标点绿色矩形20x20像素 target_px self._sim_to_pixel(self.target_pos) pygame.draw.rect( surface, (0, 200, 0), (target_px[0] - 10, target_px[1] - 10, 20, 20), ) # 智能体蓝色圆形半径8像素 agent_px self._sim_to_pixel(self.agent_pos) pygame.draw.circle( surface, (0, 0, 200), agent_px, 8, ) # 可选画一条智能体到目标点的连线方便观察方向 pygame.draw.line( surface, (200, 200, 200), agent_px, target_px, 1, )几个细节值得单独拎出来强调一下。第一pygame.draw.rect()的坐标参数是“左上角坐标 宽高”不是“中心坐标 半径”。所以为了把矩形中心对齐到目标点上要自己减去半个尺寸做偏移。这个细节很不起眼但如果不做偏移目标点的视觉位置和实际仿真位置会差出10个像素调试的时候会误导你。第二pygame.draw.circle()接收的是中心坐标和半径这个和rect刚好相反用的时候别搞混。这两类函数参数顺序截然不同写多了就容易手滑。第三颜色值我用元组表示范围是0到255。背景选白色因为仿真逻辑里障碍物和边界都是黑色系的话白色底会让前景物体更醒目。调试可视化不是做艺术品重点是对比度要足够高——智能体用一个颜色目标用另一个颜色障碍物用第三种颜色一眼就能分清谁是谁。4.2 更复杂的绘制网格、障碍物、文字标签真实任务里环境往往远不止两个物体。比如迷宫环境需要画网格线机器人导航需要画障碍物而且调试时往往需要在画面上显示关键数值——当前步数、实时距离、动作值等等。网格线的画法非常简单就是一组横线和竖线def _draw_grid(self, surface, grid_size40): for x in range(0, 400, grid_size): pygame.draw.line(surface, (220, 220, 220), (x, 0), (x, 400)) for y in range(0, 400, grid_size): pygame.draw.line(surface, (220, 220, 220), (0, y), (400, y))障碍物的画法也就是一个rect只不过颜色深一些但有一个关键点障碍物的坐标也需要经过_sim_to_pixel()转换。如果你有一个障碍物列表每个障碍物是(x1, y1, x2, y2)的仿真坐标矩形那绘制循环就是def _draw_obstacles(self, surface, obstacles): for obs in obstacles: top_left self._sim_to_pixel((obs[0], obs[1])) bottom_right self._sim_to_pixel((obs[2], obs[3])) w bottom_right[0] - top_left[0] h bottom_right[1] - top_left[1] pygame.draw.rect(surface, (80, 80, 80), (top_left[0], top_left[1], w, h))文字标签是调试利器。在屏幕上显示当前距离值、步数等你就能实时观察智能体是否在逼近目标而不是反复看一圈圈连不上的轨迹。用pygame.font模块def _draw_text(self, surface, text, position): font pygame.font.SysFont(Arial, 20) img font.render(text, True, (0, 0, 0)) surface.blit(img, position)我在自己的环境里习惯把三个信息一直挂在画布角落当前步数、智能体坐标、目标坐标。调试的时候哪怕画面暂时看不出来问题数字也能告诉我状态到底发生了什么变化。4.3 rgb_array模式下的数据转换细节rgb_array模式里那个numpy转换曾经让我在调试时浪费过一整晚。pygame的surfarray.array3d()返回的数组轴顺序是(height, width, 3)吗其实是(width, height, 3)——也就是说第一个维度是x像素坐标第二个维度是y像素坐标。而gymnasium约定的rgb_array格式是(height, width, 3)——第一个维度是行对应y轴。所以必须转置一次frame np.transpose(frame, (1, 0, 2))不转置的话直接把这个数组交给gymnasium的RecordVideo包装器生成的视频画面会是竖着拉伸的。肉眼看不出来等你看一段训练录像目标点和智能体全是“扁的”或者“位置不对”就知道这个转置有多关键了。另外surfarray.array3d()返回的数组颜色顺序是BGR——pygame底层就是这么存的——不是标准的RGB。如果你只是自己看录像问题不大反正颜色通道整体交换一下画面颜色会异常偏蓝偏红。但如果你打算用OpenCV写视频就得显式转一下frame frame[:, :, ::-1] # BGR - RGB这两个坑是rgb_array模式最常见的出错点把这行注释写清楚一点下次回来看代码不至于一头雾水。4.4 在训练循环里使用自定义渲染环境写好了接下来就是在训练代码里正确调用。这里要特别注意render_mode的传参方式和调用时机。env PointNavEnv(render_modergb_array) # 训练用 # env PointNavEnv(render_modehuman) # 调试用 obs, _ env.reset() for step in range(200): action env.action_space.sample() obs, reward, terminated, truncated, _ env.step(action) frame env.render() # 收集frame到列表之后用imageio写视频 if terminated or truncated: break训练循环里一帧一帧调用render()拿回数组攒到列表里用imageio写成一个gif或mp4就能看到完整的回合轨迹了。我个人习惯是每隔一定步数比如1000步随机抽一个回合来录像这样能直观掌握训练的推进节奏。我用一个简单的gif写入参考需要先安装imageiopip install imageio imageio-ffmpegimport imageio frames [] # 训练循环里收集frame这里省略 imageio.mimsave(episode.gif, frames, fps30)5. 进阶思路仿真与渲染分离才是真正好维护的方案5.1 为什么要单独维护一个渲染器类前面说了解耦这里我把它落到实处。如果你只有一个环境类代码里render()、_draw_scene()、_sim_to_pixel()都堆在一起当环境逻辑变复杂的时候这个类的职责会变得非常混乱。更好的做法是把渲染器单独拆成一个类环境通过接口去调用它。我写渲染器类时一般长这样class PointNavRenderer: def __init__(self, render_mode, canvas_size400): self.render_mode render_mode self.canvas_size canvas_size self.screen None if self.render_mode human: pygame.init() self.screen pygame.display.set_mode((canvas_size, canvas_size)) pygame.display.set_caption(Point Nav) def render(self, agent_pos, target_pos): if self.render_mode rgb_array: return self._render_rgb(agent_pos, target_pos) elif self.render_mode human: self._render_human(agent_pos, target_pos) def _render_rgb(self, agent_pos, target_pos): surface pygame.Surface((self.canvas_size, self.canvas_size)) self._draw(surface, agent_pos, target_pos) frame pygame.surfarray.array3d(surface) return np.transpose(frame, (1, 0, 2)) def _render_human(self, agent_pos, target_pos): self._draw(self.screen, agent_pos, target_pos) pygame.event.pump() pygame.display.flip() def _draw(self, surface, agent_pos, target_pos): # 画图逻辑写在这里 pass然后环境类里只需要持有一个渲染器实例render()里调用它就行def render(self): return self.renderer.render(self.agent_pos, self.target_pos)这个拆分的好处等你往环境里加第二个任务、或者把同一套渲染逻辑用到另一个环境上时就会体会到。渲染器只关心“怎么画指定的状态”至于这个状态从哪里来、奖励怎么算它一概不关心。测试起来也方便你可以不创建一个完整环境直接构造一个渲染器手动设置几个坐标点看看画面是否符合预期。5.2 性能考量训练时别滥用human模式训练性能是另一个容易被忽视的方面。human模式需要在屏幕上刷新窗口这一步在很多系统上要消耗几十毫秒。如果你在训练循环里每步都调用human模式渲染训练时间会翻好几倍。特别是PPO这类需要大量采样样本的算法一个回合几百步每一步都刷窗口跑一晚上可能只训练了几个千步。我的实践经验是训练时用rgb_array模式并且不一定每一步都渲染。可以每10步才采集一帧画面这样录像虽然帧率低一点但能看出整体的行为模式够用了。调试算法时用human模式单独跑一个回合看实时表现。两边互不干扰。另外rgb_array模式如果也不需要录像那干脆不要调用render()——完全省掉绘制开销纯跑仿真逻辑。很多人在训练时习惯性地每步调用render其实cpu在干大量无用功。5.3 状态空间越来越大时分层渲染思路环境状态一大——比如要画地图、要画多个智能体、还要显示时序曲线——渲染逻辑就会越来越拗。我现在的做法是给渲染器增加“图层”概念背景层画地图和静态元素实体层画智能体、目标点等动态元素信息层画文字标签和曲线。每层分别用一个方法绘制render()里依次调用。def _draw(self, surface, state): self._draw_background(surface) self._draw_entities(surface, state) self._draw_overlay(surface, state)这样做的好处是你想临时加一个调试信息只用改overlay层根本不会动到实体层的代码。有了这个分层习惯就算环境复杂度翻倍渲染代码的维护压力也能控制得住。6. 踩坑实录自定义环境可视化最常见的6个问题6.1 pygame窗口弹不出来 / 报错这个问题高居榜首。常见原因有两个一是环境在rgb_array模式下还调用了pygame.display.set_mode()。前面说过这个函数需要真实显示器。解决办法就是在人类可交互的模式下才初始化窗口rgb_array一律用Surface离屏绘制。二是没有调用pygame.event.pump()。在human模式下如果不定时处理事件队列里的窗口事件macOS系统上窗口会直接“假死”——显示出来但不刷新。我习惯在每帧渲染末尾强制调用一次pygame.event.pump()保证窗口响应正常。6.2 画面一闪而过/白屏很多新手初始化pygame窗口后画完图形不调用pygame.display.flip()画面就根本不会刷新到屏幕上。因为pygame的绘制是先画到后备缓冲区只有flip之后才会交换到前面显示。这个跟OpenGL的双缓冲机制是一个道理。记住一条human模式每帧渲染后必须flip()。6.3 坐标错位物体位置不对这个问题排查起来最疼因为bug可以藏得很深。我的排查顺序是先看看仿真坐标范围是否正确——打印出agent_pos确认是不是在预期范围内。再检查_sim_to_pixel()映射是否与仿真坐标范围匹配——比如仿真坐标是[-1,1]但映射按[0,1]算画面位置就会整体偏到左上角或右下角。最后检查绘制参数——rect用的是左上角还是中心点圆用的是中心点还是左上角。这三个环节任何一个出错都会导致画面位置错位。我调试时习惯先在目标点位置画一个十字标记两条交叉线这样更容易判断中心是否对得上。6.4 rgb_array输出的画面方向和颜色不对这个话题前面已经讲过这里再总结一次。要记牢两个硬规则pygame.surfarray.array3d()返回的轴顺序是(width, height, 3)要转成gymnasium要求的(height, width, 3)必须np.transpose(frame, (1, 0, 2))。颜色通道是BGR不是RGB如果需要标准RGB用frame[:, :, ::-1]。6.5 渲染慢导致训练半天跑不完前面提过每步都render()很费时间。但还有一个隐蔽的性能问题rgb_array模式下每次创建新的Surface对象也有开销。如果你不需要连续录像只求偶尔抽帧别每次render()都新建Surface可以预先创建好一个Surface反复使用只在render()开头fill()清空即可这样能省下不少内存分配时间。6.6 在线训练时渲染会影响随机种子复现这个坑挺冷门。如果你在训练循环里每步都调用render()而pygame内部用了一些随机操作比如字体渲染时的抗锯齿那可能间接影响全局的随机数状态。虽然概率很低但对于强调可复现的实验建议在同一条件下做两组对照一组渲染一组不渲染看看训练曲线是否一致。如果不一致就把渲染从训练循环里彻底隔离出来用独立的回合去录像。7. 一套完整可复跑的示例机械臂二连杆任务可视化前面都是零散的技术点最后给一个稍微进阶一点的完整示例把本章所有技术串起来。我来设计一个简单但很典型的环境一个二连杆机械臂在二维平面内末端执行器需要触达目标点。这个任务的核心状态是两段臂的关节角度但可视化时你需要根据正运动学把角度换算成末端坐标。环境类骨架如下import gymnasium as gym import numpy as np from gymnasium import spaces import pygame class TwoLinkArmEnv(gym.Env): def __init__(self, render_modergb_array): super().__init__() self.render_mode render_mode self.canvas_size 400 self.arm_length 80 # 每段臂的长度像素 # 观察两个关节角度 目标点坐标 self.observation_space spaces.Box( lownp.array([-np.pi, -np.pi, 0.0, 0.0]), highnp.array([np.pi, np.pi, 1.0, 1.0]), dtypenp.float32, ) self.action_space spaces.Box( lownp.array([-0.1, -0.1]), highnp.array([0.1, 0.1]), dtypenp.float32, ) self.angles None self.target None self.renderer TwoLinkArmRenderer(render_mode, self.canvas_size) def reset(self, seedNone, optionsNone): super().reset(seedseed) self.angles np.array([0.0, 0.0], dtypenp.float32) self.target np.array([0.8, 0.8], dtypenp.float32) obs np.concatenate([self.angles, self.target], dtypenp.float32) return obs, {} def step(self, action): self.angles self.angles action self.angles np.clip(self.angles, -np.pi, np.pi) # 正运动学算末端坐标 end_effector self._forward_kinematics(self.angles) dist np.linalg.norm(end_effector - self.target) terminated dist 0.08 reward -dist if terminated: reward 20.0 obs np.concatenate([self.angles, self.target], dtypenp.float32) return obs, reward, terminated, False, {} def _forward_kinematics(self, angles): # 把关节角度映射到单位仿真空间 x 0.5 self.arm_length * np.cos(angles[0]) / self.canvas_size y 0.5 self.arm_length * np.sin(angles[0]) / self.canvas_size x self.arm_length * np.cos(angles[0] angles[1]) / self.canvas_size y self.arm_length * np.sin(angles[0] angles[1]) / self.canvas_size return np.array([x, y], dtypenp.float32) def render(self): return self.renderer.render(self.angles, self.target)渲染器类与前面介绍的类似但需要在画布上同时画出两段杆、关节圆点和目标点class TwoLinkArmRenderer: def __init__(self, render_mode, canvas_size400): self.render_mode render_mode self.canvas_size canvas_size self.screen None if render_mode human: pygame.init() self.screen pygame.display.set_mode((canvas_size, canvas_size)) pygame.display.set_caption(Two-link Arm) def render(self, angles, target): if self.render_mode rgb_array: return self._render_rgb(angles, target) else: self._render_human(angles, target) def _render_rgb(self, angles, target): surface pygame.Surface((self.canvas_size, self.canvas_size)) self._draw(surface, angles, target) frame pygame.surfarray.array3d(surface) return np.transpose(frame, (1, 0, 2)) def _render_human(self, angles, target): self._draw(self.screen, angles, target) pygame.event.pump() pygame.display.flip() def _draw(self, surface, angles, target): surface.fill((255, 255, 255)) # 目标点 target_px (int(target[0] * self.canvas_size), int(target[1] * self.canvas_size)) pygame.draw.rect(surface, (0, 200, 0), (target_px[0] - 5, target_px[1] - 5, 10, 10)) # 臂的起点固定在画布中心 origin (self.canvas_size // 2, self.canvas_size // 2) joint1 ( int(origin[0] 80 * np.cos(angles[0])), int(origin[1] 80 * np.sin(angles[0])), ) end_effector ( int(joint1[0] 80 * np.cos(angles[0] angles[1])), int(joint1[1] 80 * np.sin(angles[0] angles[1])), ) # 第一段杆 pygame.draw.line(surface, (0, 0, 200), origin, joint1, 6) # 第二段杆 pygame.draw.line(surface, (200, 0, 0), joint1, end_effector, 6) # 关节圆点 pygame.draw.circle(surface, (0, 0, 0), joint1, 5) pygame.draw.circle(surface, (0, 0, 0), end_effector, 5)这个例子比前面点导航复杂但整体结构完全一致环境只管仿真逻辑渲染器只管画图两者通过“状态传入”解耦。运行这个环境的训练循环你可以很直观地看到机械臂末端在训练过程中逐渐学会伸向目标点。8. 可视化调试的一些进阶扩展思路8.1 同时录制多个视角当你调试复杂任务时单看一个主视图往往不够。比如机械臂任务既想看主视角的完整场景又想看末端执行器与目标点的距离曲线。我的做法是让render()返回一个包含多个子画面的组合帧——主视图画左上角距离曲线画右下角用pygame.Surface.blit()拼接即可。这样自动生成的回放视频信息量比单视图丰富得多。8.2 用快捷键切换模式human模式下可以监听键盘事件来做交互调试。比如按空格键切换是否显示轨迹按数字键切换不同的观察视角。pygame的event.get()配合KEYDOWN事件就能实现这套逻辑适合做演示或者手动验证边界情况。监听键盘事件的代码最好放在render()外部由训练循环管理避免环境类承担过多交互职责。8.3 梯度可视化与特征热力如果你的状态包含高维数据——比如像素观察——可以顺便把特征图或者价值估计画出来。最简单的是用pygame.surfarray.make_surface()把一个二维数组转成画面叠加在主视图旁边。这在调试基于图像的强化学习任务时特别有用能直观看到模型关注了哪些区域。9. 收尾我这个环境可视化方案踩出来的经验沉淀如果你从头看到这里从搭仿真骨架、实现render、写pygame绘制到踩坑排查、拆分渲染器再到一个完整的二连杆例子整套自定义环境可视化闭环基本就建立起来了。我在实际项目中反复用这个套路最大的体会就是可视化不是做完环境之后的附属品而应该是写环境的第一天就规划进去的模块架构。先想清楚“这个环境要画什么、怎么展示给用户、哪些信息需要调试时可见”再来写仿真逻辑反而更不容易跑偏。另外一个很实用的建议是别追求一上来就把画面做得非常漂亮。先用最简陋的方块、圆形、连线把画面撑起来确认坐标正确、逻辑正确之后再慢慢加颜色、贴图、文字。调试可视化本身就是迭代过程一上来追求美感的代价是bug藏得更深。如果你只是想让训练过程看起来直观不需要录像或交互那就把render_mode固定成human写完环境直接跑一个随机策略的回合看看智能体是怎么动的——这一步能发现的环境逻辑bug比你看一百个reward数值都管用。我第一次跑通那个二连杆环境的时候就是因为看到末端一直在目标点附近绕来绕去却没有足够的奖励信号才回去改稀疏奖励设计的。一个能看见的环境就是强化学习实验里最值钱的那块“放大镜”。