Diffusion Policy Visuomotor Policy Learning via Action Diffusion官方项目解读（二）（2）

官方项目地址：https://diffusion-policy.cs.columbia.edu/
Colab代码：vision-based environment

---

系列子文章：

运行官方代码库中提供的Colab代码：vision-based environment（二）（1）
运行官方代码库中提供的Colab代码：vision-based environment（二）（2）
运行官方代码库中提供的Colab代码：vision-based environment（二）（3）
运行官方代码库中提供的Colab代码：vision-based environment（二）（4）
运行官方代码库中提供的Colab代码：vision-based environment（二）（5）
运行官方代码库中提供的Colab代码：vision-based environment（二）（6）
运行官方代码库中提供的Colab代码：vision-based environment（二）（7）

---

七、函数 pymunk_to_shapely

---

def pymunk_to_shapely(body, shapes):
    geoms = list()
    for shape in shapes:
        if isinstance(shape, pymunk.shapes.Poly):
            verts = [body.local_to_world(v) for v in shape.get_vertices()]
            verts += [verts[0]]
            geoms.append(sg.Polygon(verts))
        else:
            raise RuntimeError(f'Unsupported shape type {type(shape)}')
    geom = sg.MultiPolygon(geoms)
    return geom

---

def pymunk_to_shapely(body, shapes):

作用：定义一个函数 pymunk_to_shapely，它接收两个参数 body 和 shapes。

• 参数说明：
  • body：通常为一个 pymunk 中的物体（Body），它代表刚体在物理仿真中的状态。这个对象提供了将局部坐标转换为世界坐标的方法，例如 body.local_to_world(v)。
  • shapes：一个列表或其他可迭代对象，包含一个或多个 pymunk 的形状（例如多边形 Poly），这些形状通常与 body 关联。
• 目的：将 pymunk 的形状转换为 Shapely 的几何对象（Polygon 或 MultiPolygon），方便后续进行几何计算和空间操作。
• 输出：该函数最终返回一个 Shapely 的 MultiPolygon 对象。

    geoms = list()

作用：创建一个空列表 geoms，用于存储转换后的 Shapely 几何对象。

• 示例：
  • 执行后，geoms = []。
• 意义：列表用来逐个保存由每个 pymunk 形状转换得到的 Shapely 多边形，最终组合成一个 MultiPolygon。

    for shape in shapes:

遍历每个形状
作用：对传入的 shapes 中的每个形状进行循环处理。

• 示例：
  • 如果 shapes 包含两个形状，例如 shapes = [shape1, shape2]，那么循环会先处理 shape1，再处理 shape2。
• 意义：对每个形状进行单独转换，保证所有形状都能转为 Shapely 的几何对象。

        if isinstance(shape, pymunk.shapes.Poly):

类型判断
作用：判断当前的 shape 是否为 pymunk 的多边形类型（pymunk.shapes.Poly）。

• 示例：
  • 如果 shape 是一个多边形，例如其内部顶点可能为 [(0,0), (10,0), (10,10), (0,10)]，则条件成立。
• 意义：当前函数只支持多边形类型，如果传入的形状不是多边形，则会抛出异常（见后面的 else 部分）。

            verts = [body.local_to_world(v) for v in shape.get_vertices()]

转换顶点坐标
作用：对当前多边形 shape 的所有顶点进行坐标转换。

• 具体过程：
  • 调用 shape.get_vertices() 获取多边形在局部坐标系下的顶点列表。例如，shape.get_vertices() 返回 [(0,0), (10,0), (10,10), (0,10)]（假设单位为像素或其他单位）。
  • 对于每个顶点 v，调用 body.local_to_world(v) 将局部坐标转换为世界坐标。例如，假设 body 位于 (100, 200) 的位置，则：
    • 对 v = (0, 0) 可能返回 (100, 200)；
    • 对 v = (10, 0) 返回 (110, 200)；
    • 对 v = (10, 10) 返回 (110, 210)；
    • 对 v = (0, 10) 返回 (100, 210)。
  • 最终得到的 verts 列表为 [(100,200), (110,200), (110,210), (100,210)].
• 意义：转换后得到的顶点处于全局坐标系中，方便 Shapely 几何计算与其它空间数据的对接。

            verts += [verts[0]]

闭合多边形
作用：将多边形的第一个顶点添加到顶点列表末尾，确保多边形闭合。

• 示例：
  • 原来 verts = [(100,200), (110,200), (110,210), (100,210)]，执行后变为 [(100,200), (110,200), (110,210), (100,210), (100,200)]。
• 意义：Shapely 的 Polygon 需要闭合路径（首尾一致）才能正确表示一个封闭的多边形区域。

            geoms.append(sg.Polygon(verts))

创建 Shapely 多边形对象
作用：利用转换后的顶点列表 verts 创建一个 Shapely 的 Polygon 对象，并添加到 geoms 列表中。

• 具体示例：
  • 使用上一步的顶点列表，执行 sg.Polygon(verts) 得到一个多边形对象，该对象表示一个矩形，其边界为上述顶点。
• 意义：把 pymunk 的形状转换为 Shapely 格式后，可以利用 Shapely 提供的丰富几何运算和空间查询功能。

        else:

处理不支持的类型
作用：如果当前 shape 不是 pymunk.shapes.Poly 类型，则进入 else 分支。

• 意义：明确指出该函数目前只支持多边形，对于其他形状类型不做转换，保证代码的健壮性。

            raise RuntimeError(f'Unsupported shape type {type(shape)}')

抛出异常
作用：对不支持的形状类型，抛出 RuntimeError 异常，并输出提示信息，显示不支持的类型。

• 示例：
  • 如果 shape 是一个圆（Circle），那么 type(shape) 可能显示为 <class 'pymunk.shapes.Circle'>，错误信息为：
    "Unsupported shape type <class 'pymunk.shapes.Circle'>"
• 意义：明确错误原因，防止后续代码继续处理不支持的类型，保证数据正确性。

    geom = sg.MultiPolygon(geoms)

创建 Shapely MultiPolygon
作用：将之前收集的所有 Shapely 多边形对象组合成一个 MultiPolygon 对象。

• 示例：
  • 如果 geoms 列表中有两个多边形，例如一个矩形和另一个三角形，则 sg.MultiPolygon(geoms) 会创建一个包含这两个多边形的复合几何对象。
  • 意义：将多个几何体组合成一个整体，方便后续的空间运算、碰撞检测或区域分析。

    return geom

返回结果
作用：将最终生成的 MultiPolygon 对象作为函数的输出返回。

• 意义：调用者可以利用返回的 Shapely 几何对象进行进一步的空间计算或可视化操作。

八、类 PushTEnv，继承自gym.Env

该类继承自 Gym 环境（gym.Env），用于构建一个基于 PyMunk 物理仿真的具身任务环境，任务目标通常与推动“T”型积木相关。环境内部不仅负责物理模拟、控制更新、状态与奖励计算，还支持图形渲染（包括人机交互模式）以及状态重置、种子设置等功能。

---

---

# env
class PushTEnv(gym.Env):
    metadata = {"render.modes": ["human", "rgb_array"], "video.frames_per_second": 10}
    reward_range = (0., 1.)

• class PushTEnv(gym.Env):

  • 作用：声明一个名为 PushTEnv 的类，继承自 Gym 提供的 Env 基类，使其能被 Gym 框架识别为环境。
  • 意义：在具身智能研究中，环境是算法与物理实体交互的平台，继承 gym.Env 方便调用标准 API（如 reset、step、render 等）。

• metadata = {"render.modes": ["human", "rgb_array"], "video.frames_per_second": 10}

  • 作用：定义环境的元数据，指定支持的渲染模式（human：实时窗口；rgb_array：返回 RGB 数组）及视频帧率。
  • 示例：视频帧率为 10 帧/秒，适合观察低速运动的仿真效果。

• reward_range = (0., 1.)

  • 作用：规定奖励的取值范围为 0 到 1。
  • 意义：对强化学习算法来说，知道奖励范围有助于稳定性和调参。

---

八.1 def __init__()

    def __init__(self,
            legacy=False,
            block_cog=None, damping=None,
            render_action=True,
            render_size=96,
            reset_to_state=None
        ):

• 作用：构造函数，用于初始化环境实例。
• 参数说明：
  • legacy (bool)：是否采用旧版状态设置（例如顺序不同），默认 False。
  • block_cog：指定积木（block）的重心（Center Of Gravity），如传入 (x, y) 坐标；若为 None 则不改变。
  • damping：物理空间阻尼系数，若不为 None，则覆盖默认阻尼。
  • render_action (bool)：是否在渲染时显示动作信息（如鼠标点击标记）。
  • render_size (int)：渲染图像的大小（像素），例如 96 表示生成 96×96 的图像。
  • reset_to_state：重置时预设的状态，如果为 None，则随机生成状态。

        self._seed = None

• 作用：初始化环境种子变量为 None。
• 意义：后续调用 seed() 方法设置种子，用于随机性控制。

        self.seed()

• 作用：调用自身的 seed 方法，设置随机种子。
• 示例：若 seed() 随机生成种子 1234，则后续随机数均基于此种子产生，保证实验可重复性。

        self.window_size = ws = 512  # The size of the PyGame window

• 作用：将窗口尺寸设置为 512 像素，并赋值给局部变量 ws 及实例变量 self.window_size。
• 示例：生成的 PyGame 窗口为 512×512 像素。
• 意义：决定了渲染区域的尺寸。

        self.render_size = render_size

• 作用：将构造函数传入的 render_size（默认 96）赋给实例变量，用于后续图像缩放。
• 意义：方便在低分辨率下渲染环境快照。

        self.sim_hz = 100

• 作用：设置物理仿真的频率为 100 赫兹。
• 示例：仿真每秒更新 100 次，步长 dt = 1/100 = 0.01 秒。

        # Local controller params.
        self.k_p, self.k_v = 100, 20    # PD control.z

• 作用：定义局部控制器的 PD 控制参数，其中比例系数 k_p=100，微分系数 k_v=20。
• 意义：在控制 agent（智能体）时，通过 PD 控制使其平滑向目标移动。
• 示例：若 agent 与目标位置误差为 0.5，则产生的控制加速度为 100×0.5=50，加上速度反馈部分 20×(0-当前速度)。

        self.control_hz = self.metadata['video.frames_per_second']

• 作用：将控制频率设置为视频帧率，即 10 赫兹。
• 意义：控制更新与渲染帧率一致，便于同步观察。

        # legcay set_state for data compatiblity
        self.legacy = legacy

• 作用：将 legacy 参数存入实例变量，用于决定状态设置的顺序（旧数据兼容）。
• 意义：确保与旧版数据处理流程兼容。

        # agent_pos, block_pos, block_angle
        self.observation_space = spaces.Box(
            low=np.array([0,0,0,0,0], dtype=np.float64),
            high=np.array([ws,ws,ws,ws,np.pi*2], dtype=np.float64),
            shape=(5,),
            dtype=np.float64
        )

• 作用：定义环境的观察空间为一个 Box（连续空间），包含 5 个数值。
• 参数解释：
  • 下界为 [0, 0, 0, 0, 0]；上界为 [512, 512, 512, 512, 2π]。
  • 通常表示：agent 的位置 (x,y)，积木的位置 (x,y) 以及积木的角度。
• 示例：观察可能为 [250.0, 400.0, 260.0, 300.0, 1.57]。

        # positional goal for agent
        self.action_space = spaces.Box(
            low=np.array([0,0], dtype=np.float64),
            high=np.array([ws,ws], dtype=np.float64),
            shape=(2,),
            dtype=np.float64
        )

• 作用：定义环境的动作空间为二维连续空间，表示 agent 的目标位置。
• 参数解释：
  • 下界为 [0,0]，上界为 [512,512]。
• 示例：动作可能为 [300.0, 350.0]。

        self.block_cog = block_cog
        self.damping = damping
        self.render_action = render_action

• 作用：将构造函数传入的参数 block_cog、damping、render_action 分别保存到实例变量中。

        """
        If human-rendering is used, `self.window` will be a reference
        to the window that we draw to. `self.clock` will be a clock that is used
        to ensure that the environment is rendered at the correct framerate in
        human-mode. They will remain `None` until human-mode is used for the
        first time.
        """

• 作用：多行字符串注释，说明 human 模式下 window、clock 和 screen 的作用和初始化时机。

        self.window = None
        self.clock = None
        self.screen = None

• 作用：初始化渲染相关变量为 None，后续在 render() 方法中根据需要创建 PyGame 窗口和时钟。
• 意义：延迟创建窗口，提高初始化速度，只有在真正渲染时才创建图形界面。

        self.space = None
        self.teleop = None
        self.render_buffer = None
        self.latest_action = None
        self.reset_to_state = reset_to_state

• 作用：初始化其他关键变量：
  • self.space：物理仿真空间，稍后由 _setup() 方法创建；
  • self.teleop：用于手动控制的标志；
  • self.render_buffer：保存渲染帧；
  • self.latest_action：记录最新动作；
  • self.reset_to_state：存储重置状态参数。
• 意义：为环境内部的各个模块（物理、控制、渲染）预留变量，后续方法将赋予具体值。

八.2 def reset()

    def reset(self):
        seed = self._seed

• 作用：将先前保存的种子赋值给局部变量 seed。
• 示例：如果 _seed 为 1234，则 seed = 1234。

        self._setup()

• 作用：调用内部 _setup() 方法，构建物理空间、添加物体和障碍等（后面会介绍到）。
• 意义：重置环境前需要重新设置物理仿真空间。

        if self.block_cog is not None:
            self.block.center_of_gravity = self.block_cog

• 作用：若传入了 block_cog 参数，则将积木（block）的重心设置为该值。
• 示例：若 block_cog 为 (260, 300)，则 self.block.center_of_gravity = (260,300)。

        if self.damping is not None:
            self.space.damping = self.damping

• 作用：若 damping 参数不为 None，则设置物理空间的阻尼系数。
• 示例：若 damping 为 0.9，则 self.space.damping = 0.9，有助于模拟摩擦或空气阻力。

        state = self.reset_to_state

• 作用：尝试从 reset_to_state 中获取预设状态。
• 意义：允许用户指定重置时的精确初始状态，便于实验复现。

        if state is None:
            rs = np.random.RandomState(seed=seed)

• 作用：若 reset_to_state 为 None，则创建一个使用指定种子的 RandomState 对象。 可以使用rs代替所有np.random
• 示例：如果 seed 为 1234，则 rs = np.random.RandomState(1234)。

            state = np.array([
                rs.randint(50, 450), rs.randint(50, 450),
                rs.randint(100, 400), rs.randint(100, 400),
                rs.randn() * 2 * np.pi - np.pi
                ])

• 作用：随机生成状态数组，包含 5 个数：
  • 第 1、2 个数：agent 位置 x, y 在 [50,450) 内；
  • 第 3、4 个数：block 位置 x, y 在 [100,400) 内；
  • 第 5 个数：block 角度，在 -π 到 π 内随机生成（由于表达的是角度，所以会在 -π 到 π 之间）。
• 示例：可能生成状态 [200, 300, 150, 350, 0.5]。

        self._set_state(state)

• 作用：调用内部 _set_state 方法，将生成的状态设置到 agent 和 block 上（后面会介绍）。
• 意义：确保环境状态与生成状态同步。

        obs = self._get_obs()

• 作用：调用 _get_obs 获取当前观察值（后面会介绍）。
• 意义：观察通常包含 agent、block 的位置信息和角度。

        info = self._get_info()

• 作用：调用 _get_info 获取额外信息，如速度、接触点数等（后面会介绍）。
• 意义：info 为调试或分析提供辅助数据。

        return obs, info

• 作用：返回观察值和 info，供外部调用 reset() 后获得环境初始状态和相关信息。

八.3 def step()

    def step(self, action):
        dt = 1.0 / self.sim_hz

• 作用：计算仿真步长 dt。
• 示例：sim_hz 为 100 时，dt = 0.01 秒。

        self.n_contact_points = 0

• 作用：重置接触点计数，用于后续碰撞检测统计。

        n_steps = self.sim_hz // self.control_hz

• 作用：计算在每个控制周期内需要进行多少物理仿真步数。
• 示例：sim_hz=100，control_hz=10，则 n_steps = 10。

        if action is not None:
            self.latest_action = action

• 作用：如果传入动作不为 None，将动作记录到 latest_action 中，便于后续渲染显示。
• 示例：如果 action 为 [300, 350]。

            for i in range(n_steps):

• 作用：循环 n_steps 次，表示在当前控制周期内进行多次物理步进。
• 示例：循环 10 次。

                acceleration = self.k_p * (action - self.agent.position) + self.k_v * (Vec2d(0, 0) - self.agent.velocity)

• 作用：根据 PD 控制公式计算加速度。

PD 控制器：PD 代表比例（Proportional）和微分（Derivative）控制。
比例部分：根据当前位置误差（目标位置与当前位置之差）产生的控制信号，与误差成正比。
微分部分：根据速度误差（目标速度与当前速度之差）产生控制信号，起到“阻尼”作用，帮助系统平滑过渡并减少超调。
设计理由：

• 比例控制：确保智能体尽快朝向目标方向移动；误差越大，加速度越大。
• 微分控制：起到阻尼作用，防止智能体运动过快或出现震荡，改善控制系统的动态响应。

输入输出：

• 输入：
  • action（目标位置），
  • self.agent.position（当前智能体位置），
  • self.agent.velocity（当前速度），
  • 控制参数 k_p 和 k_v。
• 输出：
  • 一个加速度向量，用于更新智能体的速度，从而影响其运动行为。

• 示例：
  • 若 action = (300,350)，agent.position = (256,400)，则误差 = (44, -50)；
  • k_p=100，则比例项 = (4400, -5000)；
  • 若 agent.velocity = (5, 5) 则 (0,0)-velocity = (-5,-5)，k_v=20，微分项 = (-100,-100)；
  • 合计加速度 = (4300, -5100)（单位依赖于仿真设置）。

                self.agent.velocity += acceleration * dt

• 作用：更新 agent 的速度，使用欧拉积分。
• 示例：若 dt=0.01，则速度更新约为 agent.velocity += (43, -51)。

                self.space.step(dt)

• 作用：让物理仿真空间前进一个时间步 dt，更新所有物体状态。
• 意义：通过多次调用保证仿真精度。

        goal_body = self._get_goal_pose_body(self.goal_pose)

• 作用：根据 goal_pose（目标位置和角度）构造一个临时的目标物体。
• 示例：若 goal_pose = [256,256,π/4]，则生成 body.position = [256,256]，angle = π/4。

        goal_geom = pymunk_to_shapely(goal_body, self.block.shapes)

• 作用：调用辅助函数 pymunk_to_shapely 将目标物体的形状转换为 Shapely 多边形，用于几何运算。
• 意义：方便后续计算目标与 block 的重叠区域。

        block_geom = pymunk_to_shapely(self.block, self.block.shapes)

• 作用：同上，将实际 block 的形状转换为 Shapely 多边形。

        intersection_area = goal_geom.intersection(block_geom).area

• 作用：计算目标区域与 block 之间的交集面积。
• 示例：若目标面积为 500 平方像素，交集面积为 450，则 intersection_area = 450。

        goal_area = goal_geom.area

• 作用：获得目标区域的总面积。
• 示例：goal_area = 500。

        coverage = intersection_area / goal_area

• 作用：计算覆盖率，即 block 覆盖目标区域的比例。
• 示例：450/500 = 0.9。

        reward = np.clip(coverage / self.success_threshold, 0, 1)

• 作用：根据覆盖率计算奖励，将其归一化到 [0,1] 范围。
• 示例：若 success_threshold 为 0.95，则 reward = clip(0.9/0.95, 0, 1) ≈ 0.947。

        done = coverage > self.success_threshold

• 作用：判断是否完成任务：如果覆盖率大于 success_threshold，则任务完成。
• 示例：若 0.9 > 0.95 为 False，则 done 为 False。

        terminated = done
        truncated = done

• 作用：将任务完成状态赋给 terminated 和 truncated（可能与 Gym 的新 API 兼容）。
• 意义：用于告知智能体任务是否结束。

        observation = self._get_obs()

• 作用：获取当前环境观察值。

        info = self._get_info()

• 作用：获取当前环境的附加信息（如接触点、速度等）。

        return observation, reward, terminated, truncated, info

• 作用：返回 step() 方法的五个输出，符合 Gym 的 API 标准。

八.4 def render()

    def render(self, mode):
        return self._render_frame(mode)

• 作用：调用内部 _render_frame 方法进行渲染，并返回渲染结果。
• 输入：mode 指定渲染模式（如 “human” 或 “rgb_array”）。
• 意义：统一渲染接口，方便外部调用。

八.5 def teleop_agent()

    def teleop_agent(self):
        TeleopAgent = collections.namedtuple('TeleopAgent', ['act'])

• 作用：创建一个名为 TeleopAgent 的命名元组类型，其字段为 act。
• 意义：封装手动控制（teleoperation）的接口。

        def act(obs):
            act = None

• 作用：定义内部函数 act，输入为观察 obs，初始动作 act 设为 None。
• 意义：在后续判断中决定是否接收鼠标操作。

            mouse_position = pymunk.pygame_util.from_pygame(Vec2d(*pygame.mouse.get_pos()), self.screen)

• 作用：获取当前鼠标在屏幕上的位置，并将其转换为 pymunk 坐标。
• 示例：若 pygame.mouse.get_pos() 返回 (150, 200)，转换后可能得到 Vec2d(150, 200) 或考虑坐标翻转。

            if self.teleop or (mouse_position - self.agent.position).length < 30:
                self.teleop = True
                act = mouse_position

• 作用：如果当前已处于手动控制状态或鼠标与 agent 位置的距离小于 30，则开启手动控制，并将动作设为鼠标位置。
• 示例：若 agent.position 为 (140, 190) 与鼠标位置 (150, 200) 的距离约为 14.14，则触发条件，act = mouse_position。

            return act

• 作用：返回计算得到的动作。

        return TeleopAgent(act)

• 作用：返回一个 TeleopAgent 对象，其 act 字段绑定为上述内部函数。

八.6 def _get_obs()

    def _get_obs(self):
        obs = np.array(
            tuple(self.agent.position) \
            + tuple(self.block.position) \
            + (self.block.angle % (2 * np.pi),))
        return obs

• 作用：构造环境的观察值。
• 逐行说明：
  • tuple(self.agent.position)：将 agent 的位置（例如 (256, 400)）转换为元组。
  • tuple(self.block.position)：将 block 的位置（例如 (256, 300)）转换为元组。
  • (self.block.angle % (2 * np.pi),)：将 block 的角度取模 2π，确保角度在 [0, 2π) 内。
  • 最后将以上部分组合成一个 5 元素的元组，并转换为 np.array。
• 示例：最终 obs 可能为 np.array([256, 400, 256, 300, 0.0])。

八.7 def _get_goal_pose_body()

    def _get_goal_pose_body(self, pose):
        mass = 1

• 作用：将目标位姿（pose）转换为一个临时物体，用于计算目标区域。
• 示例：设 pose 为 [256,256,π/4]，此处 mass = 1 单位质量。

        inertia = pymunk.moment_for_box(mass, (50, 100))

• 作用：计算质量为 1 的物体、尺寸为 (50, 100) 的矩形惯性。
• 意义：惯性用于物理模拟，但在此处主要用于构造 body。
• 数据类型：float

        body = pymunk.Body(mass, inertia)

• 作用：创建一个新的 pymunk.Body 对象，参数为质量和惯性。
• 示例：生成 body 对象，其质量为 1，惯性为计算结果。

        body.position = pose[:2].tolist()

• 作用：将 pose 的前两个数（例如 [256,256]）设为 body 的位置。
• 意义：目标物体在世界坐标中的位置。

        body.angle = pose[2]

• 作用：将 pose 的第三个数（例如 π/4）设为 body 的角度。

        return body

• 作用：返回构造好的 body 对象。

八.8 def _get_info()

    def _get_info(self):
        n_steps = self.sim_hz // self.control_hz

• 作用：计算每个控制周期中的仿真步数。
• 示例：100 // 10 = 10。

        n_contact_points_per_step = int(np.ceil(self.n_contact_points / n_steps))

• 作用：计算平均每步的接触点数，向上取整（np.ceil()向上取整操作）。
• 示例：若 self.n_contact_points 为 15，n_steps 为 10，则值约为 2。

        info = {
            'pos_agent': np.array(self.agent.position),
            'vel_agent': np.array(self.agent.velocity),
            'block_pose': np.array(list(self.block.position) + [self.block.angle]),
            'goal_pose': self.goal_pose,
            'n_contacts': n_contact_points_per_step}

• 作用：构造一个字典 info，包含：
  • agent 的位置、速度；
  • block 的位置和角度；
  • 当前目标位姿；
  • 平均接触点数。
• 示例：
  • pos_agent: np.array([256,400])
  • vel_agent: np.array([5,0])
  • block_pose: np.array([256,300,0.0])
  • goal_pose: [256,256,π/4]
  • n_contacts: 2

        return info

• 作用：返回 info 字典。

八.9 def _render_frame()

该方法较长，主要负责将物理空间渲染到屏幕上，并生成 RGB 图像。下面按段逐行说明。

    def _render_frame(self, mode):

• 作用：定义内部渲染方法，mode 指定渲染模式。

        if self.window is None and mode == "human":
            pygame.init()
            pygame.display.init()
            self.window = pygame.display.set_mode((self.window_size, self.window_size))

• 作用：
  • 若处于 human 模式且 window 尚未创建，则初始化 pygame，创建显示窗口。
• 示例：创建一个 512×512 的窗口。

        if self.clock is None and mode == "human":
            self.clock = pygame.time.Clock()

• 作用：初始化 pygame 时钟，用于控制帧率。

        canvas = pygame.Surface((self.window_size, self.window_size))

• 作用：创建一个离屏 Surface 对象 canvas，尺寸与窗口相同。

        canvas.fill((255, 255, 255))

• 作用：用白色填充 canvas，清空背景。
• 示例：(255,255,255) 表示白色。

        self.screen = canvas

• 作用：将 canvas 赋值给 self.screen 以便后续绘制使用。

        draw_options = DrawOptions(canvas)

• 作用：创建 DrawOptions 对象，传入 canvas，用于调试绘制 pymunk 物理对象。

        # Draw goal pose.
        goal_body = self._get_goal_pose_body(self.goal_pose)

• 作用：生成目标物体 body，基于当前 goal_pose。

        for shape in self.block.shapes:
            goal_points = [pymunk.pygame_util.to_pygame(goal_body.local_to_world(v), draw_options.surface) for v in shape.get_vertices()]
            goal_points += [goal_points[0]]
            pygame.draw.polygon(canvas, self.goal_color, goal_points)

• 作用：
  • 遍历 block 的每个形状，计算目标物体在世界坐标下的顶点；
  • 将顶点转换为 pygame 坐标；
  • 闭合多边形后用目标颜色绘制多边形。
• 函数：
  • 遍历 self.block.shapes 中的所有形状（shapes）。假设 self.block 是代表积木（block）的物体，其可能包含多个碰撞形状（例如，一个 T 型积木可能由两个多边形组合而成）。
  • shape.get_vertices() 返回该形状所有顶点的局部坐标。例如，如果 shape 表示一个矩形，其顶点可能为：[(0,0), (50,0), (50,100), (0,100)]。
  • goal_body 是一个临时构造的物体，代表目标位姿。假设 goal_body 的位置为 (256,256) 且角度为 45°（π/4）。那么对于局部顶点 (0,0)，经过 goal_body.local_to_world((0,0)) 得到的世界坐标可能为 (256,256)；而 (50,0) 经过旋转和平移后，可能得到 (256+35,256+35) 的值（具体数值依赖旋转矩阵计算）。
  • draw_options.surface 是一个指向 pygame 绘图画布（Surface）的引用，它在 DrawOptions 类的构造函数中被设置，并作为所有绘图操作的目标画布。
  • 使用 pymunk.pygame_util.to_pygame(..., draw_options.surface) 将世界坐标转换为 pygame 屏幕坐标（考虑 y 轴翻转等）。
  • goal_points += [goal_points[0]]：将列表的第一个点追加到末尾，确保多边形闭合。
  • pygame.draw.polygon(canvas, self.goal_color, goal_points)：在 pygame 的 canvas 上绘制一个多边形，顶点为 goal_points，填充颜色为 self.goal_color。
    • canvas 是一个 pygame Surface 对象，之前已经创建并填充为白色背景。
    • self.goal_color 是目标区域的颜色，通常设置为一个易于识别的颜色，如浅绿色（LightGreen）。
• 示例：若 goal_color 为 LightGreen，绘制出绿色区域。

        # Draw agent and block.
        self.space.debug_draw(draw_options)

• 作用：调用 pymunk 的 debug_draw 方法，将物理空间中 agent、block 等物体绘制到 canvas 上。

        if mode == "human":
            self.window.blit(canvas, canvas.get_rect())
            pygame.event.pump()
            pygame.display.update()

• 作用：
  • 若处于 human 模式，将 canvas 内容复制到窗口；
  • 调用 pygame.event.pump() 处理事件；
  • 更新显示窗口。

        img = np.transpose(
                np.array(pygame.surfarray.pixels3d(canvas)), axes=(1, 0, 2)
            )

• 作用：
  • 使用 pygame.surfarray 将 canvas 像素数据转为 NumPy 数组；
  • 转置轴顺序，使图像符合常见的 (height, width, channels) 格式。
• 示例：若 canvas 为 512×512 像素，则 img.shape 可能为 (512,512,3)。

        img = cv2.resize(img, (self.render_size, self.render_size))

• 作用：使用 OpenCV 将图像调整为 render_size 大小（例如 96×96）。

        if self.render_action:
            if self.render_action and (self.latest_action is not None):
                action = np.array(self.latest_action)
                coord = (action / 512 * 96).astype(np.int32)
                marker_size = int(8/96*self.render_size)
                thickness = int(1/96*self.render_size)
                cv2.drawMarker(img, coord,
                    color=(255,0,0), markerType=cv2.MARKER_CROSS,
                    markerSize=marker_size, thickness=thickness)

• 作用：
  • 若 render_action 为 True 且最新动作存在，则：
    • 将 latest_action 转为 NumPy 数组；
    • 根据窗口尺寸（512）与渲染尺寸（96）的比例调整动作坐标；
    • 计算标记大小和线宽；
    • 在图像上绘制一个红色（255,0,0）的十字标记，指示动作目标位置。
• 示例：若 latest_action = [256,256]，则 coord = ([256/51296,256/51296]) = ([48,48])，在图像中央绘制红色十字。

        return img

• 作用：返回渲染后的图像（RGB 数组）。

八.10 def close()

    def close(self):
        if self.window is not None:
            pygame.display.quit()
            pygame.quit()

• 作用：
  • 若窗口存在，退出 pygame 显示和 pygame 整体，释放资源。
• 意义：防止程序退出时窗口卡死或资源泄露。

八.11 def seed()

    def seed(self, seed=None):
        if seed is None:
            seed = np.random.randint(0,25536)

• 作用：
  • 如果未传入种子，则随机生成一个整数种子，范围 0 到 25535。
• 示例：可能生成种子 1234。

        self._seed = seed

• 作用：将生成或传入的种子保存到实例变量 _seed 中。

        self.np_random = np.random.default_rng(seed)

• 作用：利用 NumPy 的新随机数生成器（default_rng）创建一个随机数生成器实例，便于后续生成随机数。

八.12 def _handle_collision()

    def _handle_collision(self, arbiter, space, data):
        self.n_contact_points += len(arbiter.contact_point_set.points)

• 作用：
  • 在碰撞处理回调中，统计当前碰撞的接触点数，并累加到 n_contact_points。
• 示例：若 arbiter.contact_point_set.points 有 3 个点，则 n_contact_points 增加 3。

八.13 def _set_state()

    def _set_state(self, state):
        if isinstance(state, np.ndarray):
            state = state.tolist()

• 作用：
  • 如果传入 state 为 NumPy 数组，则转换为 Python 列表，便于后续处理。

        pos_agent = state[:2]
        pos_block = state[2:4]
        rot_block = state[4]

• 作用：
  • 将 state 分割为：
    • 前两个元素为 agent 位置；
    • 第 3、4 个元素为 block 位置；
    • 第 5 个元素为 block 角度。
• 示例：若 state = [256,400,256,300,0.0]。

        self.agent.position = pos_agent

• 作用：将解析出的 agent 位置赋值给 agent 对象。

        if self.legacy:
            self.block.position = pos_block
            self.block.angle = rot_block
        else:
            self.block.angle = rot_block
            self.block.position = pos_block

• 作用：
  • 根据 legacy 标志决定设置 block 位置和角度的顺序，确保兼容旧数据。
  • 顺序不同可能影响物理仿真中相对于重心的旋转效果。

        self.space.step(1.0 / self.sim_hz)

• 作用：进行一次物理仿真步进，使状态修改生效。

八.14 def _set_state_local()

    def _set_state_local(self, state_local):
        agent_pos_local = state_local[:2]
        block_pose_local = state_local[2:]

• 作用：
  • 将局部状态拆分为 agent 的局部位置和 block 的局部位姿（位置和角度）。
• 示例：若 state_local = [x1, y1, x2, y2, theta]。

        tf_img_obj = st.AffineTransform(
            translation=self.goal_pose[:2],
            rotation=self.goal_pose[2])

• 作用：
  • 创建一个仿射变换 tf_img_obj，以目标位姿的平移（前两个元素）和旋转（第三个元素）为参数。
• 示例：若 goal_pose = [256,256,π/4]。

        tf_obj_new = st.AffineTransform(
            translation=block_pose_local[:2],
            rotation=block_pose_local[2]
        )

• 作用：
  • 创建另一个仿射变换 tf_obj_new，以 block 的局部位置和角度为参数。

        tf_img_new = st.AffineTransform(
            matrix=tf_img_obj.params @ tf_obj_new.params
        )

• 作用：
  • 通过矩阵乘法组合上述两个仿射变换，得到一个新的变换 tf_img_new。
  • “@” 表示矩阵乘法运算。

        agent_pos_new = tf_img_new(agent_pos_local)

• 作用：
  • 使用组合变换将 agent 的局部位置转换为新的位置。

        new_state = np.array(
            list(agent_pos_new[0]) + list(tf_img_new.translation) \
                + [tf_img_new.rotation])

• 作用：
  • 构造新的状态数组，包含转换后的 agent 位置、tf_img_new 的平移部分（作为 block 位置）以及旋转部分（作为 block 角度）。
• 示例：可能生成新状态 [agent_x, agent_y, block_x, block_y, rotation]。

        self._set_state(new_state)

• 作用：调用 _set_state 方法，将新状态赋值给环境。

        return new_state

• 作用：返回新状态，供外部参考。

八.15 def _setup()

    def _setup(self):
        self.space = pymunk.Space()

• 作用：创建一个新的 pymunk.Space 对象，作为物理仿真的空间。

        self.space.gravity = 0, 0

• 作用：设置空间重力为 (0,0)，表示无重力环境。
• 示例：无重力适合 2D 平面上的具身操作任务。

        self.space.damping = 0

• 作用：将阻尼系数设为 0，初始时不引入额外阻力。

        self.teleop = False

• 作用：初始化手动控制标志为 False。

        self.render_buffer = list()

• 作用：初始化渲染缓冲区为空列表。

        walls = [
            self._add_segment((5, 506), (5, 5), 2),
            self._add_segment((5, 5), (506, 5), 2),
            self._add_segment((506, 5), (506, 506), 2),
            self._add_segment((5, 506), (506, 506), 2)
        ]

• 作用：
  • 构造四条边界墙，使用 _add_segment 方法创建线段。
  • 坐标示例：左边墙从 (5,506) 到 (5,5)；上边从 (5,5) 到 (506,5) 等。

        self.space.add(*walls)

• 作用：将所有墙添加到物理空间中。

        self.agent = self.add_circle((256, 400), 15)

• 作用：调用 add_circle，在位置 (256,400) 添加一个半径为 15 的 agent，返回 agent 的 body 对象。

        self.block = self.add_tee((256, 300), 0)

• 作用：调用 add_tee，在位置 (256,300) 添加一个“T”型物体，角度为 0，返回其 body 对象。

        self.goal_color = pygame.Color('LightGreen')

• 作用：将目标区域颜色设置为浅绿色。

        self.goal_pose = np.array([256,256,np.pi/4])

• 作用：设定目标位姿为 x=256, y=256, 角度=π/4。

        self.collision_handeler = self.space.add_collision_handler(0, 0)

• 作用：为 collision type 0 与 0 添加碰撞处理器。
• 意义：所有形状默认 collision type 为 0，因此此处理器将捕获所有碰撞。

        self.collision_handeler.post_solve = self._handle_collision

• 作用：设置碰撞后求解回调函数为 _handle_collision，用于统计接触点。

        self.n_contact_points = 0

• 作用：初始化接触点计数为 0。

        self.max_score = 50 * 100

• 作用：设定最大得分，数值为 5000（可能用于评估）。

        self.success_threshold = 0.95    # 95% coverage.

• 作用：设定任务成功的覆盖率阈值为 95%。

八.16 def _add_segment()

    def _add_segment(self, a, b, radius):
        shape = pymunk.Segment(self.space.static_body, a, b, radius)

• 作用：创建一个静态边界段，端点 a 与 b，线宽 radius。
• 示例：如 a=(5,506), b=(5,5), radius=2。

        shape.color = pygame.Color('LightGray')

• 作用：设置边界颜色为浅灰色。

        return shape

• 作用：返回创建的边界 shape。

八.17 def add_circle()

    def add_circle(self, position, radius):
        body = pymunk.Body(body_type=pymunk.Body.KINEMATIC)

• 作用：创建一个静态（运动学）物体作为 agent。
• 示例：position 可能为 (256,400)，radius 为 15。

        body.position = position

• 作用：设置 body 的位置为传入的 position。

        body.friction = 1

• 作用：设置摩擦系数为 1，保证物体间有足够摩擦力。

        shape = pymunk.Circle(body, radius)

• 作用：基于 body 创建一个圆形碰撞形状，半径为 radius。

        shape.color = pygame.Color('RoyalBlue')

• 作用：将 agent 的颜色设置为皇家蓝。

        self.space.add(body, shape)

• 作用：将 agent 的 body 和 shape 添加到物理空间中。

        return body

• 作用：返回创建的 agent body。

八.18 def add_box()

    def add_box(self, position, height, width):
        mass = 1

• 作用：设定 box 物体的质量为 1。

        inertia = pymunk.moment_for_box(mass, (height, width))

• 作用：计算给定质量和尺寸的盒子惯性。

        body = pymunk.Body(mass, inertia)

• 作用：创建一个动态 body。

        body.position = position

• 作用：设置位置为传入参数。

        shape = pymunk.Poly.create_box(body, (height, width))

• 作用：基于 body 创建一个矩形多边形形状。

        shape.color = pygame.Color('LightSlateGray')

• 作用：将 box 的颜色设为浅石板灰。

        self.space.add(body, shape)

• 作用：将 body 和 shape 添加到物理空间。

        return body

• 作用：返回创建的 box body。

八.19 def add_tee()

    def add_tee(self, position, angle, scale=30, color='LightSlateGray', mask=pymunk.ShapeFilter.ALL_MASKS()):
        mass = 1

• 作用：创建一个“T”型物体，质量设为 1。
• 参数：position 为放置位置，angle 为旋转角度，scale 控制尺寸，color 指定颜色，mask 指定碰撞掩码。

        length = 4

• 作用：定义 T 型结构中横杆或竖杆的基准长度为 4（单位可理解为块数）。

        vertices1 = [(-length*scale/2, scale),
                     ( length*scale/2, scale),
                     ( length*scale/2, 0),
                     (-length*scale/2, 0)]

• 作用：定义 T 型物体第一个部件（可能为横杆）的顶点。
• 示例：scale=30, length=4，则顶点计算为：
  • (-60, 30), (60,30), (60,0), (-60,0)。

        inertia1 = pymunk.moment_for_poly(mass, vertices=vertices1)

• 作用：计算第一个部件的惯性。

        vertices2 = [(-scale/2, scale),
                     (-scale/2, length*scale),
                     ( scale/2, length*scale),
                     ( scale/2, scale)]

• 作用：定义第二个部件（可能为竖杆）的顶点。
• 示例：scale=30, length=4，则顶点为：
  • (-15,30), (-15,120), (15,120), (15,30)。

        # inertia2 = pymunk.moment_for_poly(mass, vertices=vertices2)
        inertia2 = pymunk.moment_for_poly(mass, vertices=vertices1)

• 作用：注意：这里似乎存在笔误，使用了 vertices1 而非 vertices2 计算 inertia2。但按照代码，计算第二部分的惯性（数值可能不正确）。

        body = pymunk.Body(mass, inertia1 + inertia2)

• 作用：创建一个 body，其总惯性为两个部件的惯性之和。

        shape1 = pymunk.Poly(body, vertices1)
        shape2 = pymunk.Poly(body, vertices2)

• 作用：为 body 分别创建两个多边形形状，表示 T 型物体的两个组成部分。

        shape1.color = pygame.Color(color)
        shape2.color = pygame.Color(color)

• 作用：将两个形状的颜色均设置为传入颜色，例如 ‘LightSlateGray’。

        shape1.filter = pymunk.ShapeFilter(mask=mask)
        shape2.filter = pymunk.ShapeFilter(mask=mask)

• 作用：设置碰撞掩码，控制哪些物体可与之碰撞。

        body.center_of_gravity = (shape1.center_of_gravity + shape2.center_of_gravity) / 2

• 作用：将 body 的重心设为两个形状重心的平均值，保证物理行为合理。

        body.position = position

• 作用：设置 body 的位置为传入参数。

        body.angle = angle

• 作用：设置 body 的旋转角度为传入的 angle。

        body.friction = 1

• 作用：设置摩擦系数为 1。

        self.space.add(body, shape1, shape2)

• 作用：将 T 型物体的 body 和两个形状添加到物理空间。

        return body

• 作用：返回创建的 T 型物体 body。

总结

整个 PushTEnv 类构建了一个具备以下功能的环境：

• 物理仿真：利用 pymunk 构建无重力空间，添加 agent、block（T型积木）和边界墙。
• 状态与奖励：通过 _get_obs、_get_info 和 step 方法获取环境状态，并基于目标区域与 block 的重叠计算奖励。
• 渲染：支持人机交互模式与图像返回模式，通过 _render_frame 方法绘制物理仿真图像，同时展示目标区域、 agent、block 以及动作标记。
• 控制与复现：支持 PD 控制，种子设置以及状态重置，方便实验复现和调试。