[FEATURE] WIP: Make rigid body simulation fully differentiable

examples/differentiable_rigid.py

@@ -0,0 +1,92 @@
+import torch
+import genesis as gs
+
+show_viewer = False
+
+gs.init(precision="32", logging_level="info")
+
+dt = 1e-2
+horizon = 100
+substeps = 1
+goal_pos = gs.tensor([0.7, 1.0, 0.05])
+goal_quat = gs.tensor([0.3, 0.2, 0.1, 0.9])
+goal_quat = goal_quat / torch.norm(goal_quat, dim=-1, keepdim=True)
+
+scene = gs.Scene(
+    sim_options=gs.options.SimOptions(dt=dt, substeps=substeps, requires_grad=True, gravity=(0, 0, -1)),
+    rigid_options=gs.options.RigidOptions(
+        enable_collision=False,
+        enable_self_collision=False,
+        enable_joint_limit=False,
+        disable_constraint=True,
+        use_contact_island=False,
+        use_hibernation=False,
+    ),
+    viewer_options=gs.options.ViewerOptions(
+        camera_pos=(2.5, -0.15, 2.42),
+        camera_lookat=(0.5, 0.5, 0.1),
+    ),
+    show_viewer=show_viewer,
+)
+
+box = scene.add_entity(
+    gs.morphs.Box(
+        pos=(0, 0, 0),
+        size=(0.1, 0.1, 0.2),
+    ),
+    surface=gs.surfaces.Default(
+        color=(0.9, 0.0, 0.0, 1.0),
+    ),
+)
+if show_viewer:
+    target = scene.add_entity(
+        gs.morphs.Box(
+            pos=goal_pos,
+            quat=goal_quat,
+            size=(0.1, 0.1, 0.2),
+        ),
+        surface=gs.surfaces.Default(
+            color=(0.0, 0.9, 0.0, 0.5),
+        ),
+    )
+
+scene.build()
+
+num_iter = 200
+lr = 1e-2
+
+init_pos = gs.tensor([0.3, 0.1, 0.28], requires_grad=True)
+init_quat = gs.tensor([1.0, 0.0, 0.0, 0.0], requires_grad=True)
+optimizer = torch.optim.Adam([init_pos, init_quat], lr=lr)
+
+scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_iter, eta_min=1e-3)
+
+for iter in range(num_iter):
+    scene.reset()
+
+    box.set_pos(init_pos)
+    box.set_quat(init_quat)
+
+    loss = 0
+    for i in range(horizon):
+        scene.step()
+        if show_viewer:
+            target.set_pos(goal_pos)
+            target.set_quat(goal_quat)
+
+    box_state = box.get_state()
+    box_pos = box_state.pos
+    box_quat = box_state.quat
+    loss = torch.abs(box_pos - goal_pos).sum() + torch.abs(box_quat - goal_quat).sum()
+
+    optimizer.zero_grad()
+    loss.backward()  # this lets gradient flow all the way back to tensor input
+    optimizer.step()
+    scheduler.step()
+
+    with torch.no_grad():
+        init_quat.data = init_quat / torch.norm(init_quat, dim=-1, keepdim=True)
+
+    print("loss: ", loss.item())
+
+# assert_allclose(loss, 0.0, atol=1e-2)

examples/differentiable_rigid_demo_1.py

@@ -0,0 +1,92 @@
+import torch
+import genesis as gs
+
+show_viewer = False
+
+gs.init(precision="32", logging_level="info")
+
+dt = 1e-2
+horizon = 100
+substeps = 1
+goal_pos = gs.tensor([0.7, 1.0, 0.05])
+goal_quat = gs.tensor([0.3, 0.2, 0.1, 0.9])
+goal_quat = goal_quat / torch.norm(goal_quat, dim=-1, keepdim=True)
+
+scene = gs.Scene(
+    sim_options=gs.options.SimOptions(dt=dt, substeps=substeps, requires_grad=True, gravity=(0, 0, -1)),
+    rigid_options=gs.options.RigidOptions(
+        enable_collision=False,
+        enable_self_collision=False,
+        enable_joint_limit=False,
+        disable_constraint=True,
+        use_contact_island=False,
+        use_hibernation=False,
+    ),
+    viewer_options=gs.options.ViewerOptions(
+        camera_pos=(2.5, -0.15, 2.42),
+        camera_lookat=(0.5, 0.5, 0.1),
+    ),
+    show_viewer=show_viewer,
+)
+
+box = scene.add_entity(
+    gs.morphs.Box(
+        pos=(0, 0, 0),
+        size=(0.1, 0.1, 0.2),
+    ),
+    surface=gs.surfaces.Default(
+        color=(0.9, 0.0, 0.0, 1.0),
+    ),
+)
+if show_viewer:
+    target = scene.add_entity(
+        gs.morphs.Box(
+            pos=goal_pos,
+            quat=goal_quat,
+            size=(0.1, 0.1, 0.2),
+        ),
+        surface=gs.surfaces.Default(
+            color=(0.0, 0.9, 0.0, 0.5),
+        ),
+    )
+
+scene.build()
+
+num_iter = 200
+lr = 1e-2
+
+init_pos = gs.tensor([0.3, 0.1, 0.28], requires_grad=True)
+init_quat = gs.tensor([1.0, 0.0, 0.0, 0.0], requires_grad=True)
+optimizer = torch.optim.Adam([init_pos, init_quat], lr=lr)
+
+scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_iter, eta_min=1e-3)
+
+for iter in range(num_iter):
+    scene.reset()
+
+    box.set_pos(init_pos)
+    box.set_quat(init_quat)
+
+    loss = 0
+    for i in range(horizon):
+        scene.step()
+        if show_viewer:
+            target.set_pos(goal_pos)
+            target.set_quat(goal_quat)
+
+    box_state = box.get_state()
+    box_pos = box_state.pos
+    box_quat = box_state.quat
+    loss = torch.abs(box_pos - goal_pos).sum() + torch.abs(box_quat - goal_quat).sum()
+
+    optimizer.zero_grad()
+    loss.backward()  # this lets gradient flow all the way back to tensor input
+    optimizer.step()
+    scheduler.step()
+
+    with torch.no_grad():
+        init_quat.data = init_quat / torch.norm(init_quat, dim=-1, keepdim=True)
+
+    print("loss: ", loss.item())
+
+# assert_allclose(loss, 0.0, atol=1e-2)

examples/diffrigid/1_one_step.py

@@ -0,0 +1,96 @@
+"""
+One step optimization for the basic debugging of differentiable rigid simulation.
+"""
+import torch
+import genesis as gs
+import matplotlib.pyplot as plt
+
+show_viewer = False
+
+gs.init(precision="32", logging_level="warn", backend=gs.cpu)
+
+dt = 1e-2
+horizon = 1
+substeps = 1
+goal_pos = gs.tensor([0.0, 0.0, 1.0])
+
+scene = gs.Scene(
+    sim_options=gs.options.SimOptions(dt=dt, substeps=substeps, requires_grad=True),
+    rigid_options=gs.options.RigidOptions(
+        use_gjk_collision=True,
+        enable_joint_limit=False,
+    ),
+    show_viewer=show_viewer,
+)
+
+ball = scene.add_entity(
+    gs.morphs.Sphere(
+        pos=(0, 0, 0.109),  # small penetration with ground
+        radius=0.1,
+    ),
+    surface=gs.surfaces.Default(
+        color=(0.9, 0.0, 0.0, 1.0),
+    ),
+)
+
+ground = scene.add_entity(
+    gs.morphs.Box(
+        pos=(0, 0, 0),
+        size=(5.0, 5.0, 0.02),
+        #fixed=True,
+    ),
+    surface=gs.surfaces.Default(
+        color=(0.0, 0.0, 0.9, 1.0),
+    ),
+    material=gs.materials.Rigid(
+        gravity_compensation=1,
+    )
+)
+
+scene.build()
+
+num_iter = 400
+lr = 1e-4
+
+init_pos = gs.tensor([0.0, 0.0, 0.0], requires_grad=True)
+#init_quat = gs.tensor([1.0, 0.0, 0.0, 0.0], requires_grad=True)
+optimizer = torch.optim.Adam([init_pos], lr=lr)
+
+scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_iter, eta_min=1e-3)
+
+prev_loss = float('inf')
+losses = []
+for iter in range(num_iter):
+    scene.reset()
+
+    ground.set_pos(init_pos)
+ #   ground.set_quat(init_quat)
+
+    for i in range(horizon):
+        scene.step()
+
+    ball_state = ball.get_state()
+    ball_pos = ball_state.pos
+    loss = torch.abs(ball_pos - goal_pos).sum()
+
+    optimizer.zero_grad()
+    loss.backward()  # this lets gradient flow all the way back to tensor input
+    optimizer.step()
+    scheduler.step()
+
+    grad_norm = torch.nn.utils.clip_grad_norm_(init_pos.grad, 1.0)
+
+    # with torch.no_grad():
+    #     init_quat.data = init_quat / torch.norm(init_quat, dim=-1, keepdim=True)
+    # with torch.no_grad():
+    #     init_pos.data[0] = 0.0
+    #     init_pos.data[1] = 0.0
+
+    print(f"Loss: {prev_loss:.6g} -> {loss.item():.6g}")
+    prev_loss = loss.item()
+
+    losses.append(loss.item())
+
+    plt.plot(losses)
+    plt.savefig("loss.png")
+    plt.close()

examples/diffrigid/2_lift_ball.py

@@ -0,0 +1,106 @@
+"""
+One step optimization for the basic debugging of differentiable rigid simulation.
+"""
+import torch
+import genesis as gs
+import matplotlib.pyplot as plt
+import matplotlib
+matplotlib.use("Agg")
+
+show_viewer = False
+
+gs.init(precision="32", logging_level="warn", backend=gs.cpu)
+
+dt = 1e-2
+horizon = 10
+substeps = 1
+grad_window = 5 #None
+
+scene = gs.Scene(
+    sim_options=gs.options.SimOptions(dt=dt, substeps=substeps, requires_grad=True, grad_window_steps=grad_window),
+    rigid_options=gs.options.RigidOptions(
+        use_gjk_collision=True,
+        enable_joint_limit=False,
+    ),
+    show_viewer=show_viewer,
+)
+
+ball = scene.add_entity(
+    gs.morphs.Sphere(
+        pos=(0, 0, 0.109),  # small penetration with ground
+        radius=0.1,
+    ),
+    surface=gs.surfaces.Default(
+        color=(0.9, 0.0, 0.0, 1.0),
+    ),
+)
+
+ground = scene.add_entity(
+    gs.morphs.Box(
+        pos=(0, 0, 0),
+        size=(5.0, 5.0, 0.02),
+        #fixed=True,
+    ),
+    surface=gs.surfaces.Default(
+        color=(0.0, 0.0, 0.9, 1.0),
+    ),
+    material=gs.materials.Rigid(
+        gravity_compensation=1,
+    )
+)
+cam = scene.add_camera(
+    pos=(3.5, 0.5, 2.5),
+    lookat=(0.0, 0.0, 0.5),
+    fov=40,
+    GUI=False,
+)
+
+scene.build()
+
+num_iter = 10000
+lr = 1e-2
+
+force = gs.zeros((horizon, 6), requires_grad=True)
+optimizer = torch.optim.Adam([force], lr=lr)
+
+render_every = 100
+prev_loss = float('inf')
+losses = []
+for iter in range(num_iter):
+    scene.reset()
+
+    curr_losses = []
+    if iter % render_every == 0: 
+        cam.start_recording()   
+    for i in range(horizon):
+        curr_force = force[i]
+        ball.control_dofs_force(curr_force)
+        scene.step()
+
+        ball_state = ball.get_state()
+        ball_pos = ball_state.pos
+        curr_loss = -ball_pos[:, 2].sum() # make x, y, z larger
+        curr_losses.append(curr_loss)
+
+        if iter % render_every == 0: 
+            cam.render()
+    if iter % render_every == 0: 
+        cam.stop_recording(save_to_filename=f"video_{iter:06d}.mp4", fps=30)
+
+    loss = sum(curr_losses) / len(curr_losses)
+    optimizer.zero_grad()
+    loss.backward()  # this lets gradient flow all the way back to tensor input
+    optimizer.step()
+    grad_norm = torch.nn.utils.clip_grad_norm_(force.grad, 1.0)
+
+    with torch.no_grad():
+        force.data[:, 6:] = 0.0
+
+    print(f"Loss: {prev_loss:.6g} -> {loss.item():.6g} | Grad Norm: {grad_norm:.6g} | Force: {force.data.mean(dim=0).cpu().numpy().tolist()}")
+    prev_loss = loss.item()
+
+    losses.append(loss.item())
+
+    plt.plot(losses)
+    plt.savefig("loss.png")
+    plt.close()