可微分仿真#

Genesis World 可以让物理引擎以可微分模式运行,从而把定义在仿真输出上的损失一直反向传播到产生这些输出的输入。你可以用它优化控制信号、拟合物理参数,或让策略穿过动力学进行训练,而不是绕开动力学。

最小示例#

在 gs.options.SimOptions 上设置 requires_grad=True 启用可微分模式,向前运行仿真,然后对损失调用 backward():

import genesis as gs
import torch

gs.init(backend=gs.gpu)

scene = gs.Scene(
    sim_options=gs.options.SimOptions(
        dt=0.01,  # seconds
        requires_grad=True,  # enable differentiable mode
    ),
)
robot = scene.add_entity(gs.morphs.URDF(file="robot.urdf"))
scene.build()

# A control signal we want gradients for
force = torch.zeros(robot.n_dofs, dtype=gs.tc_float, device=gs.device, requires_grad=True)

for _ in range(100):
    robot.control_dofs_force(force)
    scene.step()

target = torch.tensor([1.0, 0.0, 0.5], dtype=gs.tc_float, device=gs.device)
loss = ((robot.get_state().pos - target) ** 2).sum()

loss.backward()  # propagates through every step back to `force`
print(force.grad)

自动微分的工作方式#

  • 在场景上启用一次即可: 在 gs.options.SimOptions 中设置 requires_grad=True。Genesis World 随后会记录每个时间步反向传播所需的中间子步状态。该标志默认为 False,因此除非你主动启用,仿真就是不可微分的;可微分模式的代价见限制。

  • 查询到的状态是可微分的: 对实体调用 get_state()(robot.get_state().pos),或对求解器调用(scene.rigid_solver.get_state().qpos),返回的状态其张量带有梯度。它们就是普通的 torch.Tensor。其他 getter(get_pos()、get_qpos() 等)返回的值不带梯度。

  • 输入可以是任意 torch 张量: 传给 setter 的控制力、初始位置和目标值,既可以是带 requires_grad=True 创建的叶子张量,也可以是 torch 计算的结果,例如神经网络策略。

  • 仿真是 PyTorch 计算图的一部分: 每次 scene.step() 都是自动微分图中的一个节点,因此 loss.backward() 和 torch.autograd.grad(loss, params) 都能在一次遍历中穿过物理仿真求导,一直追溯到每个输入的来源。

一次前向运行只支持一次反向传播。在下一次前向运行之前,用 scene.reset() 重置场景(例如每个优化步重置一次),并且要在重置之前调用 backward:在 scene.reset() 之后对一次运行求导会抛出错误。scene.backward(loss) 会执行反向传播,然后恢复场景当时所处的物理状态,因此 rollout 可以从那里继续。查询到的状态张量总以环境维度开头,当场景没有并行环境时该维度大小为 1([max(1, n_envs), ...])。

通过动力学训练策略#

输入可以依赖同一次 rollout 中更早查询到的状态,从而闭合策略与仿真之间的回路。梯度会经过每一步传到策略参数,动作惩罚项也可以与状态项共享同一个损失,而无需传入 retain_graph=True。下面的场景使用并行环境构建(scene.build(n_envs=16)),因此动作的每一行对应一个环境:

n_qs, n_dofs = scene.rigid_solver.n_qs, scene.rigid_solver.n_dofs
policy = torch.nn.Linear(n_qs + n_dofs, n_dofs).to(device=gs.device, dtype=gs.tc_float)
optimizer = torch.optim.Adam(policy.parameters(), lr=1e-3)
target_qpos = scene.rigid_solver.get_state().qpos.detach() + 0.1

for _ in range(100):  # optimization iterations
    scene.reset()
    loss = 0.0
    for _ in range(50):  # horizon
        state = scene.rigid_solver.get_state()
        force = policy(torch.cat((state.qpos, state.dofs_vel), dim=-1))
        robot.control_dofs_force(force)
        scene.step()
        loss = loss + 1e-4 * (force**2).sum()
    loss = loss + ((scene.rigid_solver.get_state().qpos - target_qpos) ** 2).sum()
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

如果某个输入在损失所读取的最后一个状态之后才给出,它就无法影响损失,其梯度会保持为 None,这与任何对 PyTorch 损失没有贡献的张量一样。若要阻止梯度经过某个状态回传,像处理普通 PyTorch 张量那样使用 detach()。

限制#

  • 显存随视界增长: 可微分模式会为每个时间步保存中间子步状态,因此长轨迹会按比例占用更多 GPU 显存。通过 gs.options.SimOptions 中的 substeps_local 控制保留多少子步状态;在可微分模式下,它必须能被 substeps 整除。

  • 部分接触与碰撞路径没有梯度: 经过这些路径的梯度可能为零或未定义。需要梯度时,刚体求解器使用 GJK 碰撞路径(见碰撞检测),并且不支持椭圆摩擦锥(见约束模型)。

  • 耦合器支持情况: 默认的传统耦合器可以传递梯度,SAP 耦合器则不能。见耦合。

  • 休眠: 在 requires_grad=True 时不可用;见碰撞检测。

  • 长视界下的数值稳定性: 经过许多步反向传播的梯度可能消失或爆炸,这与任何长程循环计算一样。

延伸阅读#