Differentiate Warp's Simulation with respect to mesh vertices

[lksmac]

Hi,
is it possible to differentiate Warp's simulation with respect to mesh vertices?
Thanks.

[mmacklin]

Hi @lksmac, yes in general most operations like wp.mesh_query_point(), wp.mesh_query_ray() etc support propagating gradients back to the meshes. Can you elaborate on the context some more? e.g.: are you using warp.sim or doing general mesh processing?

[lksmac]

Hi @mmacklin, thank you very much for the quick response.
Yes I am using warp.sim to perform trajectory rollouts of rigid-body mesh objects. Based on some loss I would then like to differentiate back to the mesh vertices.
I would have suspected the mesh vertices to be stored as a wp.array(mesh_vertices, requires_grad=True), but haven't found a representation like that. I am new to Warp (awesome tool btw) and still haven't fully comprehended what's happening to a wp.sim.Mesh object during the builder.finalize(requires_grad=True) step.
Naively I would think of something like optimizer.step([mesh.vertices]) to acquire the gradient after the rollouts. What is the proper way to achieve that?

[christophercrouzet]

Hi @lksmac, if you are using warp.sim, then the vertices are stored under the particle_q field inside of the model/state objects. There's also particle_qd for velocities and particle_f for forces.

[lksmac]

Hi @christophercrouzet, thanks for your reply. If I add a mesh to a rigid-body through the builder, the model.particle_q list stays empty (even after finalizing). In the source code I have only found builder.add_soft_mesh to actively add points to particle_q, but not for builder.add_shape_mesh. To make things more clear I have prepared the following pseudo code:

import warp as wp

# Create mesh
mesh = wp.sim.Mesh(vertices=..., indices=...,)

builder = wp.sim.ModelBuilder()
b = builder.add_body(origin=wp.transform((0.0, 1.0, 5.0), wp.quat_identity()))
builder.add_shape_mesh(body=b, mesh=mesh)

model = builder.finalize(requires_grad=True)

# How can I access the vertices of the mesh now?
print(model.particle_q) # -> returns empty list
print(mesh.mesh.points) # -> returns mesh points

# but ...
optimizer = wp.optim.Adam([mesh.mesh.points], lr=1e-2)
...
optimizer.step([mesh.mesh.points])
print(mesh.mesh.points.grad) # -> returns None

Are the mesh points accessible through the model somehow? (Or is mesh.mesh.points the correct way for a wp.sim.Mesh object?)
Is it possible to obtain a gradient for these mesh points? (Or is it only possible if defined as a soft mesh object?)