#
# MIT License
#
# Copyright (c) 2020-2021 NVIDIA CORPORATION.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.
#
# **********************************************************************
# The first version was licensed as "Original Source License"(see below).
# Several enhancements and bug fixes were done at NVIDIA CORPORATION
# since obtaining the first version.
#
#
#
# Original Source License:
#
# MIT License
#
# Copyright (c) Facebook, Inc. and its affiliates.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.#
# Copyright (c) Facebook, Inc. and its affiliates.
import torch
from . import utils
from .coordinate_transform import (
CoordinateTransform,
z_rot,
y_rot,
x_rot,
)
import hydra
[docs]class DifferentiableRigidBody(torch.nn.Module):
"""
Differentiable Representation of a link
"""
def __init__(self, rigid_body_params, tensor_args={'device':"cpu", 'dtype':torch.float32}):
super().__init__()
self.tensor_args = tensor_args
self.device = tensor_args['device']
self.joint_id = rigid_body_params["joint_id"]
self.name = rigid_body_params["link_name"]
# dynamics parameters
self.mass = rigid_body_params["mass"]
self.com = rigid_body_params["com"]
self.inertia_mat = rigid_body_params["inertia_mat"]
self.joint_damping = rigid_body_params["joint_damping"]
# kinematics parameters
self.trans = rigid_body_params["trans"]
self.rot_angles = rigid_body_params["rot_angles"].unsqueeze(0)
roll = self.rot_angles[:,0]
pitch = self.rot_angles[:,1]
yaw = self.rot_angles[:,2]
self.fixed_rotation = (z_rot(yaw) @ y_rot(pitch)) @ x_rot(roll)
self.joint_limits = rigid_body_params["joint_limits"]
# local joint axis (w.r.t. joint coordinate frame):
self.joint_axis = rigid_body_params["joint_axis"]
self.axis_idx = torch.nonzero(self.joint_axis.squeeze(0))
if self.axis_idx.nelement() > 0:
self.axis_idx = self.axis_idx[0]
if self.joint_axis[0, 0] == 1:
self.axis_rot_fn = x_rot
elif self.joint_axis[0, 1] == 1:
self.axis_rot_fn = y_rot
else:
self.axis_rot_fn = z_rot
self.joint_pose = CoordinateTransform(tensor_args=tensor_args) #.to(device)
self.joint_pose.set_translation(torch.reshape(self.trans, (1, 3)))
self._batch_size = -1
self._batch_trans = None
self._batch_rot = None
# local velocities and accelerations (w.r.t. joint coordinate frame):
# in spatial vector terminology: linear velocity v
self.joint_lin_vel = torch.zeros((1, 3), **self.tensor_args)
# in spatial vector terminology: angular velocity w
self.joint_ang_vel = torch.zeros((1, 3), **self.tensor_args)
# in spatial vector terminology: linear acceleration vd
self.joint_lin_acc = torch.zeros((1, 3), **self.tensor_args)
# in spatial vector terminology: angular acceleration wd
self.joint_ang_acc = torch.zeros((1, 3), **self.tensor_args)
self.update_joint_state(torch.zeros((1, 1), **self.tensor_args), torch.zeros((1, 1), **self.tensor_args))
# self.update_joint_acc(torch.zeros(1, 1).to(self.device, dtype=self.float_dtype))
self.update_joint_acc(torch.zeros((1, 1), **self.tensor_args))
self.pose = CoordinateTransform(tensor_args=self.tensor_args)
# I have different vectors for angular/linear motion/force, but they usually always appear as a pair
# meaning we usually always compute both angular/linear components.
# Maybe worthwile thinking of a structure for this - in math notation we would use the notion of spatial vectors
# drake uses some form of spatial vector implementation
self.lin_vel = torch.zeros((1, 3), **self.tensor_args)
self.ang_vel = torch.zeros((1, 3), **self.tensor_args)
self.lin_acc = torch.zeros((1, 3), **self.tensor_args)
self.ang_acc = torch.zeros((1, 3), **self.tensor_args)
# in spatial vector terminology this is the "linear force f"
self.lin_force = torch.zeros((1, 3), **self.tensor_args)
# in spatial vector terminology this is the "couple n"
self.ang_force = torch.zeros((1, 3), **self.tensor_args)
return
[docs] def update_joint_state(self, q, qd):
batch_size = q.shape[0]
self.joint_ang_vel = qd @ self.joint_axis
if(batch_size != self._batch_size):
#print('calling once', self._batch_size, batch_size)
self._batch_size = batch_size
self._batch_trans = self.trans.unsqueeze(0).repeat(self._batch_size,1)
self._batch_rot = self.fixed_rotation.repeat(self._batch_size, 1, 1)
# when we update the joint angle, we also need to update the transformation
self.joint_pose.set_translation(self._batch_trans)#
#self.joint_pose.set_translation(torch.reshape(self.trans, (1, 3)))
#print(q.shape)
rot = self.axis_rot_fn(q.squeeze(1))
self.joint_pose.set_rotation(self._batch_rot @ rot)
return
[docs] def update_joint_acc(self, qdd):
# local z axis (w.r.t. joint coordinate frame):
self.joint_ang_acc = qdd @ self.joint_axis
return
[docs] def multiply_inertia_with_motion_vec(self, lin, ang):
mass, com, inertia_mat = self._get_dynamics_parameters_values()
mcom = com * mass
com_skew_symm_mat = utils.vector3_to_skew_symm_matrix(com)
inertia = inertia_mat + mass * (
com_skew_symm_mat @ com_skew_symm_mat.transpose(-2, -1)
)
batch_size = lin.shape[0]
new_lin_force = mass * lin - utils.cross_product(
mcom.repeat(batch_size, 1), ang
)
new_ang_force = (inertia.repeat(batch_size, 1, 1) @ ang.unsqueeze(2)).squeeze(
2
) + utils.cross_product(mcom.repeat(batch_size, 1), lin)
return new_lin_force, new_ang_force
[docs] def _get_dynamics_parameters_values(self):
return self.mass, self.com, self.inertia_mat
[docs] def get_joint_limits(self):
return self.joint_limits
[docs] def get_joint_damping_const(self):
return self.joint_damping
[docs]class LearnableRigidBody(DifferentiableRigidBody):
r"""
Learnable Representation of a link
"""
def __init__(self, learnable_rigid_body_config, gt_rigid_body_params, device="cpu", float_dtype=torch.float32):
super().__init__(rigid_body_params=gt_rigid_body_params, device=device)
# we overwrite dynamics parameters
if "mass" in learnable_rigid_body_config.learnable_dynamics_params:
self.mass_fn = hydra.utils.instantiate(
learnable_rigid_body_config.mass_parametrization, device=device
)
else:
self.mass_fn = lambda: self.mass
if "com" in learnable_rigid_body_config.learnable_dynamics_params:
self.com_fn = hydra.utils.instantiate(
learnable_rigid_body_config.com_parametrization, device=device
)
else:
self.com_fn = lambda: self.com
if "inertia_mat" in learnable_rigid_body_config.learnable_dynamics_params:
self.inertia_mat_fn = hydra.utils.instantiate(learnable_rigid_body_config.inertia_parametrization)
else:
self.inertia_mat_fn = lambda: self.inertia_mat
self.joint_damping = gt_rigid_body_params["joint_damping"]
# kinematics parameters
if "trans" in learnable_rigid_body_config.learnable_kinematics_params:
self.trans = torch.nn.Parameter(
torch.rand_like(gt_rigid_body_params["trans"])
)
self.joint_pose.set_translation(torch.reshape(self.trans, (1, 3)))
if "rot_angles" in learnable_rigid_body_config.learnable_kinematics_params:
self.rot_angles = torch.nn.Parameter(gt_rigid_body_params["rot_angles"])
return
[docs] def _get_dynamics_parameters_values(self):
return self.mass_fn(), self.com_fn(), self.inertia_mat_fn()