Clean up some code

Signed-off-by: An Thai Le <an.thai.le97@gmail.com>
ASUDCT · Sep 4, 2021 · 8a11fc8 · 8a11fc8
1 parent e06b61b
commit 8a11fc8
Show file tree

Hide file tree

Showing 7 changed files with 33 additions and 26 deletions.
diff --git a/data/tasks/test/models/dynamics_test3.p b/data/tasks/test/models/dynamics_test3.p
diff --git a/data/tasks/test/models/tphsmm_test3.p b/data/tasks/test/models/tphsmm_test3.p
diff --git a/scripts/test_tprmp_2d.py b/scripts/test_tprmp_2d.py
@@ -30,7 +30,7 @@
 NUM_COMP = 5
 alpha, beta = 0., 0.
 stiff_scale = 1.
-mass_scale = 0.1
+mass_scale = 1.
 tau = 0.5
 delta = 2.
 potential_method = 'huber'
@@ -40,7 +40,8 @@
 energy = 0.
 var_scale = 2.
 res = 0.05
-margin = 0.2
+max_z = 1000
+margin = 0.5
 verbose = False
 # load data
 data = load(data_file)
@@ -63,9 +64,14 @@
     model.train(demos, alpha=alpha, beta=beta, d_min=d_min, train_method=train_method, energy=energy, verbose=verbose)
     model.save(name=args.data)
 # model.model.plot_model(demos, tagging=False, var_scale=var_scale, three_d=False)
-plot_potential_field(model, frames, only_global=True, margin=margin, var_scale=var_scale, three_d=True, res=res, new_fig=True, show=False)
+plot_potential_field(model, frames, only_global=True, margin=margin, var_scale=var_scale, max_z=max_z, three_d=True, res=res, new_fig=True, show=False)
 plot_dissipation_field(model, frames, only_global=True, margin=margin, var_scale=var_scale, res=res, new_fig=True, show=True)
 # execution
 x0, dx0 = sample.traj[:, 0], 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()
+x0, dx0 = np.array([1., 2.]), 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()
+x0, dx0 = np.array([2., 1.]), 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.])
diff --git a/tprmp/models/coriolis.py b/tprmp/models/coriolis.py
@@ -6,20 +6,21 @@
 def compute_coriolis_force(x, dx, mvns, mass_scale=1.):
     weights = compute_obsrv_prob(x, mvns)
     scale = compute_scale(weights, x, mvns)
-    return compute_dMdt(weights, scale, dx, mvns, mass_scale=mass_scale) @ dx - 0.5 * compute_dTdx(weights, scale, dx, mvns, mass_scale=mass_scale)
+    c = compute_dMdt(weights, scale, dx, mvns) @ dx - 0.5 * compute_dTdx(weights, scale, dx, mvns)
+    return (mass_scale**2) * c
 
 
-def compute_dMdt(weights, scale, dx, mvns, mass_scale=1.):
+def compute_dMdt(weights, scale, dx, mvns):
     scale = scale @ dx
-    Ms = np.array([(mass_scale**2) * comp.cov_inv for comp in mvns])
+    Ms = np.array([comp.cov_inv for comp in mvns])
     dMdt = Ms.T @ (weights * scale)
     return dMdt
 
 
-def compute_dTdx(weights, scale, dx, mvns, mass_scale=1.):
+def compute_dTdx(weights, scale, dx, mvns):
     dTdx = np.zeros_like(dx)
     for k in range(len(mvns)):
-        dTdx += weights[k] * (dx.T @ ((mass_scale**2) * mvns[k].cov_inv) @ dx) * scale[k]
+        dTdx += weights[k] * (dx.T @ mvns[k].cov_inv @ dx) * scale[k]
     return dTdx
 
 

diff --git a/tprmp/models/rmp.py b/tprmp/models/rmp.py
@@ -40,20 +40,17 @@ def compute_potential_term(weights, phi0, x, mvns, **kwargs):
     mean_pull = weights.T @ pulls
     for k in range(num_comp):
         Ps += weights[k] * phi[k] * (pulls[k] - mean_pull)
+        v = manifold.log_map(x, base=mvns[k].mean)
+        norm = np.sqrt((stiff_scale**2) * v.T @ pulls[k])
         if potential_method == 'quadratic':
             Ps += -weights[k] * (stiff_scale**2) * pulls[k]
         elif potential_method == 'tanh':
-            v = manifold.log_map(x, base=mvns[k].mean)
-            norm = np.sqrt((stiff_scale**2) * v.T @ pulls[k])
             Ps += -weights[k] * np.tanh(tau * norm) * (stiff_scale**2) * pulls[k] / norm
-        elif potential_method == 'huber':  # huber potential does not use mvns variance shape
-            v = manifold.log_map(x, base=mvns[k].mean)
-            quadratic = (stiff_scale**2) * v.T @ np.eye(manifold.dim_T) @ v
-            norm = np.sqrt(quadratic)
+        elif potential_method == 'huber':
             if norm <= delta:
-                Ps += -weights[k] * (stiff_scale**2) * v
+                Ps += -weights[k] * (stiff_scale**2) * pulls[k]
             else:
-                Ps += -weights[k] * (stiff_scale**2) * delta * v / norm
+                Ps += -weights[k] * (stiff_scale**2) * delta * pulls[k] / norm
         else:
             raise ValueError(f'Potential method {potential_method} is unrecognized!')
     return Ps
@@ -84,8 +81,7 @@ def compute_potentials(phi0, x, mvns, **kwargs):
         elif potential_method == 'tanh':
             norm = np.sqrt(quadratic)
             P[k] = 1 / tau * (np.exp(tau * norm) + np.exp(-tau * norm))
-        elif potential_method == 'huber':  # huber potential does not use mvns variance shape
-            quadratic = (stiff_scale**2) * v.T @ np.eye(manifold.dim_T) @ v
+        elif potential_method == 'huber':
             norm = np.sqrt(quadratic)
             if norm <= delta:
                 P[k] = 0.5 * quadratic

diff --git a/tprmp/optimizer/dynamics.py b/tprmp/optimizer/dynamics.py
@@ -20,7 +20,7 @@ def optimize_dynamics(tp_gmm, demos, **kwargs):
     d_min = kwargs.get('d_min', 0.)
     energy = kwargs.get('energy', 0.)
     verbose = kwargs.get('verbose', False)
-    phi0 = optimize_potentials(tp_gmm, demos, alpha=alpha, stiff_scale=stiff_scale, tau=tau, delta=delta, potential_method=potential_method, energy=energy, verbose=verbose)
+    phi0 = optimize_potentials(tp_gmm, demos, alpha=alpha, stiff_scale=stiff_scale, mass_scale=mass_scale, tau=tau, delta=delta, potential_method=potential_method, energy=energy, verbose=verbose)
     d0 = optimize_dissipation(tp_gmm, demos, phi0, beta=beta, stiff_scale=stiff_scale, mass_scale=mass_scale,
                               tau=tau, delta=delta, potential_method=potential_method, train_method=train_method, d_min=d_min, verbose=verbose)
     return phi0, d0
@@ -29,6 +29,7 @@ def optimize_dynamics(tp_gmm, demos, **kwargs):
 def optimize_potentials(tp_gmm, demos, **kwargs):
     alpha = kwargs.get('alpha', 1e-5)
     stiff_scale = kwargs.get('stiff_scale', 1.)
+    mass_scale = kwargs.get('mass_scale', 1.)
     tau = kwargs.get('tau', 1.)
     delta = kwargs.get('delta', 1.)
     potential_method = kwargs.get('potential_method', 'quadratic')
@@ -44,11 +45,13 @@ def optimize_potentials(tp_gmm, demos, **kwargs):
         for t in range(x.shape[1]):
             weights = compute_obsrv_prob(x[:, t], mvns)
             f = compute_potential_term(weights, phi0, x[:, t], mvns, stiff_scale=stiff_scale, tau=tau, delta=delta, potential_method=potential_method)
+            M = compute_riemannian_metric(x[:, t], mvns, mass_scale=mass_scale)
+            M_inv = np.linalg.inv(M)
             v = dx[:, t]
             norm_v = np.linalg.norm(v)
             if norm_v > eps:
                 v = v / norm_v
-            loss += (cp.norm(v - f) / x.shape[1])
+            loss += (cp.norm(v - M_inv @ f) / x.shape[1])
     loss /= len(demos)
     if alpha > 0.:
         loss += alpha * cp.pnorm(phi0, p=2)**2  # L2 regularization

diff --git a/tprmp/visualization/dynamics.py b/tprmp/visualization/dynamics.py
@@ -138,21 +138,22 @@ def plot_potential_field(tprmp, frames, **kwargs):
     three_d = kwargs.get('three_d', False)
     margin = kwargs.get('margin', 0.5)
     res = kwargs.get('res', 0.1)
+    max_z = kwargs.get('max_z', 1000)
     new_fig = kwargs.get('new_fig', False)
     show = kwargs.get('show', False)
     mid, ranges = _get_data_ranges(frames, tprmp.model.manifold.get_origin(), margin=margin)
     if new_fig:
         plt.figure()
-    _plot_potential_field_global(tprmp, frames, mid, ranges, plot_gaussian=plot_gaussian, var_scale=var_scale, three_d=three_d, res=res)
+    _plot_potential_field_global(tprmp, frames, mid, ranges, plot_gaussian=plot_gaussian, var_scale=var_scale, three_d=three_d, res=res, max_z=max_z)
     if not only_global:
         if three_d:
             plt.figure()
-        _plot_potential_field_frames(tprmp, frames, ranges, plot_gaussian=plot_gaussian, var_scale=var_scale, three_d=three_d, res=res)
+        _plot_potential_field_frames(tprmp, frames, ranges, plot_gaussian=plot_gaussian, var_scale=var_scale, three_d=three_d, res=res, max_z=max_z)
     if show:
         plt.show()
 
 
-def _plot_potential_field_global(tprmp, frames, mid, ranges, plot_gaussian=True, var_scale=1., three_d=False, res=0.1, alpha=0.7):
+def _plot_potential_field_global(tprmp, frames, mid, ranges, plot_gaussian=True, var_scale=1., three_d=False, res=0.1, max_z=1000, alpha=0.7):
     if three_d:
         ax = plt.subplot(111, projection='3d')
     else:
@@ -164,7 +165,7 @@ def _plot_potential_field_global(tprmp, frames, mid, ranges, plot_gaussian=True,
     tprmp.generate_global_gmm(frames)
     for i in range(X.shape[0]):
         for j in range(X.shape[1]):
-            Z[i, j] = tprmp.compute_potential_field(np.array([X[i, j], Y[i, j]]))
+            Z[i, j] = min(tprmp.compute_potential_field(np.array([X[i, j], Y[i, j]])), max_z)  # constraining plot
     if three_d:
         c = ax.plot_surface(X, Y, Z, cmap='RdBu', vmin=0., vmax=Z.max(), alpha=alpha)
     else:
@@ -177,7 +178,7 @@ def _plot_potential_field_global(tprmp, frames, mid, ranges, plot_gaussian=True,
         _plot_gmm_global(tprmp.model, frames, var_scale=var_scale, three_d=False, new_ax=False)
 
 
-def _plot_potential_field_frames(tprmp, frames, ranges, axs=None, plot_gaussian=True, var_scale=1., three_d=False, res=0.1, alpha=0.7):
+def _plot_potential_field_frames(tprmp, frames, ranges, axs=None, plot_gaussian=True, var_scale=1., three_d=False, res=0.1, max_z=1000, alpha=0.7):
     if axs is None:
         axs = {}
         if three_d:
@@ -196,7 +197,7 @@ def _plot_potential_field_frames(tprmp, frames, ranges, axs=None, plot_gaussian=
         Z[f_key] = np.zeros_like(X)
         for i in range(X.shape[0]):
             for j in range(X.shape[1]):
-                Z[f_key][i, j] = tprmp.compute_potential_field_frame(np.array([X[i, j], Y[i, j]]), f_key)
+                Z[f_key][i, j] = min(tprmp.compute_potential_field_frame(np.array([X[i, j], Y[i, j]]), f_key), max_z)
         z_max = max(z_max, Z[f_key].max())
     for f_key in frames:
         if three_d: