Spaces:
Sleeping
Sleeping
| # Copyright (c) 2022 NVIDIA CORPORATION. All rights reserved. | |
| # NVIDIA CORPORATION and its licensors retain all intellectual property | |
| # and proprietary rights in and to this software, related documentation | |
| # and any modifications thereto. Any use, reproduction, disclosure or | |
| # distribution of this software and related documentation without an express | |
| # license agreement from NVIDIA CORPORATION is strictly prohibited. | |
| ########################################################################### | |
| # Example Sim Rigid Kinematics | |
| # | |
| # Tests rigid body forward and backwards kinematics through the | |
| # wp.sim.eval_ik() and wp.sim.eval_fk() methods. Shows how to connect | |
| # gradients from Warp to PyTorch, through custom autograd nodes. | |
| # | |
| ########################################################################### | |
| import os | |
| import numpy as np | |
| import torch | |
| import warp as wp | |
| import warp.sim | |
| import warp.sim.render | |
| wp.init() | |
| class ForwardKinematics(torch.autograd.Function): | |
| def forward(ctx, joint_q, joint_qd, model): | |
| # ensure Torch operations complete before running Warp | |
| wp.synchronize_device() | |
| ctx.tape = wp.Tape() | |
| ctx.model = model | |
| ctx.joint_q = wp.from_torch(joint_q) | |
| ctx.joint_qd = wp.from_torch(joint_qd) | |
| # allocate output | |
| ctx.state = model.state() | |
| with ctx.tape: | |
| wp.sim.eval_fk(model, ctx.joint_q, ctx.joint_qd, None, ctx.state) | |
| # ensure Warp operations complete before returning data to Torch | |
| wp.synchronize_device() | |
| return (wp.to_torch(ctx.state.body_q), wp.to_torch(ctx.state.body_qd)) | |
| def backward(ctx, adj_body_q, adj_body_qd): | |
| # ensure Torch operations complete before running Warp | |
| wp.synchronize_device() | |
| # map incoming Torch grads to our output variables | |
| ctx.state.body_q.grad = wp.from_torch(adj_body_q, dtype=wp.transform) | |
| ctx.state.body_qd.grad = wp.from_torch(adj_body_qd, dtype=wp.spatial_vector) | |
| ctx.tape.backward() | |
| # ensure Warp operations complete before returning data to Torch | |
| wp.synchronize_device() | |
| # return adjoint w.r.t. inputs | |
| return (wp.to_torch(ctx.tape.gradients[ctx.joint_q]), wp.to_torch(ctx.tape.gradients[ctx.joint_qd]), None) | |
| class Example: | |
| def __init__(self, stage, device=None, verbose=False): | |
| self.verbose = verbose | |
| self.frame_dt = 1.0 / 60.0 | |
| self.render_time = 0.0 | |
| builder = wp.sim.ModelBuilder() | |
| builder.add_articulation() | |
| chain_length = 4 | |
| chain_width = 1.0 | |
| for i in range(chain_length): | |
| if i == 0: | |
| parent = -1 | |
| parent_joint_xform = wp.transform([0.0, 0.0, 0.0], wp.quat_identity()) | |
| else: | |
| parent = builder.joint_count - 1 | |
| parent_joint_xform = wp.transform([chain_width, 0.0, 0.0], wp.quat_identity()) | |
| # create body | |
| b = builder.add_body(origin=wp.transform([i, 0.0, 0.0], wp.quat_identity()), armature=0.1) | |
| builder.add_joint_revolute( | |
| parent=parent, | |
| child=b, | |
| axis=wp.vec3(0.0, 0.0, 1.0), | |
| parent_xform=parent_joint_xform, | |
| child_xform=wp.transform_identity(), | |
| limit_lower=-np.deg2rad(60.0), | |
| limit_upper=np.deg2rad(60.0), | |
| target_ke=0.0, | |
| target_kd=0.0, | |
| limit_ke=30.0, | |
| limit_kd=30.0, | |
| ) | |
| if i == chain_length - 1: | |
| # create end effector | |
| builder.add_shape_sphere(pos=wp.vec3(0.0, 0.0, 0.0), radius=0.1, density=10.0, body=b) | |
| else: | |
| # create shape | |
| builder.add_shape_box( | |
| pos=wp.vec3(chain_width * 0.5, 0.0, 0.0), hx=chain_width * 0.5, hy=0.1, hz=0.1, density=10.0, body=b | |
| ) | |
| # finalize model | |
| self.model = builder.finalize(device) | |
| self.model.ground = False | |
| self.torch_device = wp.device_to_torch(self.model.device) | |
| self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=50.0) | |
| self.target = torch.from_numpy(np.array((2.0, 1.0, 0.0))).to(self.torch_device) | |
| self.body_q = None | |
| self.body_qd = None | |
| # optimization variable | |
| self.joint_q = torch.zeros(len(self.model.joint_q), requires_grad=True, device=self.torch_device) | |
| self.joint_qd = torch.zeros(len(self.model.joint_qd), requires_grad=True, device=self.torch_device) | |
| self.train_rate = 0.01 | |
| def update(self): | |
| (self.body_q, self.body_qd) = ForwardKinematics.apply(self.joint_q, self.joint_qd, self.model) | |
| l = torch.norm(self.body_q[self.model.body_count - 1][0:3] - self.target) ** 2.0 | |
| l.backward() | |
| if self.verbose: | |
| print(l) | |
| print(self.joint_q.grad) | |
| with torch.no_grad(): | |
| self.joint_q -= self.joint_q.grad * self.train_rate | |
| self.joint_q.grad.zero_() | |
| def render(self): | |
| s = self.model.state() | |
| s.body_q = wp.from_torch(self.body_q, dtype=wp.transform, requires_grad=False) | |
| s.body_qd = wp.from_torch(self.body_qd, dtype=wp.spatial_vector, requires_grad=False) | |
| self.renderer.begin_frame(self.render_time) | |
| self.renderer.render(s) | |
| self.renderer.render_sphere(name="target", pos=self.target, rot=wp.quat_identity(), radius=0.1) | |
| self.renderer.end_frame() | |
| self.render_time += self.frame_dt | |
| if __name__ == "__main__": | |
| stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_sim_fk_grad.usd") | |
| example = Example(stage_path, device=wp.get_preferred_device(), verbose=True) | |
| train_iters = 512 | |
| for _ in range(train_iters): | |
| example.update() | |
| example.render() | |
| example.renderer.save() | |