Skip to content

Commit

Permalink
Implement huber potential
Browse files Browse the repository at this point in the history 
Signed-off-by: An Thai Le <an.thai.le97@gmail.com>
  • Loading branch information
@anindex
anindex committed Aug 29, 2021
1 parent 3783c4b commit 2a24455
Show file tree
Hide file tree
Showing 10 changed files with 102 additions and 50 deletions.
Binary file added data/tasks/test/demos/test3.p
View file
Binary file not shown.
Binary file added data/tasks/test/models/dynamics_test3.p
View file
Binary file not shown.
Binary file added data/tasks/test/models/tphsmm_test3.p
View file
Binary file not shown.
20 changes: 11 additions & 9 deletions scripts/test_tprmp_2d.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,24 +19,26 @@
description='Example run: python test_tprmp.py test.p')
parser.add_argument('--loading', help='Load or not', type=bool, default=True)
parser.add_argument('--task', help='The task folder', type=str, default='test')
parser.add_argument('--data', help='The data file', type=str, default='test2.p')
parser.add_argument('--data', help='The data file', type=str, default='test3.p')
args = parser.parse_args()

DATA_DIR = join(ROOT_DIR, 'data', 'tasks', args.task, 'demos')
data_file = join(DATA_DIR, args.data)
# parameters
oversteps = 1000
dt = 0.01
NUM_COMP = 30
NUM_COMP = 5
alpha, beta = 0., 0.
stiff_scale = 1.
mass_scale = 0.1
tau = 0.5
potential_method = 'tanh'
train_method = 'match_accel'
delta = 2.
potential_method = 'huber'
train_method = 'match_energy'
d_min = 0.
d_scale = 1.
energy = 38.
var_scale = 20.
energy = 0.
var_scale = 2.
res = 0.05
margin = 0.2
verbose = False
Expand All @@ -57,13 +59,13 @@
if args.loading:
model = TPRMP.load(args.task, model_name=args.data)
else:
model = TPRMP(num_comp=NUM_COMP, name=args.task, stiff_scale=stiff_scale, var_scale=var_scale, tau=tau, potential_method=potential_method, d_scale=d_scale)
model.train(demos, alpha=alpha, beta=beta, d_min=d_min, energy=energy, verbose=verbose)
model = TPRMP(num_comp=NUM_COMP, name=args.task, stiff_scale=stiff_scale, mass_scale=mass_scale, var_scale=var_scale, tau=tau, delta=delta, potential_method=potential_method, d_scale=d_scale)
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_dissipation_field(model, frames, only_global=True, margin=margin, var_scale=var_scale, res=res, new_fig=True, show=True)
# execution
x0, dx0 = demos[0].traj[:, 0], np.zeros(2)
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()
12 changes: 6 additions & 6 deletions tprmp/models/coriolis.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,23 +3,23 @@
from tprmp.models.rmp import compute_obsrv_prob


def compute_coriolis_force(x, dx, mvns):
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) @ dx - 0.5 * compute_dTdx(weights, scale, dx, 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)


def compute_dMdt(weights, scale, dx, mvns):
def compute_dMdt(weights, scale, dx, mvns, mass_scale=1.):
scale = scale @ dx
Ms = np.array([comp.cov_inv for comp in mvns])
Ms = np.array([(mass_scale**2) * comp.cov_inv for comp in mvns])
dMdt = Ms.T @ (weights * scale)
return dMdt


def compute_dTdx(weights, scale, dx, mvns):
def compute_dTdx(weights, scale, dx, mvns, mass_scale=1.):
dTdx = np.zeros_like(dx)
for k in range(len(mvns)):
dTdx += weights[k] * (dx.T @ mvns[k].cov_inv @ dx) * scale[k]
dTdx += weights[k] * (dx.T @ ((mass_scale**2) * mvns[k].cov_inv) @ dx) * scale[k]
return dTdx


Expand Down
51 changes: 39 additions & 12 deletions tprmp/models/rmp.py
View file
Original file line number Diff line number Diff line change
@@ -1,29 +1,36 @@
import numpy as np


def compute_riemannian_metric(x, mvns, eps=1e-5):
def compute_riemannian_metric(x, mvns, mass_scale=1., eps=1e-5):
weights = compute_obsrv_prob(x, mvns)
Ms = np.array([comp.cov_inv for comp in mvns])
Ms = np.array([(mass_scale**2) * comp.cov_inv for comp in mvns])
M = Ms.T @ weights if weights.sum() > 0. else np.eye(mvns[0].manifold.dim_T)
return M + eps * np.eye(mvns[0].manifold.dim_T)


def compute_hamiltonian(phi0, x, dx, mvns, stiff_scale=1., tau=1., potential_method='quadratic', manifold=None):
def compute_hamiltonian(phi0, x, dx, mvns, **kwargs):
manifold = kwargs.get('manifold', None)
mass_scale = kwargs.get('mass_scale', 1.)
weights = compute_obsrv_prob(x, mvns, manifold=manifold)
phi = compute_potentials(phi0, x, mvns, stiff_scale=stiff_scale, tau=tau, potential_method=potential_method, manifold=manifold)
M = compute_riemannian_metric(x, mvns)
phi = compute_potentials(phi0, x, mvns, **kwargs)
M = compute_riemannian_metric(x, mvns, mass_scale=mass_scale)
T = 0.5 * dx.T @ M @ dx
Phi = weights.T @ phi
return T + Phi


def compute_policy(phi0, d0, x, dx, mvns, stiff_scale=1., tau=1., potential_method='quadratic'):
def compute_policy(phi0, d0, x, dx, mvns, **kwargs):
weights = compute_obsrv_prob(x, mvns)
return compute_potential_term(weights, phi0, x, mvns, stiff_scale=stiff_scale, tau=tau, potential_method=potential_method) + compute_dissipation_term(weights, d0, dx)
return compute_potential_term(weights, phi0, x, mvns, **kwargs) + compute_dissipation_term(weights, d0, dx)


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)
def compute_potential_term(weights, phi0, x, mvns, **kwargs):
stiff_scale = kwargs.get('stiff_scale', 1.)
tau = kwargs.get('tau', 1.)
delta = kwargs.get('delta', 0.1)
potential_method = kwargs.get('potential_method', 'quadratic')
manifold = kwargs.get('manifold', None)
phi = compute_potentials(phi0, x, mvns, **kwargs)
num_comp = len(mvns)
manifold = mvns[0].manifold
Ps = np.zeros(manifold.dim_T)
Expand All @@ -39,6 +46,14 @@ def compute_potential_term(weights, phi0, x, mvns, stiff_scale=1., tau=1., poten
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)
if norm <= delta:
Ps += -weights[k] * (stiff_scale**2) * v
else:
Ps += -weights[k] * (stiff_scale**2) * delta * v / norm
else:
raise ValueError(f'Potential method {potential_method} is unrecognized!')
return Ps
Expand All @@ -49,7 +64,12 @@ def compute_dissipation_term(weights, d0, dx):
return Ds


def compute_potentials(phi0, x, mvns, stiff_scale=1., tau=1., potential_method='quadratic', manifold=None):
def compute_potentials(phi0, x, mvns, **kwargs):
stiff_scale = kwargs.get('stiff_scale', 1.)
tau = kwargs.get('tau', 1.)
delta = kwargs.get('delta', 0.1)
potential_method = kwargs.get('potential_method', 'quadratic')
manifold = kwargs.get('manifold', None)
num_comp = len(mvns)
P = np.zeros(num_comp)
if manifold is None:
Expand All @@ -58,12 +78,19 @@ def compute_potentials(phi0, x, mvns, stiff_scale=1., tau=1., potential_method='
for k in range(num_comp):
comp = mvns[k]
v = manifold.log_map(x[:d], base=comp.mean[:d])
quadratic = v.T @ ((stiff_scale**2) * comp.cov_inv[:d, :d]) @ v
quadratic = (stiff_scale**2) * v.T @ comp.cov_inv[:d, :d] @ v
if potential_method == 'quadratic':
P[k] = quadratic
P[k] = 0.5 * quadratic
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
norm = np.sqrt(quadratic)
if norm <= delta:
P[k] = 0.5 * quadratic
else:
P[k] = delta * (norm - 0.5 * delta)
else:
raise ValueError(f'Potential method {potential_method} is unrecognized!')
phi = phi0 + P
Expand Down
5 changes: 4 additions & 1 deletion tprmp/models/tp_gmm.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,10 @@ def reset_covariance(self, idxs_1, idxs_2):
def scale_covariance(self, scale):
for k in range(self.num_comp):
for f_key in self.frame_names:
self._mvns[k][f_key].cov = self._mvns[k][f_key].cov * scale**2
if isinstance(scale, list):
self._mvns[k][f_key].cov = self._mvns[k][f_key].cov * scale[k]**2
else:
self._mvns[k][f_key].cov = self._mvns[k][f_key].cov * scale**2

def get_local_gmm(self, frame, tag=None):
if frame not in self.frame_names:
Expand Down
27 changes: 15 additions & 12 deletions tprmp/models/tp_rmp.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,9 +25,11 @@ class TPRMP(object):
def __init__(self, **kwargs):
self._sigma = kwargs.pop('sigma', 1.)
self._stiff_scale = kwargs.pop('stiff_scale', 1.)
self._mass_scale = kwargs.pop('mass_scale', 1.)
self._var_scale = kwargs.pop('var_scale', 1.)
self._tau = kwargs.pop('tau', 1.)
self._potential_method = kwargs.pop('potential_method', 'quadratic')
self._delta = kwargs.pop('delta', 0.1)
self._potential_method = kwargs.pop('potential_method', 'tanh')
self._d_scale = kwargs.pop('d_scale', 1.)
self._model = TPHSMM(**kwargs)
self._global_mvns = None
Expand All @@ -40,7 +42,7 @@ def save(self, name=None):
self.model.save(name)
file = join(DATA_PATH, self.model.name, 'models', 'dynamics_' + (name if name is not None else (str(time.time()) + '.p')))
with open(file, 'wb') as f:
pickle.dump({'phi0': self._phi0, 'd0': self._d0, 'stiff_scale': self._stiff_scale, 'var_scale': self._var_scale,
pickle.dump({'phi0': self._phi0, 'd0': self._d0, 'stiff_scale': self._stiff_scale, 'mass_scale': self._mass_scale, 'var_scale': self._var_scale,
'tau': self._tau, 'potential_method': self._potential_method, 'd_scale': self._d_scale}, f)

def generate_global_gmm(self, frames):
Expand All @@ -58,13 +60,13 @@ def rmp(self, x, dx):
"""
Retrieve global RMP.
"""
f = self.compute_global_policy(x, dx) - compute_coriolis_force(x, dx, self._global_mvns)
M = compute_riemannian_metric(x, self._global_mvns)
f = self.compute_global_policy(x, dx) - compute_coriolis_force(x, dx, self._global_mvns, mass_scale=self._mass_scale)
M = compute_riemannian_metric(x, self._global_mvns, mass_scale=self._mass_scale)
return M, f

def compute_global_policy(self, x, dx):
policy = compute_policy(self._phi0, self._d_scale * self._d0, x, dx, self._global_mvns,
stiff_scale=self._stiff_scale, tau=self._tau, potential_method=self._potential_method)
stiff_scale=self._stiff_scale, tau=self._tau, delta=self._delta, potential_method=self._potential_method)
return policy

def compute_frame_weights(self, x, frames, normalized=True, eps=1e-307):
Expand Down Expand Up @@ -95,14 +97,14 @@ def compute_frame_weights(self, x, frames, normalized=True, eps=1e-307):

def compute_potential_field(self, x, manifold=None):
weights = compute_obsrv_prob(x, self._global_mvns, manifold=manifold)
phi = compute_potentials(self._phi0, x, self._global_mvns, stiff_scale=self._stiff_scale, tau=self._tau, potential_method=self._potential_method, manifold=manifold)
phi = compute_potentials(self._phi0, x, self._global_mvns, stiff_scale=self._stiff_scale, tau=self._tau, delta=self._delta, potential_method=self._potential_method, manifold=manifold)
Phi = weights.T @ phi
return Phi

def compute_potential_field_frame(self, lx, frame, manifold=None):
mvns = self.model.get_local_gmm(frame)
weights = compute_obsrv_prob(lx, mvns, manifold=manifold)
phi = compute_potentials(self.phi0[frame], lx, mvns, stiff_scale=self._stiff_scale, tau=self._tau, potential_method=self._potential_method, manifold=manifold)
phi = compute_potentials(self.phi0[frame], lx, mvns, stiff_scale=self._stiff_scale, tau=self._tau, delta=self._delta, potential_method=self._potential_method, manifold=manifold)
Phi = weights.T @ phi
return Phi

Expand Down Expand Up @@ -130,11 +132,11 @@ def train(self, demos, **kwargs):
if 'S' in self.model.manifold.name: # decouple orientation and position
pos_idx, quat_idx = self.model.manifold.get_pos_quat_indices(tangent=True)
self.model.reset_covariance(pos_idx, quat_idx)
if self._var_scale > 1.:
if isinstance(self._var_scale, list) or self._var_scale > 1.:
self.model.scale_covariance(self._var_scale)
# train dynamics
self._phi0, self._d0 = optimize_dynamics(self.model, demos, alpha=alpha, beta=beta, train_method=train_method,
stiff_scale=self._stiff_scale, tau=self._tau, potential_method=self._potential_method, d_min=d_min, energy=energy, verbose=verbose)
self._phi0, self._d0 = optimize_dynamics(self.model, demos, alpha=alpha, beta=beta, train_method=train_method, mass_scale=self._mass_scale,
stiff_scale=self._stiff_scale, tau=self._tau, delta=self._delta, potential_method=self._potential_method, d_min=d_min, energy=energy, verbose=verbose)
# train local Riemannian metrics TODO: RiemannianNetwork is still under consideration
# self._R_net = optimize_riemannian_metric(self, demos, **kwargs)

Expand All @@ -152,8 +154,9 @@ def load(task_name, model_name='sample.p'):
raise ValueError(f'[TPHSMM]: File {file} does not exist!')
dynamics = load(file)
tprmp._phi0, tprmp._d0 = dynamics['phi0'], dynamics['d0']
tprmp._stiff_scale, tprmp._var_scale, tprmp._d_scale = dynamics['stiff_scale'], dynamics['var_scale'], dynamics['d_scale']
tprmp._potential_method, tprmp._tau = dynamics['potential_method'], dynamics['tau']
tprmp._stiff_scale, tprmp._mass_scale = dynamics.get('stiff_scale', 1.0), dynamics.get('mass_scale', 1.0)
tprmp._var_scale, tprmp._d_scale = dynamics.get('var_scale', 1.0), dynamics.get('d_scale', 1.0)
tprmp._potential_method, tprmp._tau = dynamics.get('potential_method', 'tanh'), dynamics.get('tau', 1.)
return tprmp

@property
Expand Down
21 changes: 15 additions & 6 deletions tprmp/optimizer/dynamics.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,21 +12,25 @@ def optimize_dynamics(tp_gmm, demos, **kwargs):
alpha = kwargs.get('alpha', 1e-5)
beta = kwargs.get('beta', 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')
train_method = kwargs.get('train_method', 'match_accel')
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, potential_method=potential_method, energy=energy, verbose=verbose)
d0 = optimize_dissipation(tp_gmm, demos, phi0, beta=beta, stiff_scale=stiff_scale, tau=tau, potential_method=potential_method, train_method=train_method, d_min=d_min, verbose=verbose)
phi0 = optimize_potentials(tp_gmm, demos, alpha=alpha, stiff_scale=stiff_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


def optimize_potentials(tp_gmm, demos, **kwargs):
alpha = kwargs.get('alpha', 1e-5)
stiff_scale = kwargs.get('stiff_scale', 1.)
tau = kwargs.get('tau', 1.)
delta = kwargs.get('delta', 1.)
potential_method = kwargs.get('potential_method', 'quadratic')
energy = kwargs.get('energy', 0.)
eps = kwargs.get('eps', 1e-4)
Expand All @@ -39,7 +43,7 @@ def optimize_potentials(tp_gmm, demos, **kwargs):
mvns = tp_gmm.generate_global_gmm(demo.get_task_parameters())
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, potential_method=potential_method)
f = compute_potential_term(weights, phi0, x[:, t], mvns, stiff_scale=stiff_scale, tau=tau, delta=delta, potential_method=potential_method)
v = dx[:, t]
norm_v = np.linalg.norm(v)
if norm_v > eps:
Expand All @@ -53,14 +57,17 @@ def optimize_potentials(tp_gmm, demos, **kwargs):
problem.solve(verbose=verbose)
logger.info('Optimizing potential...')
logger.info(f'Status: {problem.status}')
logger.info(f'Final loss: {loss.value}')
logger.info(f'Optimal phi0: {phi0.value}')
return phi0.value


def optimize_dissipation(tp_gmm, demos, phi0, **kwargs):
beta = kwargs.get('beta', 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')
train_method = kwargs.get('train_method', 'match_accel')
d_min = kwargs.get('d_min', 0.)
Expand All @@ -74,12 +81,13 @@ def optimize_dissipation(tp_gmm, demos, phi0, **kwargs):
mvns = tp_gmm.generate_global_gmm(demo.get_task_parameters())
if train_method == 'match_accel':
for t in range(x.shape[1]):
M = compute_riemannian_metric(x[:, t], mvns)
M = compute_riemannian_metric(x[:, t], mvns, mass_scale=mass_scale)
M_inv = np.linalg.inv(M)
f = compute_policy(phi0, d0, x[:, t], dx[:, t], mvns, stiff_scale=stiff_scale, tau=tau, potential_method=potential_method) - compute_coriolis_force(x[:, t], dx[:, t], mvns)
f = compute_policy(phi0, d0, x[:, t], dx[:, t], mvns, stiff_scale=stiff_scale, tau=tau, delta=delta, potential_method=potential_method)
f -= compute_coriolis_force(x[:, t], dx[:, t], mvns, mass_scale=mass_scale)
loss += cp.norm(ddx[:, t] - M_inv @ f)
elif train_method == 'match_energy':
energy = compute_hamiltonian(phi0, x[:, 0], dx[:, 0], mvns, stiff_scale=stiff_scale, tau=tau, potential_method=potential_method)
energy = compute_hamiltonian(phi0, x[:, 0], dx[:, 0], mvns, stiff_scale=stiff_scale, mass_scale=mass_scale, tau=tau, delta=delta, potential_method=potential_method)
d_energy = 0. # d_energy is negative
for t in range(x.shape[1] - 1):
weights = compute_obsrv_prob(x[:, t], mvns)
Expand All @@ -96,6 +104,7 @@ def optimize_dissipation(tp_gmm, demos, phi0, **kwargs):
problem.solve(max_iters=max_iters, verbose=verbose)
logger.info('Optimizing dissipation...')
logger.info(f'Status: {problem.status}')
logger.info(f'Final loss: {loss.value}')
logger.info(f'Optimal d0: {d0.value}')
res = d0.value
except cp.error.SolverError:
Expand Down
Loading

0 comments on commit 2a24455

Please sign in to comment.