Implement potential field visualization, tested on 2d demos

Signed-off-by: An Thai Le <an.thai.le97@gmail.com>
ASUDCT · Aug 22, 2021 · 2ed96cb · 2ed96cb
1 parent c90aabf
commit 2ed96cb
Show file tree

Hide file tree

Showing 10 changed files with 295 additions and 118 deletions.
diff --git a/scripts/test_tprmp.py b/scripts/test_tprmp.py
@@ -30,7 +30,7 @@
 T = 300
 NUM_COMP = 30
 alpha, beta = 0., 0.
-stiff_scale = 3.
+stiff_scale = 2.
 d_min = 0.
 d_scale = 3.
 energy = 0.

diff --git a/scripts/test_tprmp_2d.py b/scripts/test_tprmp_2d.py
@@ -10,6 +10,7 @@
 sys.path.append(ROOT_DIR)
 from tprmp.utils.loading import load  # noqa
 from tprmp.visualization.demonstration import plot_demo  # noqa
+from tprmp.visualization.dynamics import plot_potential_field, visualize_rmp # noqa
 from tprmp.models.tp_rmp import TPRMP  # noqa
 from tprmp.demonstrations.base import Demonstration  # noqa
 from tprmp.demonstrations.manifold import Manifold  # noqa
@@ -25,17 +26,22 @@
 DATA_DIR = join(ROOT_DIR, 'data', 'tasks', args.task, 'demos')
 data_file = join(DATA_DIR, args.data)
 # parameters
-T = 300
+oversteps = 50
 dt = 0.01
-NUM_COMP = 30
+NUM_COMP = 20
 alpha, beta = 0., 0.
-stiff_scale = 3.
+stiff_scale = 1.
+tau = 1.
+potential_method = 'quadratic'
+optimize_method = 'flow'
 d_min = 0.
-d_scale = 3.
+d_scale = 1.
 energy = 0.
 sigma = 1.
 var_scale = 1.
 r = 0.01
+res = 0.05
+verbose = False
 # load data
 data = load(data_file)
 demos = []
@@ -46,7 +52,15 @@
     demo.add_frame_from_pose(traj[:, 0], 'start')
     demo.add_frame_from_pose(traj[:, -1], 'end')
     demos.append(demo)
+plot_demo(demos, only_global=False, plot_quat=False, new_fig=True, new_ax=True, three_d=False, show=True)
 # train tprmp
-model = TPRMP(num_comp=NUM_COMP, name=args.task, sigma=sigma, stiff_scale=stiff_scale, 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)
+sample = demos[0]
+frames = sample.get_task_parameters()
+model = TPRMP(num_comp=NUM_COMP, name=args.task, sigma=sigma, stiff_scale=stiff_scale, tau=tau, potential_method=potential_method, d_scale=d_scale)
+model.train(demos, optimize_method=optimize_method, alpha=alpha, beta=beta, d_min=d_min, energy=energy, var_scale=var_scale, verbose=verbose)
+model.model.plot_model(demos, tagging=False, three_d=False)
+plot_potential_field(model, frames, only_global=False, margin=0.5, res=res, new_fig=True, show=True)
+# execution
+x0, dx0 = np.array([0.5, 2.5]), np.zeros(2)
+visualize_rmp(model, frames, x0, dx0, sample.traj.shape[1] + oversteps, dt, sample=sample, x_limits=[0., 4.], vel_limits=[-10., 10.])
+input()
diff --git a/tprmp/demonstrations/base.py b/tprmp/demonstrations/base.py
@@ -178,6 +178,14 @@ def construct_linear_map(manifold, pose):
     def traj(self):
         return self._traj
 
+    @property
+    def d_traj(self):
+        return self._d_traj
+
+    @property
+    def dd_traj(self):
+        return self._dd_traj
+
     @traj.setter
     def traj(self, value):
         if isinstance(value, list):

diff --git a/tprmp/models/rmp.py b/tprmp/models/rmp.py
@@ -5,23 +5,32 @@
 def compute_riemannian_metric(x, mvns):
     weights = compute_obsrv_prob(x, mvns)
     Ms = np.array([comp.cov_inv for comp in mvns])
-    return Ms.T @ weights
+    M = Ms.T @ weights if weights.sum() == 1. else np.eye(mvns[0].manifold.dim_T)
+    return M
 
 
-def compute_policy(phi0, d0, x, dx, mvns, stiff_scale=2., use_cp=False):
+def compute_policy(phi0, d0, x, dx, mvns, stiff_scale=1., tau=1., potential_method='quadratic', use_cp=False):
     weights = compute_obsrv_prob(x, mvns)
-    return compute_potential_term(weights, phi0, x, mvns, stiff_scale=stiff_scale) + compute_dissipation_term(weights, d0, dx, mvns, use_cp=use_cp)
+    return compute_potential_term(weights, phi0, x, mvns, stiff_scale=stiff_scale, tau=tau, potential_method=potential_method) + compute_dissipation_term(weights, d0, dx, mvns, use_cp=use_cp)
 
 
-def compute_potential_term(weights, phi0, x, mvns, stiff_scale=2.):
-    phi = compute_potentials(phi0, x, mvns)
+def compute_potential_term(weights, phi0, x, mvns, stiff_scale=1., tau=1., potential_method='quadratic'):
+    phi = compute_potentials(phi0, x, mvns, stiff_scale=stiff_scale, tau=tau, potential_method=potential_method)
     Phi = weights.T @ phi if len(x.shape) == 1 else np.diag(weights.T @ phi)
     num_comp = len(mvns)
     manifold = mvns[0].manifold
     Ps = np.zeros(manifold.dim_T) if len(x.shape) == 1 else np.zeros((manifold.dim_T, x.shape[1]))
     for k in range(num_comp):
-        Ps += weights[k] * (phi[k] - Phi) * (mvns[k].cov_inv @ manifold.log_map(x, base=mvns[k].mean))
-        Ps += -weights[k] * (stiff_scale * mvns[k].cov_inv @ manifold.log_map(x, base=mvns[k].mean))
+        v = manifold.log_map(x, base=mvns[k].mean)
+        pull = stiff_scale * mvns[k].cov_inv @ v
+        Ps += weights[k] * (phi[k] - Phi) * pull
+        if potential_method == 'quadratic':
+            Ps += -weights[k] * pull
+        elif potential_method == 'tanh':
+            norm = np.sqrt(v.T @ pull)
+            Ps += -weights[k] * np.tanh(tau * norm) * pull / norm
+        else:
+            raise ValueError(f'Potential method {potential_method} is unrecognized!')
     return Ps
 
 
@@ -40,14 +49,20 @@ def compute_dissipation_term(weights, d0, dx, mvns, use_cp=False):
     return Ds
 
 
-def compute_potentials(phi0, x, mvns):
+def compute_potentials(phi0, x, mvns, stiff_scale=1., tau=1., potential_method='quadratic'):
     num_comp = len(mvns)
     P = np.zeros(num_comp) if len(x.shape) == 1 else np.zeros((num_comp, x.shape[1]))
     manifold = mvns[0].manifold
     for k in range(num_comp):
         comp = mvns[k]
         v = manifold.log_map(x, base=comp.mean)
-        P[k] = (v * (comp.cov_inv @ v)).sum(0)
+        quadratic = (v * (stiff_scale * comp.cov_inv @ v)).sum(0)
+        if potential_method == 'quadratic':
+            P[k] = quadratic
+        elif potential_method == 'tanh':
+            P[k] = 1 / tau * (np.exp(tau * np.sqrt(quadratic)) + np.exp(-tau * np.sqrt(quadratic)))
+        else:
+            raise ValueError(f'Potential method {potential_method} is unrecognized!')
     phi = phi0 + P if len(x.shape) == 1 else np.expand_dims(phi0, axis=1) + P
     return phi
 
@@ -59,12 +74,8 @@ def compute_obsrv_prob(x, mvns, normalized=True, eps=1e-30):
         prob[k] = mvns[k].pdf(x)
     if normalized:
         s = prob.sum()
-        if s < eps:  # use distance as metric to compute prob
-            manifold = mvns[0].manifold
-            dist = np.array([np.linalg.norm(manifold.log_map(x, base=mvns[k].mean)) for k in range(len(mvns))])
-            prob = dist / dist.sum()
-        else:
-            prob /= prob.sum()
+        if s > eps:
+            prob /= s
     return prob
 
 
@@ -73,7 +84,7 @@ def compute_obsrv_prob(x, mvns, normalized=True, eps=1e-30):
     sys.path.append('.')
     from tprmp.demonstrations.probability import ManifoldGaussian
     from tprmp.demonstrations.manifold import Manifold
-    from tprmp.visualization.dynamics import visualize_rmp
+    # from tprmp.visualization.dynamics import visualize_rmp
     manifold = Manifold.get_euclidean_manifold(2)
     mvns = [ManifoldGaussian(manifold, np.ones(2), np.eye(2)),
             ManifoldGaussian(manifold, 3 * np.ones(2), 2 * np.eye(2)),
@@ -89,5 +100,5 @@ def compute_obsrv_prob(x, mvns, normalized=True, eps=1e-30):
     x = np.zeros((2, T))
     print(compute_riemannian_metric(x, mvns).shape)
     # # test dynamics
-    x, dx = np.zeros(2), np.zeros(2)
-    visualize_rmp(phi0, d0, mvns, x, dx, T, dt, limit=10)
+    # x, dx = np.zeros(2), np.zeros(2)
+    # visualize_rmp(phi0, d0, mvns, x, dx, T, dt, limit=10)
diff --git a/tprmp/models/tp_hsmm.py b/tprmp/models/tp_hsmm.py
@@ -104,15 +104,16 @@ def set_params_with_tag(self, tag_to_params):
         model_params['end_states'] = model_params['end_states'] / model_params['end_states'].sum()
         self.set_params(model_params)
 
-    def plot_model(self, demos, tagging=True, plot_quat=False, show=True):
+    def plot_model(self, demos, tagging=True, plot_transition=False, plot_quat=False, three_d=True, show=True):
         if isinstance(demos, Demonstration):
             demos = [demos]
-        plot_hsmm(self, new_fig=True, show=show)
+        if plot_transition:
+            plot_hsmm(self, new_fig=True, show=show)
         tag_to_demos = TPHSMM.compute_tag_to_demo(demos)
         for tag in tag_to_demos:  # plot one demo for each tag
             demo = tag_to_demos[tag][0]
-            plot_demo(demo, plot_quat=plot_quat, new_fig=True, new_ax=True, show=False)
-            plot_gmm(self, demo.get_task_parameters(), plot_quat=plot_quat, tag=tag if tagging else None, new_fig=False, show=show)
+            plot_demo(demo, plot_quat=plot_quat, new_fig=True, new_ax=True, three_d=three_d, show=False)
+            plot_gmm(self, demo.get_task_parameters(), plot_quat=plot_quat, tag=tag if tagging else None, new_fig=False, three_d=three_d, show=show)
 
     def set_params(self, model_params):
         super(TPHSMM, self).set_params(model_params)

diff --git a/tprmp/models/tp_rmp.py b/tprmp/models/tp_rmp.py
@@ -6,7 +6,7 @@
 import time
 
 from tprmp.models.tp_hsmm import TPHSMM
-from tprmp.models.rmp import compute_policy, compute_riemannian_metric
+from tprmp.models.rmp import compute_policy, compute_riemannian_metric, compute_potentials, compute_obsrv_prob
 from tprmp.models.coriolis import compute_coriolis_force
 from tprmp.optimizer.dynamics import optimize_dynamics
 from tprmp.utils.loading import load
@@ -24,7 +24,9 @@ class TPRMP(object):
 
     def __init__(self, **kwargs):
         self._sigma = kwargs.pop('sigma', 1.)
-        self._stiff_scale = kwargs.pop('stiff_scale', 2.)
+        self._stiff_scale = kwargs.pop('stiff_scale', 1.)
+        self._tau = kwargs.pop('tau', 1.)
+        self._potential_method = kwargs.pop('potential_method', 'quadratic')
         self._d_scale = kwargs.pop('d_scale', 1.)
         self._model = TPHSMM(**kwargs)
         self._global_mvns = None
@@ -68,7 +70,8 @@ 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], self._d_scale * self._d0[f_key], lx, ldx, self.model.get_local_gmm(f_key), stiff_scale=self._stiff_scale)
+            local_policy = compute_policy(self._phi0[f_key], self._d_scale * self._d0[f_key], lx, ldx, self.model.get_local_gmm(f_key),
+                                          stiff_scale=self._stiff_scale, tau=self._tau, potential_method=self._potential_method)
             policies[i] = weights[f_key] * frames[f_key].transform_tangent(local_policy)
         return policies.sum(0)
 
@@ -82,18 +85,29 @@ def compute_frame_weights(self, x, frames, normalized=True, eps=1e-30):
             weights[f] = w
         s = sum(weights.values())
         if normalized:
-            if s < eps:  # use distance as metric to compute prob
-                for f in weights:
-                    d = np.linalg.norm(self.model.manifold.log_map(x, base=frame_origins[f]))
-                    weights[f] = d
-                norm_dist = sum(weights.values())
-                for f in weights:
-                    weights[f] /= norm_dist
-            else:
+            if s > eps:
                 for f in weights:
                     weights[f] /= s
         return weights
 
+    def compute_potential_field(self, x, frames):
+        frame_weights = self.compute_frame_weights(x, frames)
+        Phi = 0.
+        for f_key in frames:
+            lx = frames[f_key].pullback(x)
+            mvns = self.model.get_local_gmm(f_key)
+            weights = compute_obsrv_prob(lx, mvns)
+            phi = compute_potentials(self.phi0[f_key], lx, mvns, stiff_scale=self._stiff_scale, tau=self._tau, potential_method=self._potential_method)
+            Phi += frame_weights[f_key] * (weights.T @ phi)
+        return Phi
+
+    def compute_potential_field_frame(self, lx, frame):
+        mvns = self.model.get_local_gmm(frame)
+        weights = compute_obsrv_prob(lx, mvns)
+        phi = compute_potentials(self.phi0[frame], lx, mvns, stiff_scale=self._stiff_scale, tau=self._tau, potential_method=self._potential_method)
+        Phi = weights.T @ phi
+        return Phi
+
     def train(self, demos, **kwargs):
         """
         Trains the TP-RMP with a given set of demonstrations.
@@ -102,11 +116,13 @@ def train(self, demos, **kwargs):
         ----------
         :param demos: list of Demonstration objects
         """
+        optimize_method = kwargs.get('optimize_method', 'flow')
         alpha = kwargs.get('alpha', 1e-5)
         beta = kwargs.get('beta', 1e-5)
         min_d = kwargs.get('min_d', 20.)
         energy = kwargs.get('energy', 0.)
         var_scale = kwargs.get('var_scale', 1.)
+        verbose = kwargs.get('verbose', False)
         # train TP-HSMM/TP-GMM
         self.model.train(demos, **kwargs)
         if 'S' in self.model.manifold.name:  # decouple orientation and position
@@ -115,7 +131,8 @@ def train(self, demos, **kwargs):
         if var_scale > 1.:
             self.model.scale_covariance(var_scale)
         # train dynamics
-        self._phi0, self._d0 = optimize_dynamics(self.model, demos, alpha=alpha, beta=beta, stiff_scale=self._stiff_scale, min_d=min_d, energy=energy)
+        self._phi0, self._d0 = optimize_dynamics(self.model, demos, alpha=alpha, beta=beta, optimize_method=optimize_method,
+                                                 stiff_scale=self._stiff_scale, tau=self._tau, potential_method=self._potential_method, min_d=min_d, energy=energy, verbose=verbose)
         # train local Riemannian metrics TODO: RiemannianNetwork is still under consideration
         # self._R_net = optimize_riemannian_metric(self, demos, **kwargs)
 

diff --git a/tprmp/optimizer/dynamics.py b/tprmp/optimizer/dynamics.py
@@ -2,17 +2,21 @@
 import numpy as np
 import logging
 
-from tprmp.models.rmp import compute_policy  # , compute_riemannian_metric
+from tprmp.models.rmp import compute_policy, compute_riemannian_metric
 
 logger = logging.getLogger(__name__)
 
 
 def optimize_dynamics(tp_hsmm, demos, **kwargs):
     alpha = kwargs.get('alpha', 1e-5)
     beta = kwargs.get('beta', 1e-5)
-    stiff_scale = kwargs.get('stiff_scale', 2.)
+    stiff_scale = kwargs.get('stiff_scale', 1.)
+    tau = kwargs.get('tau', 1.)
+    potential_method = kwargs.get('potential_method', 'quadratic')
     d_min = kwargs.get('d_min', 0.)
-    energy = kwargs.get('alpha', 0.)
+    energy = kwargs.get('energy', 0.)
+    optimize_method = kwargs.get('optimize_method', 'flow')
+    verbose = kwargs.get('verbose', False)
     frames = demos[0].frame_names
     gap = energy / tp_hsmm.num_comp
     phi0_frames = {}
@@ -26,17 +30,23 @@ def optimize_dynamics(tp_hsmm, demos, **kwargs):
             x, dx, ddx = trajs['traj'], trajs['d_traj'], trajs['dd_traj']
             for t in range(x.shape[1]):
                 mvns = tp_hsmm.get_local_gmm(frame)
-                # M = compute_riemannian_metric(x[:, t], mvns)
-                f = compute_policy(phi0, d0, x[:, t], dx[:, t], mvns, stiff_scale=stiff_scale, use_cp=True)
-                loss += cp.norm(ddx[:, t] - f) / x.shape[1]
+                if optimize_method == 'flow':
+                    M_inv = np.eye(tp_hsmm.manifold.dim_T)
+                elif optimize_method == 'riemannian':
+                    M = compute_riemannian_metric(x[:, t], mvns)
+                    M_inv = np.linalg.inv(M)
+                else:
+                    raise ValueError(f'Optimizing method {optimize_method} is not regconized!')
+                f = compute_policy(phi0, d0, x[:, t], dx[:, t], mvns, stiff_scale=stiff_scale, tau=tau, potential_method=potential_method, use_cp=True)
+                loss += cp.norm(ddx[:, t] - M_inv @ f) / x.shape[1]
         loss /= len(demos)
         if alpha > 0.:
             loss += alpha * cp.pnorm(phi0, p=2)**2
         if beta > 0.:
             loss += beta * cp.pnorm(d0, p=2)**2
         objective = cp.Minimize(loss)  # L2 regularization
         problem = cp.Problem(objective, field_constraints(phi0, d0, gap, d_min))
-        problem.solve()
+        problem.solve(verbose=verbose)
         logger.info(f'Opimizing dynamics for frame {frame}...')
         logger.info(f'Status: {problem.status}')
         logger.info(f'Optimal phi0: {phi0.value}')
@@ -59,7 +69,7 @@ def field_constraints(phi0, d0, gap=0., d_min=0.):
     from tprmp.demonstrations.manifold import Manifold
     from tprmp.demonstrations.base import Demonstration
     from tprmp.models.tp_gmm import TPGMM
-    from tprmp.visualization.dynamics import visualize_rmp
+    # from tprmp.visualization.dynamics import visualize_rmp
     logging.basicConfig()
     logging.getLogger().setLevel(logging.INFO)
 
@@ -88,5 +98,5 @@ def simple_demo(T, dt):
     # test training
     phi0, d0 = optimize_dynamics(model, [demo], alpha=1e-5, beta=1e-5)
     # test retrieval
-    x0, dx0 = np.array([0, 0.5]), np.zeros(2)
-    visualize_rmp(phi0['obj'], d0['obj'], model.get_local_gmm('obj'), x0, dx0, T, dt, limit=10)
+    # x0, dx0 = np.array([0, 0.5]), np.zeros(2)
+    # visualize_rmp(phi0['obj'], d0['obj'], model.get_local_gmm('obj'), x0, dx0, T, dt, limit=10)