569 lines
20 KiB
Python
569 lines
20 KiB
Python
|
import abc
|
||
|
|
import time
|
||
|
|
|
||
|
|
import einops
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
from torch import nn
|
||
|
|
|
||
|
|
from mp_baselines.planners.costs.cost_functions import CostGPTrajectory
|
||
|
|
from mp_baselines.planners.costs.factors.mp_priors_multi import MultiMPPrior
|
||
|
|
from torch_robotics.torch_planning_objectives.fields.distance_fields import interpolate_points_v1
|
||
|
|
from torch_robotics.torch_utils.torch_utils import to_torch
|
||
|
|
|
||
|
|
|
||
|
|
class GuideManagerTrajectories(nn.Module):
|
||
|
|
|
||
|
|
def __init__(self, dataset, cost, clip_grad=False, clip_grad_rule='norm', max_grad_norm=1., max_grad_value=0.1,
|
||
|
|
interpolate_trajectories_for_collision=False,
|
||
|
|
num_interpolated_points_for_collision=128,
|
||
|
|
use_velocity_from_finite_difference=False,
|
||
|
|
start_state_pos=None,
|
||
|
|
goal_state_pos=None,
|
||
|
|
num_steps=100,
|
||
|
|
robot=None,
|
||
|
|
n_samples=1,
|
||
|
|
tensor_args=None,
|
||
|
|
**kwargs):
|
||
|
|
super().__init__()
|
||
|
|
self.cost = cost
|
||
|
|
self.dataset = dataset
|
||
|
|
|
||
|
|
self.interpolate_trajectories_for_collision = interpolate_trajectories_for_collision
|
||
|
|
self.num_interpolated_points_for_collision = num_interpolated_points_for_collision
|
||
|
|
|
||
|
|
self.clip_grad = clip_grad
|
||
|
|
self.clip_grad_rule = clip_grad_rule
|
||
|
|
self.max_grad_norm = max_grad_norm
|
||
|
|
self.max_grad_value = max_grad_value
|
||
|
|
|
||
|
|
# velocity
|
||
|
|
self.use_velocity_from_finite_difference = use_velocity_from_finite_difference
|
||
|
|
self.robot = robot
|
||
|
|
self.start_state_pos = start_state_pos
|
||
|
|
self.goal_state_pos = goal_state_pos
|
||
|
|
# initialize velocity trajectory with a constant velocity
|
||
|
|
self.velocity = self.robot.get_velocity(
|
||
|
|
MultiMPPrior.const_vel_trajectory(
|
||
|
|
start_state_pos,
|
||
|
|
goal_state_pos,
|
||
|
|
robot.dt,
|
||
|
|
num_steps,
|
||
|
|
self.robot.q_dim,
|
||
|
|
set_initial_final_vel_to_zero=True,
|
||
|
|
tensor_args=tensor_args)
|
||
|
|
)
|
||
|
|
self.velocity = to_torch(self.velocity, **tensor_args)
|
||
|
|
self.velocity = einops.repeat(self.velocity, "H D -> B H D", B=n_samples)
|
||
|
|
|
||
|
|
def forward(self, x_pos_normalized):
|
||
|
|
x_pos = x_pos_normalized.clone()
|
||
|
|
with torch.enable_grad():
|
||
|
|
x_pos.requires_grad_(True)
|
||
|
|
self.velocity.requires_grad_(True)
|
||
|
|
|
||
|
|
# unnormalize x
|
||
|
|
# x is normalized, but the guides are defined on unnormalized trajectory space
|
||
|
|
x_pos = self.dataset.unnormalize_trajectories(x_pos)
|
||
|
|
|
||
|
|
if self.interpolate_trajectories_for_collision:
|
||
|
|
# finer interpolation of trajectory for better collision avoidance
|
||
|
|
x_interpolated = interpolate_points_v1(x_pos, num_interpolated_points=self.num_interpolated_points_for_collision)
|
||
|
|
else:
|
||
|
|
x_interpolated = x_pos
|
||
|
|
|
||
|
|
# compute costs
|
||
|
|
# append the current velocity trajectory to the position trajectory only for non-interpolated trajectories
|
||
|
|
if self.use_velocity_from_finite_difference:
|
||
|
|
x_vel = self.robot.get_velocity(x_pos)
|
||
|
|
x_pos_vel = torch.cat((x_pos, x_vel), dim=-1)
|
||
|
|
else:
|
||
|
|
x_pos_vel = torch.cat((x_pos, self.velocity), dim=-1)
|
||
|
|
|
||
|
|
cost_l, weight_grad_cost_l = self.cost(x_pos_vel, x_interpolated=x_interpolated, return_invidual_costs_and_weights=True)
|
||
|
|
grad = 0
|
||
|
|
grad_velocity = 0
|
||
|
|
for cost, weight_grad_cost in zip(cost_l, weight_grad_cost_l):
|
||
|
|
if torch.is_tensor(cost):
|
||
|
|
# y.sum() is a surrogate to compute gradients of independent quantities over the batch dimension
|
||
|
|
# x are the support points. Compute gradients wrt x, not x_interpolated
|
||
|
|
if self.use_velocity_from_finite_difference:
|
||
|
|
grad_cost = torch.autograd.grad([cost.sum()], [x_pos], retain_graph=True)[0]
|
||
|
|
else:
|
||
|
|
grad_cost, grad_cost_velocity = torch.autograd.grad([cost.sum()], [x_pos, self.velocity])
|
||
|
|
|
||
|
|
# clip gradients
|
||
|
|
grad_cost_clipped = self.clip_gradient(grad_cost)
|
||
|
|
if not self.use_velocity_from_finite_difference:
|
||
|
|
grad_cost_velocity_clipped = self.clip_gradient(grad_cost_velocity)
|
||
|
|
|
||
|
|
# zeroing gradients at start and goal
|
||
|
|
grad_cost_clipped[..., 0, :] = 0.
|
||
|
|
grad_cost_clipped[..., -1, :] = 0.
|
||
|
|
if not self.use_velocity_from_finite_difference:
|
||
|
|
grad_cost_velocity_clipped[..., 0, :] = 0.
|
||
|
|
grad_cost_velocity_clipped[..., -1, :] = 0.
|
||
|
|
|
||
|
|
# combine gradients
|
||
|
|
grad_cost_clipped_weighted = weight_grad_cost * grad_cost_clipped
|
||
|
|
grad += grad_cost_clipped_weighted
|
||
|
|
|
||
|
|
if not self.use_velocity_from_finite_difference:
|
||
|
|
grad_cost_velocity_clipped_weighted = weight_grad_cost * grad_cost_velocity_clipped
|
||
|
|
grad_velocity += grad_cost_velocity_clipped_weighted
|
||
|
|
|
||
|
|
# Update the velocity
|
||
|
|
if not self.use_velocity_from_finite_difference:
|
||
|
|
self.velocity = self.velocity - grad_velocity
|
||
|
|
|
||
|
|
# gradient ascent
|
||
|
|
grad = -1. * grad
|
||
|
|
return grad
|
||
|
|
|
||
|
|
def clip_gradient(self, grad):
|
||
|
|
if self.clip_grad:
|
||
|
|
if self.clip_grad_rule == 'norm':
|
||
|
|
return self.clip_grad_by_norm(grad)
|
||
|
|
elif self.clip_grad_rule == 'value':
|
||
|
|
return self.clip_grad_by_value(grad)
|
||
|
|
else:
|
||
|
|
raise NotImplementedError
|
||
|
|
else:
|
||
|
|
return grad
|
||
|
|
|
||
|
|
def clip_grad_by_norm(self, grad):
|
||
|
|
# clip gradient by norm
|
||
|
|
if self.clip_grad:
|
||
|
|
grad_norm = torch.linalg.norm(grad + 1e-6, dim=-1, keepdims=True)
|
||
|
|
scale_ratio = torch.clip(grad_norm, 0., self.max_grad_norm) / grad_norm
|
||
|
|
grad = scale_ratio * grad
|
||
|
|
return grad
|
||
|
|
|
||
|
|
def clip_grad_by_value(self, grad):
|
||
|
|
# clip gradient by value
|
||
|
|
if self.clip_grad:
|
||
|
|
grad = torch.clip(grad, -self.max_grad_value, self.max_grad_value)
|
||
|
|
return grad
|
||
|
|
|
||
|
|
|
||
|
|
class GuideManagerTrajectoriesWithVelocity(nn.Module):
|
||
|
|
|
||
|
|
def __init__(self, dataset, cost, clip_grad=False, clip_grad_rule='norm', max_grad_norm=1., max_grad_value=0.1,
|
||
|
|
interpolate_trajectories_for_collision=False,
|
||
|
|
num_interpolated_points_for_collision=128,
|
||
|
|
start_state_pos=None,
|
||
|
|
goal_state_pos=None,
|
||
|
|
num_steps=100,
|
||
|
|
robot=None,
|
||
|
|
n_samples=1,
|
||
|
|
tensor_args=None,
|
||
|
|
**kwargs):
|
||
|
|
super().__init__()
|
||
|
|
self.cost = cost
|
||
|
|
self.dataset = dataset
|
||
|
|
|
||
|
|
self.interpolate_trajectories_for_collision = interpolate_trajectories_for_collision
|
||
|
|
self.num_interpolated_points_for_collision = num_interpolated_points_for_collision
|
||
|
|
|
||
|
|
self.clip_grad = clip_grad
|
||
|
|
self.clip_grad_rule = clip_grad_rule
|
||
|
|
self.max_grad_norm = max_grad_norm
|
||
|
|
self.max_grad_value = max_grad_value
|
||
|
|
|
||
|
|
def forward(self, x_normalized):
|
||
|
|
x = x_normalized.clone()
|
||
|
|
with torch.enable_grad():
|
||
|
|
x.requires_grad_(True)
|
||
|
|
|
||
|
|
# unnormalize x
|
||
|
|
# x is normalized, but the guides are defined on unnormalized trajectory space
|
||
|
|
x = self.dataset.unnormalize_trajectories(x)
|
||
|
|
|
||
|
|
if self.interpolate_trajectories_for_collision:
|
||
|
|
# finer interpolation of trajectory for better collision avoidance
|
||
|
|
x_interpolated = interpolate_points_v1(x, num_interpolated_points=self.num_interpolated_points_for_collision)
|
||
|
|
else:
|
||
|
|
x_interpolated = x
|
||
|
|
|
||
|
|
# compute costs
|
||
|
|
# append the current velocity trajectory to the position trajectory only for non-interpolated trajectories
|
||
|
|
cost_l, weight_grad_cost_l = self.cost(x, x_interpolated=x_interpolated, return_invidual_costs_and_weights=True)
|
||
|
|
grad = 0
|
||
|
|
for cost, weight_grad_cost in zip(cost_l, weight_grad_cost_l):
|
||
|
|
if torch.is_tensor(cost):
|
||
|
|
# y.sum() is a surrogate to compute gradients of independent quantities over the batch dimension
|
||
|
|
# x are the support points. Compute gradients wrt x, not x_interpolated
|
||
|
|
grad_cost = torch.autograd.grad([cost.sum()], [x], retain_graph=True)[0]
|
||
|
|
|
||
|
|
# clip gradients
|
||
|
|
grad_cost_clipped = self.clip_gradient(grad_cost)
|
||
|
|
|
||
|
|
# zeroing gradients at start and goal
|
||
|
|
grad_cost_clipped[..., 0, :] = 0.
|
||
|
|
grad_cost_clipped[..., -1, :] = 0.
|
||
|
|
|
||
|
|
# combine gradients
|
||
|
|
grad_cost_clipped_weighted = weight_grad_cost * grad_cost_clipped
|
||
|
|
grad += grad_cost_clipped_weighted
|
||
|
|
|
||
|
|
# gradient ascent
|
||
|
|
grad = -1. * grad
|
||
|
|
return grad
|
||
|
|
|
||
|
|
def clip_gradient(self, grad):
|
||
|
|
if self.clip_grad:
|
||
|
|
if self.clip_grad_rule == 'norm':
|
||
|
|
return self.clip_grad_by_norm(grad)
|
||
|
|
elif self.clip_grad_rule == 'value':
|
||
|
|
return self.clip_grad_by_value(grad)
|
||
|
|
else:
|
||
|
|
raise NotImplementedError
|
||
|
|
else:
|
||
|
|
return grad
|
||
|
|
|
||
|
|
def clip_grad_by_norm(self, grad):
|
||
|
|
# clip gradient by norm
|
||
|
|
if self.clip_grad:
|
||
|
|
grad_norm = torch.linalg.norm(grad + 1e-6, dim=-1, keepdims=True)
|
||
|
|
scale_ratio = torch.clip(grad_norm, 0., self.max_grad_norm) / grad_norm
|
||
|
|
grad = scale_ratio * grad
|
||
|
|
return grad
|
||
|
|
|
||
|
|
def clip_grad_by_value(self, grad):
|
||
|
|
# clip gradient by value
|
||
|
|
if self.clip_grad:
|
||
|
|
grad = torch.clip(grad, -self.max_grad_value, self.max_grad_value)
|
||
|
|
return grad
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
class GuideBase(nn.Module, abc.ABC):
|
||
|
|
def __init__(self, scale=1e-3, tensor_args=None, **kwargs):
|
||
|
|
super().__init__()
|
||
|
|
self.tensor_args = tensor_args
|
||
|
|
self.scale = scale
|
||
|
|
|
||
|
|
@abc.abstractmethod
|
||
|
|
def forward(self, x):
|
||
|
|
raise NotImplementedError
|
||
|
|
|
||
|
|
def gradients(self, x):
|
||
|
|
x.requires_grad_()
|
||
|
|
y = self(x)
|
||
|
|
# y.sum() is a surrogate to compute gradients of independent quantities over the batch dimension
|
||
|
|
grad = torch.autograd.grad([y.sum()], [x])[0]
|
||
|
|
x.detach()
|
||
|
|
return y, grad
|
||
|
|
|
||
|
|
|
||
|
|
class GuideCollisionAvoidance(GuideBase):
|
||
|
|
"""
|
||
|
|
Computes the collision sdf for all points in a trajectory
|
||
|
|
"""
|
||
|
|
def __init__(self, env, **kwargs):
|
||
|
|
super().__init__(**kwargs)
|
||
|
|
self.env = env
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
collision_cost = self.env.compute_collision_cost(x, field_type='sdf')
|
||
|
|
cost = collision_cost.sum(-1)
|
||
|
|
return -1 * cost # maximize
|
||
|
|
|
||
|
|
|
||
|
|
class GuideSmoothnessFiniteDifferenceVelocity(GuideBase):
|
||
|
|
"""
|
||
|
|
Smoothness cost as the central finite difference of velocity, aka acceleration
|
||
|
|
"""
|
||
|
|
def __init__(self, env, method='central', **kwargs):
|
||
|
|
super().__init__(**kwargs)
|
||
|
|
self.env = env
|
||
|
|
self.method = method
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
if self.method == 'central':
|
||
|
|
vel = self.env.get_q_velocity(x)
|
||
|
|
acc = 0.5 * (vel[..., 1:, :] - vel[..., :-1, :])
|
||
|
|
# minimize sum of accelerations along trajectory
|
||
|
|
cost = torch.linalg.norm(acc, dim=-1).sum(-1)
|
||
|
|
else:
|
||
|
|
raise NotImplementedError
|
||
|
|
return -1 * cost # maximize
|
||
|
|
|
||
|
|
|
||
|
|
class GuideSmoothnessGPPrior(GuideBase):
|
||
|
|
"""
|
||
|
|
Smoothness cost as a GP Prior
|
||
|
|
"""
|
||
|
|
def __init__(self, n_dofs, n_support_points_des, start_state, dt, cost_sigmas, **kwargs):
|
||
|
|
super().__init__(**kwargs)
|
||
|
|
# self.cost_prior = CostGP(n_dofs, n_support_points_des, start_state, dt, cost_sigmas, **kwargs)
|
||
|
|
self.cost_prior = CostGPTrajectory(n_dofs, n_support_points_des, start_state, dt, cost_sigmas, **kwargs)
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
cost = self.cost_prior(x)
|
||
|
|
return -1 * cost # maximize
|
||
|
|
|
||
|
|
|
||
|
|
class GuideStateGoal(GuideBase):
|
||
|
|
"""
|
||
|
|
State desired cost
|
||
|
|
"""
|
||
|
|
def __init__(self, state_des, idx, **kwargs):
|
||
|
|
super().__init__(**kwargs)
|
||
|
|
self.state_des = state_des
|
||
|
|
self.idx = idx
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
assert x.ndim >= 3
|
||
|
|
cost = torch.linalg.norm(x[..., self.idx, :] - self.state_des, dim=-1)
|
||
|
|
return -1 * cost # maximize
|
||
|
|
|
||
|
|
|
||
|
|
class GuideTrajectorySmoothnessFirstOrder(GuideBase):
|
||
|
|
"""
|
||
|
|
Computes the first-order smoothness of a trajectory
|
||
|
|
d(q_0,...,q_{H-1}) = sum_{i=1}^{H-1}(||q_i-q_{i-1}||^2)
|
||
|
|
"""
|
||
|
|
def __init__(self, **kwargs):
|
||
|
|
super().__init__(**kwargs)
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
# x: [batch, horizon, dim]
|
||
|
|
position_diff = torch.diff(x, dim=1)
|
||
|
|
distance = torch.linalg.norm(position_diff, dim=-1).sum(-1)
|
||
|
|
return -1. * distance # maximize the negative distance = minimize distance
|
||
|
|
|
||
|
|
|
||
|
|
class GuideSE3OrientationGoal(GuideBase):
|
||
|
|
"""
|
||
|
|
End effector orientation goal
|
||
|
|
"""
|
||
|
|
def __init__(self, env, rot_des=None, **kwargs):
|
||
|
|
super().__init__(**kwargs)
|
||
|
|
self.env = env
|
||
|
|
if rot_des is not None:
|
||
|
|
self.rot_des = rot_des.to(**self.tensor_args)
|
||
|
|
else:
|
||
|
|
self.rot_des = torch.eye(1).to(**self.tensor_args)
|
||
|
|
|
||
|
|
self.w_rot = 1.
|
||
|
|
|
||
|
|
def forward(self, q):
|
||
|
|
b = 1
|
||
|
|
h = 1
|
||
|
|
if q.ndim == 1:
|
||
|
|
q = q.unsqueeze(0).unsqueeze(0)
|
||
|
|
elif q.ndim == 2:
|
||
|
|
b = q.shape[0]
|
||
|
|
q = q.unsqueeze(0)
|
||
|
|
elif q.ndim == 3:
|
||
|
|
b = q.shape[0]
|
||
|
|
h = q.shape[1]
|
||
|
|
elif q.ndim > 3:
|
||
|
|
raise NotImplementedError
|
||
|
|
|
||
|
|
# batch, trajectory length, q dimension
|
||
|
|
q = einops.rearrange(q, 'b h d -> (b h) d')
|
||
|
|
# link_tensor = self.diff_panda.compute_forward_kinematics_all_links(q)
|
||
|
|
link_tensor_EE = self.env.diff_panda.compute_forward_kinematics_link_list(q, link_list=[self.env.link_name_ee])
|
||
|
|
# reshape to batch, trajectory, link poses
|
||
|
|
link_tensor_EE = einops.rearrange(link_tensor_EE, '(b h) 1 d1 d2 -> b h d1 d2', b=b, h=h)
|
||
|
|
|
||
|
|
H_des = link_tensor_EE.clone()
|
||
|
|
H_des[..., :3, :3] = self.rot_des
|
||
|
|
|
||
|
|
cost = SE3_distance(link_tensor_EE, H_des, w_rot=self.w_rot)
|
||
|
|
cost = cost.sum(-1)
|
||
|
|
return -1 * cost # maximize
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
class GuideTrajectoryLastPoint(GuideBase):
|
||
|
|
"""
|
||
|
|
https://arxiv.org/pdf/2301.06015.pdf
|
||
|
|
"""
|
||
|
|
def __init__(self, goal_state, **kwargs):
|
||
|
|
super().__init__(**kwargs)
|
||
|
|
self._goal_state = goal_state
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
# x: [batch, horizon, dim]
|
||
|
|
l1_dist_goal_state = torch.abs(self._goal_state - x).sum(-1)
|
||
|
|
# Equation A5
|
||
|
|
# objective = torch.exp(1/l1_dist_goal_state).sum(-1)
|
||
|
|
# Equation A8
|
||
|
|
objective = -1. * l1_dist_goal_state.sum(-1)
|
||
|
|
return objective
|
||
|
|
|
||
|
|
|
||
|
|
class GuideTrajectoryObstacleAvoidanceMultiSphere(GuideBase):
|
||
|
|
"""
|
||
|
|
Computes the sdf for all points in a trajectory
|
||
|
|
"""
|
||
|
|
def __init__(self, obstacles, tensor_args=None, **kwargs):
|
||
|
|
super().__init__(**kwargs)
|
||
|
|
obst_params = dict(obst_type='sdf')
|
||
|
|
self.shape = MultiSphere(tensor_args, **obst_params)
|
||
|
|
self.convert_obstacles_to_shape(obstacles)
|
||
|
|
self._max_sdf = 0.02
|
||
|
|
|
||
|
|
def convert_obstacles_to_shape(self, obstacles):
|
||
|
|
centers = []
|
||
|
|
radii = []
|
||
|
|
for obstacle in obstacles:
|
||
|
|
assert isinstance(obstacle, ObstacleSphere), "Only accepts circles for now"
|
||
|
|
centers.append(obstacle.get_center())
|
||
|
|
radii.append(obstacle.radius)
|
||
|
|
centers = np.array(centers)
|
||
|
|
radii = np.array(radii)
|
||
|
|
self.shape.set_obst(centers=centers, radii=radii)
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
sdf_points = self.shape.compute_cost(x)
|
||
|
|
cost_points = torch.relu(self._max_sdf - sdf_points)
|
||
|
|
# cost_points = smooth_distance_penalty(sdf_points)
|
||
|
|
cost_trajectory = cost_points.sum(-1)
|
||
|
|
return -1 * cost_trajectory # maximize
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
def smooth_distance_penalty(d, max_sdf=0.1):
|
||
|
|
# TODO this does not allow the gradient to backpropagate because of the torch.where operator
|
||
|
|
raise NotImplementedError
|
||
|
|
# https://www.ri.cmu.edu/pub_files/2009/5/icra09-chomp.pdf
|
||
|
|
# Figure 2
|
||
|
|
d_new = torch.zeros_like(d)
|
||
|
|
idxs1 = torch.where(d < 0)
|
||
|
|
d[idxs1] = -d[idxs1] + 0.5*max_sdf
|
||
|
|
idxs2 = torch.where(torch.logical_and(d >= 0, d <= max_sdf))
|
||
|
|
d[idxs2] = 1/(2*max_sdf)*(d[idxs2] - max_sdf)**2
|
||
|
|
return d_new
|
||
|
|
|
||
|
|
|
||
|
|
class GuideStochGPMP(GuideBase):
|
||
|
|
|
||
|
|
def __init__(self, env, start_state, goal_state, tensor_args, **kwargs):
|
||
|
|
super().__init__(**kwargs)
|
||
|
|
|
||
|
|
self.env = env
|
||
|
|
self.start_state = start_state
|
||
|
|
self.goal_state = goal_state
|
||
|
|
|
||
|
|
self.tensor_args = tensor_args
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
pass
|
||
|
|
|
||
|
|
def gradients(self, x):
|
||
|
|
# One step of local StochGPMP
|
||
|
|
y, grad = self.plan_sgpmp(x)
|
||
|
|
return y, grad
|
||
|
|
|
||
|
|
def plan_sgpmp(self, prior_traj, sgpmp_opt_iters=1, step_size=0.05, return_trajs=False):
|
||
|
|
prior_traj = to_torch(prior_traj, **self.tensor_args)
|
||
|
|
|
||
|
|
b, h, d = prior_traj.shape
|
||
|
|
n_support_points_des = h
|
||
|
|
|
||
|
|
# SGPMP planner
|
||
|
|
n_dofs = self.env.n_dofs
|
||
|
|
# Add velocities to state
|
||
|
|
start_state = torch.cat((self.start_state, torch.zeros(n_dofs, **self.tensor_args)))
|
||
|
|
multi_goal_states = einops.repeat(torch.cat((self.goal_state, torch.zeros(n_dofs, **self.tensor_args))),
|
||
|
|
'd -> b d', b=b)
|
||
|
|
|
||
|
|
dt = 0.001 * 64 / n_support_points_des
|
||
|
|
num_particles_per_goal = 1
|
||
|
|
num_samples = 50
|
||
|
|
sgpmp_opt_iters = sgpmp_opt_iters
|
||
|
|
|
||
|
|
# Construct cost functions
|
||
|
|
cost_sigmas = dict(
|
||
|
|
sigma_start=0.0001,
|
||
|
|
sigma_gp=0.3,
|
||
|
|
)
|
||
|
|
cost_prior = CostGP(
|
||
|
|
n_dofs, n_support_points_des, start_state, dt,
|
||
|
|
cost_sigmas, self.tensor_args
|
||
|
|
)
|
||
|
|
|
||
|
|
sigma_goal_prior = 0.0001
|
||
|
|
cost_goal_prior = CostGoalPrior(
|
||
|
|
n_dofs,
|
||
|
|
n_support_points_des,
|
||
|
|
multi_goal_states=multi_goal_states,
|
||
|
|
num_particles_per_goal=num_particles_per_goal,
|
||
|
|
num_samples=num_samples,
|
||
|
|
sigma_goal_prior=sigma_goal_prior,
|
||
|
|
tensor_args=self.tensor_args
|
||
|
|
)
|
||
|
|
|
||
|
|
sigma_coll = 1e-5
|
||
|
|
cost_obst_2D = CostCollision(
|
||
|
|
n_dofs, n_support_points_des,
|
||
|
|
field=self.env.obstacle_map,
|
||
|
|
sigma_coll=sigma_coll,
|
||
|
|
tensor_args=self.tensor_args
|
||
|
|
)
|
||
|
|
cost_func_list = [cost_prior, cost_goal_prior, cost_obst_2D]
|
||
|
|
cost_composite = CostComposite(n_dofs, n_support_points_des, cost_func_list)
|
||
|
|
|
||
|
|
# Prior mean trajectory
|
||
|
|
initial_particle_means = None
|
||
|
|
if prior_traj is not None:
|
||
|
|
# Interpolate and smooth to desired trajectory length
|
||
|
|
# prior_traj = smoothen_trajectory(prior_traj, n_support_points_des, self.tensor_args, smooth_factor=0.01)
|
||
|
|
# Reshape for sgpmp interface
|
||
|
|
prior_traj = einops.repeat(prior_traj, 'b h d -> b n h d', n=num_particles_per_goal)
|
||
|
|
# Add velocities to state
|
||
|
|
initial_particle_means = torch.cat((prior_traj, torch.zeros_like(prior_traj)), dim=-1)
|
||
|
|
|
||
|
|
sgpmp_params = dict(
|
||
|
|
start_state=start_state,
|
||
|
|
multi_goal_states=multi_goal_states,
|
||
|
|
cost=cost_composite,
|
||
|
|
initial_particle_means=initial_particle_means,
|
||
|
|
num_particles_per_goal=num_particles_per_goal,
|
||
|
|
num_samples=num_samples,
|
||
|
|
n_support_points=n_support_points_des,
|
||
|
|
dt=dt,
|
||
|
|
n_dof=n_dofs,
|
||
|
|
opt_iters=sgpmp_opt_iters,
|
||
|
|
temp=0.1,
|
||
|
|
step_size=step_size,
|
||
|
|
sigma_start_init=1e-3,
|
||
|
|
sigma_goal_init=1e-3,
|
||
|
|
sigma_gp_init=5.,
|
||
|
|
sigma_start_sample=1e-3,
|
||
|
|
sigma_goal_sample=1e-3,
|
||
|
|
sigma_gp_sample=5.,
|
||
|
|
tensor_args=self.tensor_args,
|
||
|
|
)
|
||
|
|
|
||
|
|
sgpmp_planner = StochGPMP(**sgpmp_params)
|
||
|
|
obs = {}
|
||
|
|
|
||
|
|
# Optimize
|
||
|
|
sgpmp_time_start = time.time()
|
||
|
|
_, _, _, _, costs, grad = sgpmp_planner.optimize(**obs)
|
||
|
|
|
||
|
|
# costs of all particles
|
||
|
|
costs = costs.sum(1)
|
||
|
|
|
||
|
|
# gradient includes velocities
|
||
|
|
grad = grad[..., :2]
|
||
|
|
|
||
|
|
sgpmp_time_finish = time.time()
|
||
|
|
|
||
|
|
controls, _, trajectories, trajectory_means, weights = sgpmp_planner.get_recent_samples()
|
||
|
|
|
||
|
|
# print(f'SGPMP time: {sgpmp_time_finish - sgpmp_time_start:.4f} sec')
|
||
|
|
|
||
|
|
if return_trajs:
|
||
|
|
return trajectories
|
||
|
|
|
||
|
|
return costs, grad
|