#
# 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.#
import torch
import torch.autograd.profiler as profiler
from ..cost import DistCost, PoseCost, ProjectedDistCost, JacobianCost, ZeroCost, EEVelCost, StopCost, FiniteDifferenceCost
from ..cost.bound_cost import BoundCost
from ..cost.manipulability_cost import ManipulabilityCost
from ..cost import CollisionCost, VoxelCollisionCost, PrimitiveCollisionCost
from ..model import URDFKinematicModel
from ...util_file import join_path, get_assets_path
from ...differentiable_robot_model.coordinate_transform import matrix_to_quaternion, quaternion_to_matrix
from ...mpc.model.integration_utils import build_fd_matrix
from ...mpc.rollout.rollout_base import RolloutBase
from ..cost.robot_self_collision_cost import RobotSelfCollisionCost
[docs]class ArmBase(RolloutBase):
"""
This rollout function is for reaching a cartesian pose for a robot
Todo:
1. Update exp_params to be kwargs
"""
def __init__(self, exp_params, tensor_args={'device':"cpu", 'dtype':torch.float32}, world_params=None):
self.tensor_args = tensor_args
self.exp_params = exp_params
mppi_params = exp_params['mppi']
model_params = exp_params['model']
robot_params = exp_params['robot_params']
assets_path = get_assets_path()
#print('EE LINK',exp_params['model']['ee_link_name'])
# initialize dynamics model:
dynamics_horizon = mppi_params['horizon'] * model_params['dt']
#Create the dynamical system used for rollouts
self.dynamics_model = URDFKinematicModel(join_path(assets_path,exp_params['model']['urdf_path']),
dt=exp_params['model']['dt'],
batch_size=mppi_params['num_particles'],
horizon=dynamics_horizon,
tensor_args=self.tensor_args,
ee_link_name=exp_params['model']['ee_link_name'],
link_names=exp_params['model']['link_names'],
dt_traj_params=exp_params['model']['dt_traj_params'],
control_space=exp_params['control_space'],
vel_scale=exp_params['model']['vel_scale'])
self.dt = self.dynamics_model.dt
self.n_dofs = self.dynamics_model.n_dofs
# rollout traj_dt starts from dt->dt*(horizon+1) as tstep 0 is the current state
#self.traj_dt = torch.arange(self.dt, (mppi_params['horizon'] + 1) * self.dt, self.dt, device=device, dtype=float_dtype)
self.traj_dt = self.dynamics_model.traj_dt
#print(self.traj_dt)
self.fd_matrix = build_fd_matrix(10 - self.exp_params['cost']['smooth']['order'], device=self.tensor_args['device'], dtype=self.tensor_args['dtype'], PREV_STATE=True, order=self.exp_params['cost']['smooth']['order'])
self.goal_state = None
self.goal_ee_pos = None
self.goal_ee_rot = None
device = self.tensor_args['device']
float_dtype = self.tensor_args['dtype']
self.jacobian_cost = JacobianCost(ndofs=self.n_dofs, device=device,
float_dtype=float_dtype,
retract_weight=exp_params['cost']['retract_weight'])
self.null_cost = ProjectedDistCost(ndofs=self.n_dofs, device=device, float_dtype=float_dtype,
**exp_params['cost']['null_space'])
self.manipulability_cost = ManipulabilityCost(ndofs=self.n_dofs, device=device,
float_dtype=float_dtype,
**exp_params['cost']['manipulability'])
self.zero_vel_cost = ZeroCost(device=device, float_dtype=float_dtype, **exp_params['cost']['zero_vel'])
self.zero_acc_cost = ZeroCost(device=device, float_dtype=float_dtype, **exp_params['cost']['zero_acc'])
self.stop_cost = StopCost(**exp_params['cost']['stop_cost'],
tensor_args=self.tensor_args,
traj_dt=self.traj_dt)
self.stop_cost_acc = StopCost(**exp_params['cost']['stop_cost_acc'],
tensor_args=self.tensor_args,
traj_dt=self.traj_dt)
self.retract_state = torch.tensor([self.exp_params['cost']['retract_state']], device=device, dtype=float_dtype)
# collision model:
# build robot collision model
if self.exp_params['cost']['smooth']['weight'] > 0:
self.smooth_cost = FiniteDifferenceCost(**self.exp_params['cost']['smooth'],
tensor_args=self.tensor_args)
if(self.exp_params['cost']['voxel_collision']['weight'] > 0):
self.voxel_collision_cost = VoxelCollisionCost(robot_params=robot_params,
tensor_args=self.tensor_args,
**self.exp_params['cost']['voxel_collision'])
if(exp_params['cost']['primitive_collision']['weight'] > 0.0):
self.primitive_collision_cost = PrimitiveCollisionCost(world_params=world_params, robot_params=robot_params, tensor_args=self.tensor_args, **self.exp_params['cost']['primitive_collision'])
if(exp_params['cost']['robot_self_collision']['weight'] > 0.0):
self.robot_self_collision_cost = RobotSelfCollisionCost(robot_params=robot_params, tensor_args=self.tensor_args, **self.exp_params['cost']['robot_self_collision'])
self.ee_vel_cost = EEVelCost(ndofs=self.n_dofs,device=device, float_dtype=float_dtype,**exp_params['cost']['ee_vel'])
bounds = torch.cat([self.dynamics_model.state_lower_bounds[:self.n_dofs * 3].unsqueeze(0),self.dynamics_model.state_upper_bounds[:self.n_dofs * 3].unsqueeze(0)], dim=0).T
self.bound_cost = BoundCost(**exp_params['cost']['state_bound'],
tensor_args=self.tensor_args,
bounds=bounds)
self.link_pos_seq = torch.zeros((1, 1, len(self.dynamics_model.link_names), 3), **self.tensor_args)
self.link_rot_seq = torch.zeros((1, 1, len(self.dynamics_model.link_names), 3, 3), **self.tensor_args)
[docs] def cost_fn(self, state_dict, action_batch, no_coll=False, horizon_cost=True):
ee_pos_batch, ee_rot_batch = state_dict['ee_pos_seq'], state_dict['ee_rot_seq']
state_batch = state_dict['state_seq']
lin_jac_batch, ang_jac_batch = state_dict['lin_jac_seq'], state_dict['ang_jac_seq']
link_pos_batch, link_rot_batch = state_dict['link_pos_seq'], state_dict['link_rot_seq']
prev_state = state_dict['prev_state_seq']
prev_state_tstep = state_dict['prev_state_seq'][:,-1]
retract_state = self.retract_state
J_full = torch.cat((lin_jac_batch, ang_jac_batch), dim=-2)
#null-space cost
#if self.exp_params['cost']['null_space']['weight'] > 0:
null_disp_cost = self.null_cost.forward(state_batch[:,:,0:self.n_dofs] -
retract_state[:,0:self.n_dofs],
J_full,
proj_type='identity',
dist_type='squared_l2')
cost = null_disp_cost
if(no_coll == True and horizon_cost == False):
return cost
if(self.exp_params['cost']['manipulability']['weight'] > 0.0):
cost += self.manipulability_cost.forward(J_full)
if(horizon_cost):
if self.exp_params['cost']['stop_cost']['weight'] > 0:
cost += self.stop_cost.forward(state_batch[:, :, self.n_dofs:self.n_dofs * 2])
if self.exp_params['cost']['stop_cost_acc']['weight'] > 0:
cost += self.stop_cost_acc.forward(state_batch[:, :, self.n_dofs*2 :self.n_dofs * 3])
if self.exp_params['cost']['smooth']['weight'] > 0:
order = self.exp_params['cost']['smooth']['order']
prev_dt = (self.fd_matrix @ prev_state_tstep)[-order:]
n_mul = 1
state = state_batch[:,:, self.n_dofs * n_mul:self.n_dofs * (n_mul+1)]
p_state = prev_state[-order:,self.n_dofs * n_mul: self.n_dofs * (n_mul+1)].unsqueeze(0)
p_state = p_state.expand(state.shape[0], -1, -1)
state_buffer = torch.cat((p_state, state), dim=1)
traj_dt = torch.cat((prev_dt, self.traj_dt))
cost += self.smooth_cost.forward(state_buffer, traj_dt)
if self.exp_params['cost']['state_bound']['weight'] > 0:
# compute collision cost:
cost += self.bound_cost.forward(state_batch[:,:,:self.n_dofs * 3])
if self.exp_params['cost']['ee_vel']['weight'] > 0:
cost += self.ee_vel_cost.forward(state_batch, lin_jac_batch)
if(not no_coll):
if self.exp_params['cost']['robot_self_collision']['weight'] > 0:
#coll_cost = self.robot_self_collision_cost.forward(link_pos_batch, link_rot_batch)
coll_cost = self.robot_self_collision_cost.forward(state_batch[:,:,:self.n_dofs])
cost += coll_cost
if self.exp_params['cost']['primitive_collision']['weight'] > 0:
coll_cost = self.primitive_collision_cost.forward(link_pos_batch, link_rot_batch)
cost += coll_cost
if self.exp_params['cost']['voxel_collision']['weight'] > 0:
coll_cost = self.voxel_collision_cost.forward(link_pos_batch, link_rot_batch)
cost += coll_cost
return cost
[docs] def rollout_fn(self, start_state, act_seq):
"""
Return sequence of costs and states encountered
by simulating a batch of action sequences
Parameters
----------
action_seq: torch.Tensor [num_particles, horizon, d_act]
"""
# rollout_start_time = time.time()
#print("computing rollout")
#print(act_seq)
#print('step...')
with profiler.record_function("robot_model"):
state_dict = self.dynamics_model.rollout_open_loop(start_state, act_seq)
#link_pos_seq, link_rot_seq = self.dynamics_model.get_link_poses()
with profiler.record_function("cost_fns"):
cost_seq = self.cost_fn(state_dict, act_seq)
sim_trajs = dict(
actions=act_seq,#.clone(),
costs=cost_seq,#clone(),
ee_pos_seq=state_dict['ee_pos_seq'],#.clone(),
#link_pos_seq=link_pos_seq,
#link_rot_seq=link_rot_seq,
rollout_time=0.0
)
return sim_trajs
[docs] def update_params(self, retract_state=None):
"""
Updates the goal targets for the cost functions.
"""
if(retract_state is not None):
self.retract_state = torch.as_tensor(retract_state, **self.tensor_args).unsqueeze(0)
return True
def __call__(self, start_state, act_seq):
return self.rollout_fn(start_state, act_seq)
[docs] def get_ee_pose(self, current_state):
current_state = current_state.to(**self.tensor_args)
ee_pos_batch, ee_rot_batch, lin_jac_batch, ang_jac_batch = self.dynamics_model.robot_model.compute_fk_and_jacobian(current_state[:,:self.dynamics_model.n_dofs], current_state[:, self.dynamics_model.n_dofs: self.dynamics_model.n_dofs * 2], self.exp_params['model']['ee_link_name'])
ee_quat = matrix_to_quaternion(ee_rot_batch)
state = {'ee_pos_seq':ee_pos_batch, 'ee_rot_seq':ee_rot_batch,
'lin_jac_seq': lin_jac_batch, 'ang_jac_seq': ang_jac_batch,
'ee_quat_seq':ee_quat}
return state
[docs] def current_cost(self, current_state, no_coll=True):
current_state = current_state.to(**self.tensor_args)
curr_batch_size = 1
num_traj_points = 1 #self.dynamics_model.num_traj_points
ee_pos_batch, ee_rot_batch, lin_jac_batch, ang_jac_batch = self.dynamics_model.robot_model.compute_fk_and_jacobian(current_state[:,:self.dynamics_model.n_dofs], current_state[:, self.dynamics_model.n_dofs: self.dynamics_model.n_dofs * 2], self.exp_params['model']['ee_link_name'])
link_pos_seq = self.link_pos_seq
link_rot_seq = self.link_rot_seq
# get link poses:
for ki,k in enumerate(self.dynamics_model.link_names):
link_pos, link_rot = self.dynamics_model.robot_model.get_link_pose(k)
link_pos_seq[:,:,ki,:] = link_pos.view((curr_batch_size, num_traj_points,3))
link_rot_seq[:,:,ki,:,:] = link_rot.view((curr_batch_size, num_traj_points,3,3))
if(len(current_state.shape) == 2):
current_state = current_state.unsqueeze(0)
ee_pos_batch = ee_pos_batch.unsqueeze(0)
ee_rot_batch = ee_rot_batch.unsqueeze(0)
lin_jac_batch = lin_jac_batch.unsqueeze(0)
ang_jac_batch = ang_jac_batch.unsqueeze(0)
state_dict = {'ee_pos_seq':ee_pos_batch, 'ee_rot_seq':ee_rot_batch,
'lin_jac_seq': lin_jac_batch, 'ang_jac_seq': ang_jac_batch,
'state_seq': current_state,'link_pos_seq':link_pos_seq,
'link_rot_seq':link_rot_seq,
'prev_state_seq':current_state}
cost = self.cost_fn(state_dict, None,no_coll=no_coll, horizon_cost=False, return_dist=True)
return cost, state_dict