Add damping scale

Signed-off-by: An Thai Le <an.thai.le97@gmail.com>
ASUDCT · Aug 19, 2021 · d386661 · d386661
1 parent 204ff3f
commit d386661
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 14 deletions.
diff --git a/scripts/test_tprmp.py b/scripts/test_tprmp.py
@@ -15,7 +15,7 @@
 from tprmp.models.tp_rmp import TPRMP
 from tprmp.demonstrations.base import Demonstration
 from tprmp.demonstrations.frame import Frame
-from tprmp.demonstrations.quaternion import q_to_euler, q_convert_wxyz, q_from_euler, q_convert_xyzw  # noqa
+from tprmp.demonstrations.quaternion import q_convert_wxyz, q_from_euler, q_convert_xyzw  # noqa
 from tprmp.envs.gym import Environment # noqa
 from tprmp.envs.tasks import PalletizingBoxes # noqa
 
@@ -29,30 +29,32 @@
 data_file = join(DATA_DIR, args.data)
 # parameters
 T = 300
-dt = 0.001
+dt = 0.0001
 NUM_COMP = 30
 alpha, beta = 1e-2, 0.
-min_d = 0.
+d_min = 0.
+d_scale = 150.
 energy = 0.
 sigma = 0.5
 var_scale = 2.
+r = 0.01
 # load data
 demos = load_demos(data_file, tag='pick_side')
-demos = [demos[-1]]
+demos.pop(0)
 manifold = demos[0].manifold
-# plot_demo(demos, only_global=False, plot_quat=False, new_fig=True, new_ax=True, show=True)
+plot_demo(demos, only_global=False, plot_quat=False, new_fig=True, new_ax=True, show=True)
 # train tprmp
-model = TPRMP(num_comp=NUM_COMP, name=args.task, sigma=sigma)
-model.train(demos, alpha=alpha, beta=beta, min_d=min_d, energy=energy, var_scale=var_scale)
-# model.model.plot_model(demos)
+model = TPRMP(num_comp=NUM_COMP, name=args.task, sigma=sigma, d_scale=d_scale)
+model.train(demos, alpha=alpha, beta=beta, d_min=d_min, energy=energy, var_scale=var_scale)
+model.model.plot_model(demos)
 # test tprmp
 env = Environment(task=PalletizingBoxes(), disp=True)
 ee_pose = np.array(env.home_pose)
 ee_pose[3:] = q_convert_wxyz(ee_pose[3:])
 A, b = Demonstration.construct_linear_map(manifold, ee_pose)
 ee_frame = Frame(A, b, manifold=manifold)
 box_id = env.task.goals[0][0][0]
-position = np.array([0.5, -0.25, env.task.box_size[0] / 2])
+position = np.array([0.5, -0.25, env.task.box_size[0] / 2]) + np.random.uniform(low=-r, high=r) * np.array([1, 1, 0])
 rotation = q_convert_xyzw(q_from_euler(np.array([np.pi/2, 0., 0.])))
 p.resetBasePositionAndOrientation(box_id, position, rotation)
 target = p.getBasePositionAndOrientation(box_id)

diff --git a/tprmp/demonstrations/trajectory.py b/tprmp/demonstrations/trajectory.py
@@ -100,3 +100,9 @@ def compute_traj_velocity(traj, dt, manifold=None):
     if d_traj.shape[1] != length:
         raise ValueError('[Trajectory]: Length of d_traj %s is not consistent with input traj length %s' % (d_traj.shape[1], length))
     return d_traj
+
+
+def interpolate(p1, p2, d_param, dist):
+    """ interpolate between points """
+    alpha = min(max(d_param / dist, 0.), 1.)
+    return (1. - alpha) * p1 + alpha * p2
diff --git a/tprmp/models/tp_rmp.py b/tprmp/models/tp_rmp.py
@@ -24,6 +24,7 @@ class TPRMP(object):
 
     def __init__(self, **kwargs):
         self._sigma = kwargs.pop('sigma', 1.)
+        self._d_scale = kwargs.pop('d_scale', 150.)
         self._model = TPHSMM(**kwargs)
         self._global_mvns = None
         self._phi0 = None
@@ -59,7 +60,7 @@ def compute_global_policy(self, x, dx, frames):
             # compute local policy
             lx = frames[f_key].pullback(x)
             ldx = frames[f_key].pullback_tangent(dx)
-            local_policy = compute_policy(self._phi0[f_key], 150 * self._d0[f_key], lx, ldx, self.model.get_local_gmm(f_key))
+            local_policy = compute_policy(self._phi0[f_key], self._d_scale * self._d0[f_key], lx, ldx, self.model.get_local_gmm(f_key))
             policies[i] = weights[f_key] * frames[f_key].transform_tangent(local_policy)
         return policies.sum(0)
 

diff --git a/tprmp/optimizer/dynamics.py b/tprmp/optimizer/dynamics.py
@@ -6,7 +6,7 @@
 logger = logging.getLogger(__name__)
 
 
-def optimize_dynamics(tp_hsmm, demos, alpha=1e-5, beta=1e-5, min_d=20., energy=10.):
+def optimize_dynamics(tp_hsmm, demos, alpha=1e-5, beta=1e-5, d_min=20., energy=10.):
     frames = demos[0].frame_names
     gap = energy / tp_hsmm.num_comp
     phi0_frames = {}
@@ -21,7 +21,7 @@ def optimize_dynamics(tp_hsmm, demos, alpha=1e-5, beta=1e-5, min_d=20., energy=1
             for t in range(x.shape[1]):
                 loss += cp.sum_squares(ddx[:, t] - compute_policy(phi0, d0, x[:, t], dx[:, t], tp_hsmm.get_local_gmm(frame), use_cp=True)) / x.shape[1]
         objective = cp.Minimize(loss / len(demos) + alpha * cp.pnorm(phi0, p=2)**2 + beta * cp.pnorm(d0, p=2)**2)  # L2 regularization
-        problem = cp.Problem(objective, field_constraints(phi0, d0, gap, min_d))
+        problem = cp.Problem(objective, field_constraints(phi0, d0, gap, d_min))
         problem.solve()
         logger.info(f'Opimizing dynamics for frame {frame}...')
         logger.info(f'Status: {problem.status}')
@@ -32,11 +32,11 @@ def optimize_dynamics(tp_hsmm, demos, alpha=1e-5, beta=1e-5, min_d=20., energy=1
     return phi0_frames, d0_frames
 
 
-def field_constraints(phi0, d0, gap=0., min_d=1.):
+def field_constraints(phi0, d0, gap=0., d_min=1.):
     constraints = []
     for k in range(phi0.size - 1):
         constraints.append(phi0[k] >= phi0[k + 1] + gap)
-    constraints.extend([phi0 >= 0, d0 >= min_d])
+    constraints.extend([phi0 >= 0, d0 >= d_min])
     return constraints