mirror of
https://github.com/wassname/pyrobolearn.git
synced 2026-08-23 12:30:20 +08:00
58 lines
2.0 KiB
Python
58 lines
2.0 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
"""Provide some examples using ProMPs.
|
|
"""
|
|
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
|
|
from pyrobolearn.models.promp.promp import DiscreteProMP, plot_state, plot_proba_state, plot_weighted_basis
|
|
|
|
|
|
# create data and plot it
|
|
N = 8
|
|
t = np.linspace(0., 1., 100)
|
|
# eps = 0.1
|
|
# y = np.array([np.sin(2*np.pi*t) + eps * np.random.rand(len(t)) for _ in range(N)]) # shape: NxT
|
|
# dy = np.array([2*np.pi*np.cos(2*np.pi*t) + eps * np.random.rand(len(t)) for _ in range(N)]) # shape: NxT
|
|
phi = np.random.uniform(low=-1., high=1., size=N)
|
|
y = np.array([np.sin(2 * np.pi * t + phi[i]) for i in range(int(N/2))]) # shape: NxT
|
|
y1 = np.array([np.cos(2 * np.pi * t + phi[i]) for i in range(int(N/2))])
|
|
y = np.vstack((y, y1))
|
|
dy = np.array([2 * np.pi * np.cos(2 * np.pi * t + phi[i]) for i in range(int(N/2))]) # shape: NxT
|
|
dy1 = np.array([2 * np.pi * np.sin(2 * np.pi * t + phi[i]) for i in range(int(N/2))])
|
|
dy = np.vstack((dy, dy1))
|
|
Y = np.dstack((y, dy)) # N,T,2D --> why not N,2D,T
|
|
plot_state(Y, title='Training data')
|
|
plt.show()
|
|
|
|
# create discrete and rhythmic ProMP
|
|
promp = DiscreteProMP(num_dofs=1, num_basis=10, basis_width=1./20)
|
|
|
|
# plot the basis function activations
|
|
plt.plot(promp.Phi(t)[:, :, 0].T)
|
|
plt.title('basis functions')
|
|
plt.show()
|
|
|
|
# plot ProMPs
|
|
y_pred = promp.rollout()
|
|
fig, ax = plt.subplots(1, 2)
|
|
plot_state(y_pred[None], ax=ax, title='ProMP prediction before learning', linewidth=2.) # shape: N,T,2D
|
|
plot_weighted_basis(t, promp, ax=ax)
|
|
plt.show()
|
|
|
|
# learn from demonstrations
|
|
promp.imitate(Y)
|
|
y_pred = promp.rollout()
|
|
fig, ax = plt.subplots(1, 2)
|
|
plot_state(y_pred[None], ax=ax, title='ProMP prediction after learning', linewidth=3.) # N,T,2D
|
|
plot_weighted_basis(t, promp, ax=ax)
|
|
plt.show()
|
|
|
|
method = 'marginal'
|
|
means, covariances = promp.rollout_proba(method=method, return_gaussian=False)
|
|
fig, ax = plt.subplots(1, 2)
|
|
# plot_state(Y, ax=ax, title='Training data')
|
|
plot_proba_state(means, covariances, ax=ax, title='ProMP prediction after learning', linewidth=3.)
|
|
plt.show()
|