108 lines
3.7 KiB
Python
108 lines
3.7 KiB
Python
#
|
|
# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
#
|
|
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
|
|
# property and proprietary rights in and to this material, related
|
|
# documentation and any modifications thereto. Any use, reproduction,
|
|
# disclosure or distribution of this material and related documentation
|
|
# without an express license agreement from NVIDIA CORPORATION or
|
|
# its affiliates is strictly prohibited.
|
|
#
|
|
|
|
# Standard Library
|
|
from dataclasses import dataclass
|
|
from itertools import product
|
|
|
|
# Third Party
|
|
import torch
|
|
|
|
# Local Folder
|
|
from .cost_base import CostBase, CostConfig
|
|
|
|
|
|
@dataclass
|
|
class ManipulabilityCostConfig(CostConfig):
|
|
use_joint_limits: bool = False
|
|
|
|
def __post_init__(self):
|
|
return super().__post_init__()
|
|
|
|
|
|
class ManipulabilityCost(CostBase, ManipulabilityCostConfig):
|
|
def __init__(self, config: ManipulabilityCostConfig):
|
|
ManipulabilityCostConfig.__init__(self, **vars(config))
|
|
CostBase.__init__(self)
|
|
self.i_mat = torch.ones(
|
|
(6, 1), device=self.tensor_args.device, dtype=self.tensor_args.dtype
|
|
)
|
|
|
|
self.delta_vector = torch.zeros(
|
|
(64, 1, 1, 6, 1), device=self.tensor_args.device, dtype=self.tensor_args.dtype
|
|
)
|
|
x = [i for i in product(range(2), repeat=6)]
|
|
self.delta_vector[:, 0, 0, :, 0] = torch.as_tensor(
|
|
x, device=self.tensor_args.device, dtype=self.tensor_args.dtype
|
|
)
|
|
self.delta_vector[self.delta_vector == 0] = -1.0
|
|
|
|
if self.cost_fn is None:
|
|
if self.use_joint_limits and self.joint_limits is not None:
|
|
self.cost_fn = self.joint_limited_manipulability_delta
|
|
else:
|
|
self.cost_fn = self.manipulability
|
|
|
|
def forward(self, jac_batch, q, qdot):
|
|
b, h, n = q.shape
|
|
if self.use_nn:
|
|
q = q.view(b * h, n)
|
|
score = self.cost_fn(q, jac_batch, qdot)
|
|
if self.use_nn:
|
|
score = score.view(b, h)
|
|
score[score > self.hinge_value] = self.hinge_value
|
|
score = (self.hinge_value / score) - 1
|
|
cost = self.weight * score
|
|
|
|
return cost
|
|
|
|
def manipulability(self, q, jac_batch, qdot=None):
|
|
with torch.cuda.amp.autocast(enabled=False):
|
|
J_J_t = torch.matmul(jac_batch, jac_batch.transpose(-2, -1))
|
|
score = torch.sqrt(torch.det(J_J_t))
|
|
|
|
score[score != score] = 0.0
|
|
|
|
return score
|
|
|
|
def joint_limited_manipulability_delta(self, q, jac_batch, qdot=None):
|
|
# q is [b,h,dof]
|
|
q_low = q - self.joint_limits[:, 0]
|
|
q_high = q - self.joint_limits[:, 1]
|
|
|
|
d_h_1 = torch.square(self.joint_limits[:, 1] - self.joint_limits[:, 0]) * (q_low + q_high)
|
|
d_h_2 = 4.0 * (torch.square(q_low) * torch.square(q_high))
|
|
d_h = torch.div(d_h_1, d_h_2)
|
|
|
|
dh_term = 1.0 / torch.sqrt(1 + torch.abs(d_h))
|
|
f_ten = torch.tensor(1.0, **self.tensor_args)
|
|
q_low = torch.abs(q_low)
|
|
q_high = torch.abs(q_high)
|
|
p_plus = torch.where(q_low > q_high, dh_term, f_ten).unsqueeze(-2)
|
|
p_minus = torch.where(q_low > q_high, f_ten, dh_term).unsqueeze(-2)
|
|
|
|
j_sign = torch.sign(jac_batch)
|
|
l_delta = torch.sign(self.delta_vector) * j_sign
|
|
|
|
L = torch.where(l_delta < 0.0, p_minus, p_plus)
|
|
|
|
with torch.cuda.amp.autocast(enabled=False):
|
|
w_J = L * jac_batch
|
|
J_J_t = torch.matmul(w_J, w_J.transpose(-2, -1))
|
|
score = torch.sqrt(torch.det(J_J_t))
|
|
|
|
# get actual score:
|
|
min_score = torch.min(score, dim=0)[0]
|
|
max_score = torch.max(score, dim=0)[0]
|
|
score = min_score / max_score
|
|
score[score != score] = 0.0
|
|
return score
|