检查点与仿真状态#

一次仿真有两样值得保存的东西。第一是它的动态状态:物理求解器每步推进的那些数值,例如关节位置、速度和粒子场。捕获该状态并在之后恢复,可以让你确定性地回退仿真,或在回合之间重置环境,并且从你离开时的确切状态开始。第二是场景本身:实体、它们的几何形状,以及创建场景时使用的每个选项。把它写入文件,别人就能在没有该场景所需任何资产的机器上,打开你构建的那个场景。

Genesis World 提供四种保存方式,分别满足不同需求。

  • 内存中的动态状态: scene.get_state() 返回一个 SimState 对象,scene.reset(state=...) 将其写回。速度快,是强化学习中回合重置的基础。

  • 内存中的整个场景及其状态: pickle 协议。pickle.dumps(scene) 会生成一个已构建的副本,其步进行为与原始场景完全一致;scene.__getstate__() 和 scene.__setstate__() 则就地保存和恢复已构建的场景。可用它为一次运行创建检查点,或把场景交给另一个进程。

  • 磁盘上的场景: scene.export(path) 写出 .gscene 文件,gs.Scene.load(path) 将其打开为新场景,随时可以构建。可用它共享场景、把场景附加到缺陷报告,或在不带资产的情况下重建场景。

  • 磁盘上的运行: scene.save_checkpoint(path) 写出场景及其仿真的每个数组,scene.start_recording(gs.recorders.TrajectoryFile(...)) 每步写出一个帧。gs.Scene.load_checkpoint(path) 和 gs.Scene.load_trajectory(path) 会在已构建的场景中打开它们,以便在任何机器和后端上继续、定位和重放。

状态快照假定它来自的场景已经构建。场景文件不携带仿真状态,因此从中打开的场景处于其实体被赋予的配置,状态必须通过再次步进来复现。检查点或轨迹同时携带两者,因此它打开的场景会停在被记录场景所在的位置。

状态模型#

快照只捕获动态状态:随着仿真步进而变化的字段。它不捕获场景的结构:实体、它们的 morph、求解器选项或环境数量。结构由你构建场景的方式固定,恢复快照时假定结构已经就位。

  • SimState: 由 scene.get_state() 返回的对象。它为每个活动求解器保存一个逐求解器状态对象,并按环境进行批处理。

  • 动态状态: 位置、速度,以及每个求解器积分的内部字段。快照保存的就是这些。

  • 静态结构: 实体、morph、几何形状和求解器配置。快照不包含它,场景文件则是传递它的方式(见下文)。

由于结构不属于快照,快照只对拍摄它的场景或用相同方式构建的场景有效。恢复到实体或求解器选项不同的场景中,行为是未定义的。

在内存中快照与恢复#

scene.get_state() 将当前状态读入 SimState。scene.reset() 将场景恢复到存储的初始状态,scene.reset(state=...) 则恢复任意快照:

scene.build()

for _ in range(100):
    scene.step()

state = scene.get_state()  # snapshot the state at step 100

for _ in range(50):
    scene.step()

scene.reset(state=state)  # rewind to the snapshot; the physics continue from there

无论恢复哪个快照,重置都会将它所涉及环境的仿真时间设回零,因此 scene.get_time() 从重置时开始计时,而不是从构建时开始。

Warning

将 state 传给 reset() 还会把它注册为场景的初始状态。之后直接调用 scene.reset() 会返回到这个快照,而不是场景在构建时的状态。如果你同时需要两者,请为构建时的状态保留单独的引用。

逐求解器状态对象只是普通的属性容器。刚体求解器的状态包含 qpos、dofs_vel、links_pos 和 links_quat;读取其中一个字段的方式如下:

state = scene.get_state()
rigid_state = state.solvers_state[scene.solvers.index(scene.rigid_solver)]
qpos = rigid_state.qpos  # shape ([n_envs,] n_qs)

可微分示例 examples/deformable/differentiable_push.py 使用了这种模式:它调用 scene.reset(),从固定的初始状态重新开始每一轮优化,并读取 scene.get_state().solvers_state[...],在 rollout 中途根据粒子位置计算损失。

在并行仿真中重置环境#

在并行仿真中,状态按环境批处理,reset() 接受 envs_idx 参数,因此你可以重置一个子集而不影响其余环境。这正是强化学习中逐环境回合重置的机制:当部分环境结束时,只把这些环境恢复到初始状态,让其他环境继续运行。

scene.build(n_envs=4096)

init_state = scene.get_state()  # the state all environments reset to

for step in range(episode_length):
    scene.step()
    obs, reward, done = get_observations()

    if done.any():
        done_envs = torch.where(done)[0]  # indices of finished environments
        scene.reset(state=init_state, envs_idx=done_envs)
  • envs_idx: 要重置的环境,可以是任何类似数组的索引形式。None(默认值)会重置所有环境。

  • 部分重置: 使用 envs_idx 时,只有被选中的环境接受新状态;其他环境会不受打断地继续推进。

envs_idx 仅适用于使用环境构建的场景,因此在非并行场景中使用它会抛出异常。

将场景作为文件共享#

scene.export(path) 写出场景的创作依据以及构建时解析出的内容:每个实体的描述(包括其几何形状、纹理和物理系数),以及创建场景时使用的每个选项。文件中不包含任何文件系统路径,因此它可以在不持有该场景所用任何网格、模型文件或纹理的机器上打开。gs.Scene.load(path) 会创建该场景,包含创建它时使用的每个实体,并等待构建:

# Wherever the scene was authored.
scene = gs.Scene()
scene.add_entity(gs.morphs.Plane())
robot = scene.add_entity(gs.morphs.MJCF(file="xml/franka_emika_panda/panda.xml"))
scene.export("franka.gscene")
# Anywhere else, with no asset on disk.
scene = gs.Scene.load("franka.gscene")
scene.build(n_envs=16)
scene.step()

添加实体会解析其描述,因此场景在构建前导出和构建后导出一样容易,而环境数量由构建所加载场景的人选择。场景文件非常适合附加到缺陷报告中:维护者打开它并步进你当时的确切场景,既不需要你的脚本,也不需要你的资产。

该文件指明自己包含的内容,而不是携带要执行的代码。Genesis World 只创建文件声明的选项、描述和网格,并拒绝与所声明类型相矛盾的值,因此打开陌生人提供的场景是安全的。加载时还会将文件中每个类的布局与当前布局比较,并拒绝在某个类含义不同的时期写出的文件,同时指出该类名。由另一个 Genesis World 版本写出、但类仍然匹配的文件可以加载,并会给出警告:它所描述的仿真可能以不同方式运行。

创建场景时使用的每个选项都作为一个 SceneOptions 对象传递,可在任何场景上通过 scene.options 访问。将它传给 gs.Scene(options=other.options),就能按另一个场景创建时的配置创建场景,这也是加载的场景恢复其选项的方式。物理选项保持记录时的值,因为它们决定记录状态所填充数组的大小。加载时传入的查看器和可视化器选项会覆盖记录值中它们所设置的字段,其他字段则保留记录值;加载时传入的渲染器会整体替换记录的那个。因此,以无头方式记录的场景可以在调试会话所需的查看器中打开:

scene = gs.Scene.load(
    "franka.gscene",
    show_viewer=True,
    viewer_options=gs.options.ViewerOptions(camera_pos=(2.0, 1.5, 1.0), camera_lookat=(0.0, 0.0, 0.3)),
)

描述携带什么,什么就会传递;其余内容,export 会告诉你:

  • 按名称拒绝: 任何会改变仿真、但描述未包含的内容,因为缺少它的文件会恢复出不同的物理行为:发射器、力场,以及任何不带描述的实体——即所有既非刚体也非运动学的实体。

  • 不写出,并给出点名警告: 观察或绘制仿真、而不是塑造仿真的内容。包括相机、记录器、传感器、Genesis World 每步调用的回调、从 HDR 或 EXR 文件读取的纹理,以及运行时赋予实体的可视顶点。请在加载后的场景中重新添加这些内容。

为已构建的场景创建检查点#

SimState 保存的是你读回的位置和速度。检查点则保存一步到下一步之间、每步都会读取的所有内容:降维空间状态和控制输入、接触和约束解、热启动缓存,以及上一步推导出的连杆和几何体位姿。因此,在每个后端上步进恢复后的场景,都会逐位落在原始场景到达的位置。

检查点就是 pickle 协议。scene.__getstate__() 返回一个 SceneCheckpoint:场景描述、构建场景时使用的环境布局,以及每个求解器的数组;这些数组在其所在位置复制,因此在同一进程内保存和恢复会留在设备上。scene.__setstate__(checkpoint) 会将按相同描述和布局构建的场景恢复到该状态,并拒绝来自其他构建的检查点,同时指出差异所在。对场景执行 pickle 会读取同一记录;反 pickle 时会创建并构建场景,然后恢复该状态:

import pickle

scene.build(n_envs=16)
for _ in range(100):
    scene.step()

checkpoint = scene.__getstate__()  # a device-resident copy of the state at step 100
twin = pickle.loads(pickle.dumps(scene))  # a built copy of the scene, standing at step 100

for _ in range(50):
    scene.step()
scene.__setstate__(checkpoint)  # back to step 100, and the next 50 steps land where they did

恢复还会像重置一样重新启动状态周边的组件:梯度磁带、传感器和记录器会忘记导致该状态的那次运行,查看器会重绘。每个环境的仿真时间会按记录恢复,因为它属于状态。可微分场景会一并携带其数组的梯度,因此在反向传播之后取得的检查点会恢复这些梯度。

检查点文件与轨迹文件#

scene.save_checkpoint(path) 写出 .gstraj 文件,其中包含 export 所写出的场景,以及仿真的每个数组(包括暂存缓冲区),从而保留失败运行的确切状态以供检查。gs.Scene.load_checkpoint(path) 会创建并构建场景,并恢复该状态:

scene.save_checkpoint("failing.gstraj")
scene = gs.Scene.load_checkpoint("failing.gstraj", show_viewer=True)
scene.step()  # continues from the saved state

轨迹就是每步一个检查点。在构建之前用 scene.start_recording 注册 TrajectoryFile,每一步都会追加一帧,直到调用 scene.stop_recording();后者还会记录运行停止时的状态:

scene.start_recording(gs.recorders.TrajectoryFile(filename="run.gstraj"))
scene.build()
for _ in range(1000):
    scene.step()
scene.stop_recording()

每一帧都会保留。每一步只会把自己的帧复制到一旁,不等待 GPU 完成该步;同时有一个独立线程压缩帧,另一个线程写入文件,因此在单个环境上录制可以保持实时。两种模式在大小与精确性之间取舍,由 mode 选择:

  • Exact(精确): 重放是逐位的,仿真会从任意帧完全按照它运行时的状态继续,但对于较大的批大小,录制会显著拖慢仿真。

  • Minimal(最小): 帧只保存最低限度的状态,因此小得多:其余内容通过运行一个完整的仿真步来恢复,可能在最低位上有差异,但休眠连杆的接触除外——这些接触取自定位之前场景当时的状态。在新加载的场景中从第一帧开始按顺序播放的重放不受影响。

默认情况下,只要一帧保持在 4 MiB 以内,轨迹就以精确模式记录;否则使用最小模式,因为每个被记录的步都会复制整个帧。

hz 控制每隔几步记录一次,chunk_size 设置每个块包含多少帧(块越大压缩效果越好,但定位越慢),max_size 限制文件大小:达到上限时会将其作为有效日志关闭,并在超过上限后的下一步抛出异常,因此默认情况下,一个一直保持录制状态的运行最多占用 2 GiB 磁盘空间。

gs.Scene.load_trajectory(path) 会创建并构建被记录的场景,并返回一个在其中播放的 Trajectory。第 i 帧是经过 i 步后的状态;seek(i) 将场景定位到那里,play() 逐帧定位并按录制时的速度重绘查看器,frame(i) 按名称返回某一帧的数组,time(i) 返回每个环境的仿真时间,timings(i) 返回导致第 i 帧的那一步中各个阶段的墙钟时间,名称与 scene.timings 一致。负索引从末尾计数。以早期格式版本写出的文件会被拒绝,并指出其版本。最后一帧与检查点文件一样,保存除配置和常量之外的每个数组,包括求解器的暂存数据,用于分析可能依赖平台、一旦共享就无法复现的失败:

trajectory = gs.Scene.load_trajectory("run.gstraj", show_viewer=True)
trajectory.seek(500)
trajectory.scene.step()  # steps on from frame 500 as the recorded run did
trajectory.play(loop=True)  # replays until the viewer is closed

轨迹可以在另一台机器、另一个后端或另一种精度下重放:帧的形状由选项、描述和环境数量决定,其值会转换为当前会话的精度。在命令行中,gs replay run.gstraj 会在查看器中打开被记录的场景并循环播放。

快照包含的内容#

每个求解器向 SimState 贡献的动态状态:

求解器

状态字段

刚体

qpos, dofs_vel, dofs_acc, links_pos, links_quat, friction_ratio

运动学

qpos, dofs_vel, links_pos, links_quat

MPM

pos, vel, C, F, Jp, active

SPH

pos, vel, active

PBD

pos, vel, free

FEM

pos, vel, active

可复现性说明#

  • 配置必须匹配。 快照按位置将字段恢复到已构建的场景中。实体、求解器选项和环境数量必须与生成快照的场景一致。这里没有兼容性检查:不匹配会导致失败或静默破坏状态。检查点或轨迹帧按名称恢复数组并检查其形状,因此来自其他构建的检查点会被拒绝,并指出差异所在。两者还携带场景描述、其实体和物理选项的摘要,并且只会恢复到共享该摘要的场景中,无论其查看器、可视化器和渲染器选项如何。摘要会在内存中导出描述(包括网格),每构建一个场景导出一次,每打开一个轨迹导出一次,绝不会在每次恢复时导出。

  • 精度限制精确性。 Genesis World 默认使用 32 位浮点数(见 初始化与后端)。因此,通过步进加载的场景来复现运行,或恢复经过文件的快照,只能精确到大致单精度,而不是逐位精确。如果需要更严格的可复现性,请使用 precision="64" 初始化。

另请参阅#

  • 并行仿真:状态如何按环境批处理。

  • Scene API:get_state、reset、export、load、save_checkpoint、load_checkpoint 和 load_trajectory 的完整签名。

  • 记录数据:记录轨迹所用的记录器框架。