diff --git a/pyrobolearn/models/dmp/dmpytorch/README.md b/pyrobolearn/models/dmp/dmpytorch/README.md new file mode 100644 index 0000000..823cdeb --- /dev/null +++ b/pyrobolearn/models/dmp/dmpytorch/README.md @@ -0,0 +1,94 @@ +# DMPyTorch + +This repository contains the code for the *DMPyTorch* library; a PyTorch library for Dynamic Movement Primitives. If you want to use with robots you can have a look at the [`pyrobolearn` framework](https://github.com/robotlearn/pyrobolearn) + +**Warning**: The development of this framework is ongoing, and thus some substantial changes might occur. Sorry for the inconvenience. + + +## Requirements + +The framework has been tested with Python 2.7 and 3.5 on Ubuntu 16.04 and 18.04. + + +## Installation + +1. First download the `pip` Python package manager and create a virtual environment for Python as described in the following link: https://packaging.python.org/guides/installing-using-pip-and-virtualenv/ +On Ubuntu, you can install `pip` and `virtualenv` by typing in the terminal: + +- In Python 2.7: +```bash +sudo apt install python-pip +sudo pip install virtualenv +``` + +- In Python 3.5: +```bash +sudo apt install python3-pip +sudo pip install virtualenv +``` + +You can then create the virtual environment by typing: +```bash +virtualenv -p /usr/bin/python +# activate the virtual environment +source /bin/activate +``` +where `` is the python version you want to use (select between `2.7` or `3.5`), and `` is a name of your choice for the virtual environment. For instance, it can be `py2.7` or `py3.5`. + +To deactivate the virtual environment, just type: +```bash +deactivate +``` + +2. clone this repository and install the requirements by executing the setup.py + +In Python 2.7 or 3.5: +```bash +git clone https://github.com/robotlearn/dmpytorch +cd dmpytorch +pip install numpy +pip install -e . # this will install dmpytorch as well as the required packages (so no need for: pip install -r requirements.txt) +``` + +Depending on your computer configuration and the python version you use, you might need to install also the following packages through `apt-get`: +```bash +sudo apt install python-tk # if python 2.7 +sudo apt install python3-tk # if python 3.5 +``` + +## How to use it? + +Check the `README.md` file in the `examples` folder. + +## Citation + +If you use `dmpytorch`, please cite: +``` +@misc{delhaisse2019dmpytorch, + author = {Delhaisse, Brian and Rozo, Leonel, and Caldwell, Darwin}, + title = {DMPyTorch: a PyTorch Library for Dynamic Movement Primitives}, + howpublished = {\url{https://github.com/robotlearn/dmpytorch}}, + year=2019, +} +``` + +## Acknowledgements + +Parts of the code were inspired by the work from [Travis DeWolf](https://github.com/studywolf) and his library [`pydmps`](https://github.com/studywolf/pydmps). His blog which explains DMPs pretty well can be found [here](https://studywolf.wordpress.com/category/robotics/dynamic-movement-primitive/). + + +## References + +1. "Dynamical movement primitives: Learning attractor models for motor behaviors", Ijspeert et al., 2013 +2. PyDMPs (from DeWolf, 2013): https://github.com/studywolf/pydmps +3. "Biologically-inspired Dynamical Systems for Movement Generation: Automatic Real-time Goal Adaptation and Obstacle Avoidance", Hoffmann et al., 2009 +4. "Policy Search for Motor Primitives in Robotics", Kober et al., 2010 + + +## TODO + +- [ ] implement "Orientation in Cartesian Space Dynamic Movement Primitives", Ude et al., 2014 +- [ ] implement "Action Sequencing using Dynamic Movement Primitives", Nemec et al., 2011 +- [ ] implement "A Generalized Path Integral Control Approach to Reinforcement Learning", Theodorou et al., 2010 +- [ ] test DMP with NN + diff --git a/pyrobolearn/models/dmp/dmpytorch/__init__.py b/pyrobolearn/models/dmp/dmpytorch/__init__.py new file mode 100644 index 0000000..e2c2235 --- /dev/null +++ b/pyrobolearn/models/dmp/dmpytorch/__init__.py @@ -0,0 +1,15 @@ + +# import canonical systems +from .canonical_systems import * + +# import basis functions +from .basis_functions import * + +# import forcing terms +from .forcing_terms import * + +# import dynamic movement primitives +from .dmp import * +from .discrete_dmp import * +from .rhythmic_dmp import * +from .biodiscrete_dmp import * diff --git a/pyrobolearn/models/dmp/dmpytorch/basis_functions.py b/pyrobolearn/models/dmp/dmpytorch/basis_functions.py new file mode 100644 index 0000000..a2e6bc1 --- /dev/null +++ b/pyrobolearn/models/dmp/dmpytorch/basis_functions.py @@ -0,0 +1,94 @@ +#!/usr/bin/env python +"""Define basis functions used in the forcing terms in dynamic movement primitives + +This file implements basis functions used for discrete and rhythmic dynamic movement primitives. +""" + +from abc import ABCMeta, abstractmethod +import torch + +__author__ = "Brian Delhaisse" +__copyright__ = "Copyright 2018, PyRoboLearn" +__credits__ = ["Brian Delhaisse"] +__license__ = "MIT" +__version__ = "1.0.0" +__maintainer__ = "Brian Delhaisse" +__email__ = "briandelhaisse@gmail.com" +__status__ = "Development" + + +class BF(torch.nn.Module): + r"""Basis function used in the forcing terms + """ + __metaclass__ = ABCMeta + + def __init__(self): + super(BF, self).__init__() + + +class EBF(BF): + r"""Exponential basis function + + This basis function is given by the formula: + + .. math:: \psi(s) = \exp \left( - \frac{1}{2 \sigma^2} (s - c)^2 \right) + + where :math:`c` is the center, and :math:`\sigma` is the width of a normal distribution. + + This is often used for discrete DMPs. + """ + def __init__(self, center=0, sigma=1., h=None): + """Initialize basis function + + Args: + center (float, torch.Tensor): center of the distribution + sigma (float, torch.Tensor): width of the distribution + h (float, torch.Tensor): concentration/precision of the basis fct (h = 1/(2*\sigma^2)). + if h is not provided, it will check sigma. + """ + super(EBF, self).__init__() + + if isinstance(center, torch.Tensor): + pass + + self.c = center + if h is None: + self.h = 1. / (2*sigma**2) # measure the concentration + else: + self.h = h + + def forward(self, s): + if isinstance(s, torch.Tensor): + s = s[:, None] + return torch.exp(-self.h * (s - self.c)**2) + + +class CBF(BF): + r"""Circular basis function (aka von Mises basis function) + + This basis function is given by the formula: + + .. math:: \psi(s) = \exp \left( h (\cos(s - c) - 1) \right) + + where :math:`c` is the center, and :math:`h` is a measure of concentration. + + This is often used for rhythmic DMPs. + """ + + def __init__(self, center=0, h=1.): + """Initialize basis function + + Args: + center (float, torch.Tensor): center of the basis fct + h (float, torch.Tensor): concentration/precision of the basis fct + """ + super(CBF, self).__init__() + self.c = center + self.h = h + + def forward(self, s): + if isinstance(s, torch.Tensor): + s = s[:, None] + # return torch.exp(self.h * torch.cos(s - self.c) - 1) # this is bad as it is not bounded as we increase + # the number of basis functions. + return torch.exp(self.h * torch.cos(s - self.c) - self.h) diff --git a/pyrobolearn/models/dmp/dmpytorch/biodiscrete_dmp.py b/pyrobolearn/models/dmp/dmpytorch/biodiscrete_dmp.py new file mode 100644 index 0000000..68945cb --- /dev/null +++ b/pyrobolearn/models/dmp/dmpytorch/biodiscrete_dmp.py @@ -0,0 +1,268 @@ +#!/usr/bin/env python +"""Define the biologically-inspired discrete dynamic movement primitive (as described in [1,2]) + +References: + [1] "Biologically-inspired Dynamical Systems for Movement Generation: Automatic Real-time Goal Adaptation + and Obstacle Avoidance", Hoffmann et al., 2009 + [2] "Learning and Generalization of Motor Skills by Learning from Demonstration", Pastor et al., 2009 +""" + +import numpy as np +import torch + +from pyrobolearn.models.dmp.dmpytorch.discrete_dmp import DiscreteDMP + +__author__ = "Brian Delhaisse" +__copyright__ = "Copyright 2018, PyRoboLearn" +__credits__ = ["Brian Delhaisse"] +__license__ = "MIT" +__version__ = "1.0.0" +__maintainer__ = "Brian Delhaisse" +__email__ = "briandelhaisse@gmail.com" +__status__ = "Development" + + +class BioDiscreteDMP(DiscreteDMP): + r"""Biologically-inspired Discrete DMPs + + One of the main problems with the initial DMP formulation is when some goal coordinates coincide with their + corresponding initial position coordinates, it results in an inappropriate rescaling when displacing a little bit + the goal. + + To deal with this problem, a new formulation of the transformation system was proposed in [2] and is given by: + + .. math:: \tau^2 \ddot{y} = K (g - y) - D \tau \dot{y} - K(g - y_0)s + K f(s) + + where :math:`\tau` is a scaling factor that allows to slow down or speed up the reproduced movement, :math:`K` + is the stiffness coefficient, :math:`D` is the damping coefficient, :math:`y, \dot{y}, \ddot{y}` are the position, + velocity, and acceleration of a DoF, and :math:`f(s)` is the non-linear forcing term. + + The forcing term is expressed as: + + .. math:: f(s) = \frac{\sum_i \psi_i(s) w_i}{ \sum_j \psi_j(s)} s + + Properties (from [2]): + * Invariant under affine transformation + * Movement generalization to new targets + + References: + [1] "Dynamical movement primitives: Learning attractor models for motor behaviors", Ijspeert et al., 2013 + [2] "Biologically-inspired Dynamical Systems for Movement Generation: Automatic Real-time Goal Adaptation + and Obstacle Avoidance", Hoffmann et al., 2009 + [3] "Learning and Generalization of Motor Skills by Learning from Demonstration", Pastor et al., 2009 + """ + + def __init__(self, num_dmps, num_basis, dt=0.01, y0=0, goal=1, + forces=None, stiffness=None, damping=None): + """Initialize the discrete DMP + + Args: + num_dmps (int): number of DMPs + num_basis (int): number of basis functions + dt (float): step integration for Euler's method + y0 (float, np.array): initial position(s) + goal (float, np.array): goal(s) + forces (list, ForcingTerm): the forcing terms (which can have different basis functions) + stiffness (float): stiffness coefficient + damping (float): damping coefficient + """ + # if stiffness is None and damping is None: + # # from paper [2] + # stiffness = 150 * np.ones(num_dmps) + # damping = 2 * np.sqrt(stiffness) + + self.cst = 0.75 # this depends on the K and D value + + super(BioDiscreteDMP, self).__init__(num_dmps, num_basis, dt=dt, y0=y0, goal=goal, + forces=forces, stiffness=stiffness, damping=damping) + + def step(self, s=None, tau=1.0, error=0.0, forcing_term=None, new_goal=None, external_force=None, + rescale_force=True): + """Run the DMP transformation system for a single time step. + + Args: + s (None, float): the phase value. If None, it will use the canonical system. + tau (float): Increase tau to make the system slower, and decrease it to make it faster + error (float): optional system feedback + forcing_term (np.ndarray): if given, it will replace the forcing term (shape [dmp,]) + new_goal (np.ndarray): new goal (of shape [num_dmps,]) + """ + + # system feedback + error_coupling = 1.0 / (1.0 + error) + + # get phase from canonical system + if s is None: + s = self.cs.step(tau=tau, error_coupling=error_coupling) + elif not isinstance(s, (float, int)): + raise TypeError("Expecting the phase 's' to be a float or integer. Instead, I got {}".format(type(s))) + + # check if same phase as before + if s == self.prev_s: + return self.y, self.dy, self.ddy + + if new_goal is None: + new_goal = self.goal + else: + new_goal = new_goal + self.cst * (new_goal - self.goal) + + # save previous position and velocity + prev_y, prev_dy = self.y.clone(), self.dy.clone() + + # for each DMP, solve transformation system equation using Euler's method + for d in range(self.num_dmps): + + # compute forcing term + if forcing_term is None: + f = self.forces[d](s) + self.K[d] * s * (self.goal[d] - new_goal[d]) + else: + f = forcing_term[d] + + # DMP acceleration + self.ddy[d] = self.K[d]/(tau**2) * (new_goal[d] - self.y[d]) - self.D[d]/tau * self.dy[d] + f/(tau**2) + if external_force is not None: + self.ddy[d] += external_force[d] + self.dy[d] += self.ddy[d] / tau * self.dt * error_coupling + self.y[d] += self.dy[d] * self.dt * error_coupling + + # return self.y, self.dy, self.ddy + return prev_y, prev_dy, self.ddy + + def _check_offset(self): + """No need to check for an offset with this class""" + pass + + def generate_goal(self, y0=None, dy0=None, ddy0=None, f0=None): + """ + Generate the goal from the initial positions, velocities, accelerations, and forces. + + Args: + y0 (float[M], None): initial positions. If None, it will take the default initial positions. + dy0 (float[M], None): initial velocities. If None, it will take the default initial velocities. + ddy0 (float[M], None): initial accelerations. If None, it will take the default initial accerelations. + f0 (float[M], None): initial forcing terms. If None, it will compute it based on the learned weights. + You can also give `dmp.f_target[:,0]` to get the correct goal. + + Returns: + float[M]: goal position for each DMP. + """ + if y0 is None: + y0 = self.y0 + if dy0 is None: + dy0 = self.dy0 + if ddy0 is None: + ddy0 = self.ddy0 + if f0 is None: + s0 = self.cs.init_phase + f0 = self.get_forcing_term(s0) + + return 1/self.K * (ddy0 + self.D * dy0 + self.K * y0 - self.K * f0) + + +# Tests +if __name__ == '__main__': + import matplotlib.pyplot as plt + + # tests basis functions + num_basis = 100 + + # Test Biologically-inspired DMP + t = torch.linspace(0., 1., 100) + y_d = torch.sin(np.pi * t) + new_goal = torch.tensor([[0.8, -0.25], + [0.8, 0.25], + [1.2, -0.25]]) + + discrete_dmp = DiscreteDMP(num_dmps=2, num_basis=num_basis) + discrete_dmp.imitate(torch.stack([t, y_d])) + y, dy, ddy = discrete_dmp.rollout() + init_points = torch.stack([discrete_dmp.y0, discrete_dmp.goal]) + # print(discrete_dmp.generate_goal()) + # print(discrete_dmp.generate_goal(f0=discrete_dmp.f_target[:,0])) + y = y.detach().numpy() # convert to numpy + + # check with standard discrete DMP when rescaling the goal + plt.subplot(1, 3, 1) + plt.title('Initial discrete DMP') + plt.scatter(init_points[:, 0], init_points[:, 1], color='b') + plt.scatter(new_goal[:, 0], new_goal[:, 1], color='r') + plt.plot(y[0], y[1], 'b', label='original') + + plt.subplot(1, 3, 2) + plt.title('Rescaled discrete DMP') + plt.scatter(init_points[:, 0], init_points[:, 1], color='b') + plt.scatter(new_goal[:, 0], new_goal[:, 1], color='r') + plt.plot(y[0], y[1], 'b', label='original') + for g in new_goal: + y, dy, ddy = discrete_dmp.rollout(new_goal=g) + plt.plot(y[0].detach().numpy(), y[1].detach().numpy(), 'g', label='scaled') + plt.legend(['original', 'scaled']) + + # change goal with biologically-inspired DMP + new_goal = torch.tensor([[0.8, -0.25], + [0.8, 0.25], + [0.4, 0.1], + [5., 0.15], + [1.2, -0.25], + [-0.8, 0.1], + [-0.8, -0.25], + [5., -0.25]]) + bio_dmp = BioDiscreteDMP(num_dmps=2, num_basis=num_basis) + bio_dmp.imitate(torch.stack([t, y_d])) + y, dy, ddy = bio_dmp.rollout() + init_points = torch.stack([bio_dmp.y0, bio_dmp.goal]) + y = y.detach().numpy() # convert to numpy + + plt.subplot(1, 3, 3) + plt.title('Biologically-inspired DMP') + plt.scatter(init_points[:, 0], init_points[:, 1], color='b') + plt.scatter(new_goal[:, 0], new_goal[:, 1], color='r') + plt.plot(y[0], y[1], 'b', label='original') + for g in new_goal: + y, dy, ddy = bio_dmp.rollout(new_goal=g) + y = y.detach().numpy() # convert to numpy + plt.plot(y[0], y[1], 'g', label='scaled') + plt.legend(['original', 'scaled']) + plt.show() + + # changing goal at the middle + y_list = [] + for g in new_goal: + bio_dmp.reset() + y_traj = torch.zeros(2, 100) + for t in range(100): + if t < 30: + y, dy, ddy = bio_dmp.step() + else: + y, dy, ddy = bio_dmp.step(new_goal=g) + y_traj[:, t] = y + y_list.append(y_traj) + for y in y_list: + y = y.detach().numpy() # convert to numpy + plt.plot(y[0], y[1]) + plt.scatter(bio_dmp.y0[0], bio_dmp.y0[1], color='b') + plt.scatter(new_goal[:, 0], new_goal[:, 1], color='r') + plt.title('change goal at the middle') + plt.show() + + # changing goal at the middle but with a moving goal + g = torch.cat((torch.arange(1.0, 2.0, 0.1).reshape(10, -1), + torch.arange(0.0, 1.0, 0.1).reshape(10, -1)), dim=1) + + bio_dmp.reset() + y_traj = np.zeros((2, 100)) + y_list = [] + for t in range(100): + y, dy, ddy = bio_dmp.step(new_goal=g[int(t/10)]) + y = y.detach().numpy() # convert to numpy + y_traj[:, t] = y + if (t % 10) == 0: + y_list.append(y) + y_list = np.array(y_list) + + plt.plot(y_traj[0], y_traj[1]) + plt.scatter(bio_dmp.y0[0], bio_dmp.y0[1], color='b') + plt.scatter(g[:, 0], g[:, 1], color='r') + plt.scatter(y_list[:, 0], y_list[:, 1], color='g') + plt.title('moving goal') + plt.show() diff --git a/pyrobolearn/models/dmp/dmpytorch/canonical_systems.py b/pyrobolearn/models/dmp/dmpytorch/canonical_systems.py new file mode 100644 index 0000000..31057d9 --- /dev/null +++ b/pyrobolearn/models/dmp/dmpytorch/canonical_systems.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python +"""Define canonical systems for dynamic movement primitives + +This file implements canonical systems for discrete and rhythmic dynamic movement primitives. +""" + +from abc import ABCMeta, abstractmethod +import numpy as np +import torch + + +__author__ = "Brian Delhaisse" +__copyright__ = "Copyright 2018, PyRoboLearn" +__credits__ = ["Brian Delhaisse"] +__license__ = "MIT" +__version__ = "1.0.0" +__maintainer__ = "Brian Delhaisse" +__email__ = "briandelhaisse@gmail.com" +__status__ = "Development" + + +class CS(object): + r"""Canonical System. + + A canonical system (CS) drives a dynamic movement primitive (DMP) by providing a phase variable [1]. + The phase variable was introduced to avoid an explicit dependency with time in the DMP equations. Canonical + systems can be categorized in two main categories: + * discrete CS: used for discrete movements (such as reaching, pushing/pulling, hitting, etc) + * rhythmic CS: used for rhythmic movements (such as walking, running, dribbling, sewing, flipping a pancake, etc) + + Each of these systems are described by differential equations which are solved using Euler's method. + See their corresponding classes `DiscreteCS` and `RhythmicCS` for more information. + + References: + [1] "Dynamical movement primitives: Learning attractor models for motor behaviors", Ijspeert et al., 2013 + """ + + __metaclass__ = ABCMeta + + def __init__(self, dt=0.01, T=1.): + """Initialize the canonical system. + + Args: + dt (float): the time step used in Euler's method when solving the differential equation + A very small step will lead to a better accuracy but will take more time. + """ + # set variables + self.dt = dt + self.T = T + self.timesteps = int(T / self.dt) + # rescale integration step (same as torch.linspace(0.,T.,timesteps) instead of torch.arange(0,T,dt)) + self.dt = self.T / (self.timesteps - 1.) + + self.init_phase = 1.0 + self.s = 1.0 + + # reset the phase variable + self.reset() + + @abstractmethod + def step(self, tau=1.0, error_coupling=1.0): + """Perform a step using Euler's method. This needs to be implemented in the child classes.""" + raise NotImplementedError + + def reset(self): + """Reset the phase variable""" + self.s = self.init_phase + return self.s + + def rollout(self, tau=1.0, error_coupling=1.0): + """Generate phase variable in an open loop fashion. + + Args: + tau (float): Increase tau to make the system slower, and decrease it to make it faster + error_coupling (float): slow down if the error is > 1 + """ + timesteps = int(self.timesteps * tau) + self.s_track = torch.zeros(timesteps) + + # reset + self.reset() + + # roll + for t in range(timesteps): + self.s_track[t] = self.s + self.step(tau, error_coupling) + + return self.s_track + + +class DiscreteCS(CS): + r"""Discrete Canonical System. + + The discrete canonical system drives the various DMPs by providing the phase variable at each time step, and is + given by: + + .. math:: \tau \dot{s} = - \alpha_s s + + where :math:`\tau` is a scaling factor that allows to slow down or speed up the movement, :math:`s` is the phase + variable that drives the DMP, and :math:`\alpha_s` is a predefined constant. + This differential equation is solved using Euler's method. + + This version is used for discrete movements, where :math:`s` starts from 1 and converge to 0 as time progresses. + The phase variable was introduced to avoid an explicit dependency of time in the DMP equations. + + References: + [1] "Dynamical movement primitives: Learning attractor models for motor behaviors", Ijspeert et al., 2013 + """ + + def __init__(self, alpha_s=1, dt=0.01): + super(DiscreteCS, self).__init__(dt=dt, T=1.0) + self.alpha_s = alpha_s + + def reset(self): + """Reset the phase variable""" + self.s = self.init_phase + return self.s + + def step(self, tau=1.0, error_coupling=1.0): + """Generate phase value for discrete movements. + + The phase variable :math:`s` is generated by solving :math:`\tau \dot{s} = - \alpha_s s` using Euler's method. + This phase decays from 1 to 0. + + Args: + tau (float): Increase tau to make the system slower, and decrease it to make it faster + error_coupling (float): slow down if the error is > 1 + + Returns: + float: phase value + """ + s = self.s + self.s += (-self.alpha_s/tau * self.s * error_coupling) * self.dt + # return self.s + return s + + +class RhythmicCS(CS): + r"""Rhythmic Canonical System. + + The rhythmic canonical system drives the various DMPs by providing a phase variable that is periodic [1]. It is + used for rhythmic movements (such as walking, dribbling, sewing, etc.) and is given by: + + .. math:: \tau \dot{s} = 1 + + where :math:`\tau` is a scaling factor that allows to slow down or speed up the movement, :math:`s` is the phase + variable that drives the DMP. This differential equation is solved using Euler's method. + + Rhythmic canonical systems can also be coupled with each other as done in [2] to synchronize various DMPs. + + References: + [1] "Dynamical movement primitives: Learning attractor models for motor behaviors", Ijspeert et al., 2013 + [2] "A Framework for Learning Biped Locomotion with Dynamical Movement Primitives", Nakanishi et al., 2004 + """ + + def __init__(self, dt=0.01): + super(RhythmicCS, self).__init__(dt=dt, T=2*np.pi) + self.init_phase = 0.0 + + def reset(self): + """Reset the phase variable""" + self.s = self.init_phase + return self.s + + def step(self, tau=1.0, error_coupling=1.0): + r"""Generate phase value for rhythmic movements. + + The phase variable :math:`s` is generated by solving :math:`\tau \dot{s} = 1` using Euler's method. + + Args: + tau (float): Increase tau to make the system slower, and decrease it to make it faster + error_coupling (float): slow down if the error is > 1 + + Returns: + float: phase value + """ + s = self.s + self.s += (1./tau * error_coupling) * self.dt + # return self.s + return s + + +class RhythmicNetworkCS(CS): + r"""Rhythmic Network CS. + + In this version, instead of having one canonical system that drives all the various DMPs, we have several + canonical systems coupled with each other, and where each one of them is associated to a particular DMP. + + The evolution of the phase variable :math:`\phi` of the system :math:`i` is given by: + + .. math:: \dot{\phi}_i = \omega_i + \sum_j a_j w_{ij} \sin(\phi_j - \phi_i - \varphi_{ij}) + + where :math:`\omega` is the desired angular velocity (desired frequency), :math:`w_{ij}` are the coupling weights, + :math:`\varphi_{ij}` are the phase biases, and :math:`a_j` are the amplitudes of the other systems :math:`j`. + This formulation is similar to Central Pattern Generators (CPGs), see [3]. + + References: + [1] "Dynamical movement primitives: Learning attractor models for motor behaviors", Ijspeert et al., 2013 + [2] "A Framework for Learning Biped Locomotion with Dynamical Movement Primitives", Nakanishi et al., 2004 + [3] "Central pattern generators for locomotion control in animals and robots: a review", Ijspeert, 2008 + """ + def __init__(self, dt=0.01): + super(RhythmicNetworkCS, self).__init__(dt=dt) + + +# Tests +if __name__ == '__main__': + import matplotlib.pyplot as plt + + # tests canonical systems + discrete_cs = DiscreteCS() + rhythmic_cs = RhythmicCS() + + # check tau + plt.subplot(1, 2, 1) + plt.title('Discrete CS') + for tau in [1., 0.5, 2.]: + rollout = discrete_cs.rollout(tau=tau).numpy() + plt.plot(np.linspace(0, 1., len(rollout)), rollout, label='tau='+str(tau)) + plt.legend() + + plt.subplot(1, 2, 2) + plt.title('Rhythmic CS') + for tau in [1., 0.5, 2.]: + rollout = rhythmic_cs.rollout(tau=tau).numpy() + plt.plot(np.linspace(0, 1., len(rollout)), rollout, label='tau='+str(tau)) + plt.legend() + plt.show() diff --git a/pyrobolearn/models/dmp/dmpytorch/discrete_dmp.py b/pyrobolearn/models/dmp/dmpytorch/discrete_dmp.py new file mode 100644 index 0000000..b2843b5 --- /dev/null +++ b/pyrobolearn/models/dmp/dmpytorch/discrete_dmp.py @@ -0,0 +1,149 @@ +#!/usr/bin/env python +"""Define the discrete dynamic movement primitive. +""" + +import numpy as np +import torch + +from pyrobolearn.models.dmp.dmpytorch.canonical_systems import DiscreteCS +from pyrobolearn.models.dmp.dmpytorch.forcing_terms import DiscreteForcingTerm +from pyrobolearn.models.dmp.dmpytorch.dmp import DMP + +__author__ = "Brian Delhaisse" +__copyright__ = "Copyright 2018, PyRoboLearn" +__credits__ = ["Brian Delhaisse"] +__license__ = "MIT" +__version__ = "1.0.0" +__maintainer__ = "Brian Delhaisse" +__email__ = "briandelhaisse@gmail.com" +__status__ = "Development" + + +class DiscreteDMP(DMP): + r"""Discrete Dynamic Movement Primitive + + Discrete DMPs have the same mathematical formulation as general DMPs, which is given by: + + .. math:: \tau^2 \ddot{y} = K (g - y) - D \tau \dot{y} + f(s) (g - y0) + + where :math:`\tau` is a scaling factor that allows to slow down or speed up the reproduced movement, :math:`K` + is the stiffness coefficient, :math:`D` is the damping coefficient, :math:`y, \dot{y}, \ddot{y}` are the position, + velocity, and acceleration of a DoF, and :math:`f(s)` is the non-linear forcing term. + + However, the forcing term in the case of discrete DMPs is given by: + + .. math:: f(s) = \frac{\sum_i \psi_i(s) w_i}{\sum_i \psi_i(s)} s + + where :math:`w` are the learnable weight parameters, and :math:`\psi` are the basis functions evaluated at the + given input phase variable :math:`s`, :math:`g` is the goal, and :math:`y_0` is the initial position. Note that + as the phase converges to 0, the forcing term also converges to that value. + + The basis functions (in the discrete case) are given by: + + .. math:: \psi_i(s) = \exp \left( - \frac{1}{2 \sigma_i^2} (x - c_i)^2 \right) + + where :math:`c_i` is the center of the basis function :math:`i`, and :math:`\sigma_i` is its width. + + Also, the canonical system associated with this transformation system is given by: + + .. math:: \tau \dot{s} = - \alpha_s s + + where :math:`\tau` is a scaling factor that allows to slow down or speed up the movement, :math:`s` is the phase + variable that drives the DMP, and :math:`\alpha_s` is a predefined constant. + + All these differential equations are solved using Euler's method. + + References: + [1] "Dynamical movement primitives: Learning attractor models for motor behaviors", Ijspeert et al., 2013 + """ + + def __init__(self, num_dmps, num_basis, dt=0.01, y0=0, goal=1, + forces=None, stiffness=None, damping=None): + """Initialize the discrete DMP + + Args: + num_dmps (int): number of DMPs + num_basis (int, int[M]): number of basis functions, or list of number of basis functions. + dt (float): step integration for Euler's method + y0 (float, float[M]): initial position(s) + goal (float, float[M]): goal(s) + forces (list, ForcingTerm): the forcing terms (which can have different basis functions) + stiffness (float): stiffness coefficient + damping (float): damping coefficient + """ + + # create discrete canonical system + cs = DiscreteCS(dt=dt) + + # create forcing terms (each one contains the basis functions and trainable weights) + if forces is None: + if isinstance(num_basis, int): + forces = [DiscreteForcingTerm(cs, num_basis) for _ in range(num_dmps)] + else: + if not isinstance(num_basis, (list, tuple)): + raise TypeError("Expecting 'num_basis' to be an int, list, or tuple of ints.") + if len(num_basis) != num_dmps: + raise ValueError("The length of th list of number of basis doesn't match the number of DMPs") + forces = [DiscreteForcingTerm(cs, n_basis) for n_basis in num_basis] + + # call super class constructor + super(DiscreteDMP, self).__init__(canonical_system=cs, forces=forces, y0=y0, goal=goal, + stiffness=stiffness, damping=damping) + + def get_scaling_term(self, new_goal=None): + """ + Return the scaling term for the forcing term. + + Args: + new_goal (float, float[M], None): the new goal position. If None, it will be the current goal. + + Returns: + float, float[M]: scaling term + """ + if new_goal is None: + new_goal = self.goal + return (new_goal - self.y0) / (self.goal - self.y0) + + def _generate_goal(self, y_des): + """Generate the goal for path imitation. + + Args: + y_des (np.array): the desired trajectory to follow with shape [num_dmps, timesteps] + + Returns: + float[M]: goal position + """ + return torch.clone(y_des[:, -1]) + + +# Tests +if __name__ == '__main__': + import matplotlib.pyplot as plt + + # tests canonical systems + discrete_cs = DiscreteCS() + + # tests basis functions + num_basis = 20 + discrete_f = DiscreteForcingTerm(discrete_cs, num_basis) + + # tests forcing terms + force = torch.sin(torch.linspace(0, 2*np.pi, 100)) + discrete_f.train(force, plot=True) + + # Test discrete DMP + discrete_dmp = DiscreteDMP(num_dmps=1, num_basis=num_basis) + t = torch.linspace(-6, 6, 100) + y_target = 1. / (1 + torch.exp(-t)) + discrete_dmp.imitate(y_target) + y, dy, ddy = discrete_dmp.rollout() + + plt.plot(y_target.numpy(), label='y_target') + plt.plot(y[0].detach().numpy(), label='y_pred') + # plt.plot(dy[0]) + # plt.plot(ddy[0]) + y, dy, ddy = discrete_dmp.rollout(new_goal=torch.tensor([2.])) + plt.plot(y[0].detach().numpy(), label='y_scaled') + plt.title('Discrete DMP') + plt.legend() + plt.show() diff --git a/pyrobolearn/models/dmp/dmpytorch/dmp.py b/pyrobolearn/models/dmp/dmpytorch/dmp.py new file mode 100644 index 0000000..285d878 --- /dev/null +++ b/pyrobolearn/models/dmp/dmpytorch/dmp.py @@ -0,0 +1,671 @@ +#!/usr/bin/env python +"""Define the general dynamic movement primitive abstract class. + +This file implements the DMP abstract class from which all dynamic movement primitive classes inherit from. +""" + +import numpy as np +import torch +import copy +import scipy.interpolate + +from pyrobolearn.models.dmp.dmpytorch.canonical_systems import CS +from pyrobolearn.models.dmp.dmpytorch.forcing_terms import ForcingTerm + + +__author__ = "Brian Delhaisse" +__copyright__ = "Copyright 2018, PyRoboLearn" +__credits__ = ["Brian Delhaisse", "Travis DeWolf"] +__license__ = "MIT" +__version__ = "1.0.0" +__maintainer__ = "Brian Delhaisse" +__email__ = "briandelhaisse@gmail.com" +__status__ = "Development" + + +class DMP(object): + r"""Dynamic Movement Primitive + + Dynamic movement primitives (DMPs) are a set of differential equations (for each degree of freedoms (DoFs), i.e. + general coordinates) that encodes a movement [1]. It is thought that movement primitives are the building blocks + of a movement, and several evidences show that such modules exist in animals [2]. + + DMPs are often formulated as a 2nd-order differential equation: + + .. math:: \tau^2 \ddot{y} = \alpha ( \beta (g - y) - \dot{y}) + f(s) + + or sometimes, as a first-order differential system: + + .. math:: + + \tau \dot{z} &= \alpha ( \beta (g - y) - z) + f(s) \\ + \tau \dot{y} &= z + + They can also be rewritten as: + + .. math:: \tau^2 \ddot{y} = K (g - y) - D \tau \dot{y} + f(s) + + where :math:`\tau` is a scaling factor that allows to slow down or speed up the reproduced movement, :math:`K` + is the stiffness coefficient, :math:`D` is the damping coefficient, :math:`y, \dot{y}, \ddot{y}` are the position, + velocity, and acceleration of a DoF, and :math:`f(s)` is the non-linear forcing term. These equations are also + known as the transformation systems and represent, with the canonical system, DMPs. + + All of the above formulations are equivalent to each other. However, in my humble opinion, the last equation + depicts better what the transformation system constitutes; it is a unit-mass spring-damper system or PD controller + with a forcing term. This last term is non-linear and can be learned from the demonstrations. + If the forcing is zero, then the differential equation is stable, and the position :math:`y` converges to the goal. + The stiffness and damping coefficients (:math:`K` and :math:`D`) are often selected such that the whole system + (without the forcing term) is critically damped (:math:`D = 2 \sqrt{K}`). Other behaviors can be obtained by + selecting the stiffness and damping coefficient such that we obtain: + * an undamped system: :math:`D = 0` or :math:`K \rightarrow \infty` + * an underdamped system: :math:`D < 2 \sqrt{K}` + * a critically damped system: :math:`D = 2 \sqrt{K}` + * an overdamped system: :math:`D > 2 \sqrt{K}` + + Because the last formulation is more intuitive (at least for me), it will be used in this class. + Imitation is performed by learning the forcing term. + + DMPs can be categorized in two main categories: + * discrete DMP: used to represent discrete movements such as such as reaching, pushing/pulling, etc. + * rhythmic DMP: used to represent rhythmic movements such as walking, running dribbling, sewing, etc. + + DMP have the following nice properties: + * translation invariant + * linear parameters but still allows to represent non-linear movements + + Here are few limitations/shortcomings: + * hard to couple sensory information with it + * have to come up with the number of basis functions + + For a more biologically-inspired DMP [5] which allows to adapt the goal in real-time and a better rescaling, see + the `BioDMP` class. + + Note that this code was inspired by the `pydmps` code [3, 4], but differ in several ways, notably: + - we undertake a more object-oriented programming (OOP) approach + - we use pytorch instead of numpy which allows to use automatic differentiation, and allows to use deep learning + tools with DMPs. TODO: need to check if the gradient vanishes... + - the equations are a little bit differents (e.g. :math:`tau`) in which we use the ones presented in the refs + - we decouple the Euler's method time step with the time step for the number of data points + - timesteps: we go from 0 to T included, while DeWolf goes from 0 to T-1 + - we use array operation instead of iterating over each element to update them + - we enforce consistency between the various methods and data structures + - we implement `BioDMP` which allows to adapt and rescale the goal in real-time based on [4] + - we implemented DMP sequencing based on [5] + - we implemented DMP that can be used with orientations based on [7] + - phase nodes which allows to couple phases, such as done in [8] for locomotion + - it can be used with RL algorithms, notably PoWER [9] and PI^2 [10] + + References: + [1] "Dynamical movement primitives: Learning attractor models for motor behaviors", Ijspeert et al., 2013 + [2] "Motor primitives in vertebrates and invertebrates", Flash et al., 2005 + [3] Tutorials on DMP: https://studywolf.wordpress.com/category/robotics/dynamic-movement-primitive/ + [4] PyDMPs (from DeWolf, 2013): https://github.com/studywolf/pydmps + [5] "Biologically-inspired Dynamical Systems for Movement Generation: Automatic Real-time Goal Adaptation + and Obstacle Avoidance", Hoffmann et al., 2009 + [6] "Action Sequencing using Dynamic Movement Primitives", Nemec et al., 2011 + [7] "Orientation in Cartesian Space Dynamic Movement Primitives", Ude et al., 2014 + [8] "A Framework for Learning Biped Locomotion with Dynamical Movement Primitives", Nakanishi et al., 2004 + [9] "Policy Search for Motor Primitives in Robotics", Kober et al., 2010 + [10] "A Generalized Path Integral Control Approach to Reinforcement Learning", Theodorou et al., 2010 + """ + + def __init__(self, canonical_system, forces, y0=0, goal=1, stiffness=None, damping=None): + """Initialize the DMP. + + Args: + canonical_system (CS): canonical system which drives the DMP transformation system + forces (list): list of forcing terms (one forcing term for each DMP). Each forcing term can have + different number of basis functions. + y0 (float, float[M]): initial state of DMPs + goal (float, float[M]): goal state of DMPs + stiffness (float): stiffness term in the transformation system for DMPs + damping (float): damping term in the transformation system for DMPs + """ + + self.cs = canonical_system + + if isinstance(forces, ForcingTerm): + forces = [forces] + elif isinstance(forces, (list, tuple)): + for force in forces: + if not isinstance(force, ForcingTerm): + raise TypeError("An item in the iterable is not an instance of `ForcingTerm`.") + else: + raise TypeError("Expecting forcing term to be an instance of ForcingTerm or a list/tuple of ForcingTerm") + + self.forces = forces + self.num_dmps = len(forces) + self.dt = self.cs.dt + self.timesteps = self.cs.timesteps + + # check initial and goal positions # TODO use property to set them + self.y0 = y0 + self.dy0 = torch.zeros(self.num_dmps) + self.ddy0 = torch.zeros(self.num_dmps) + self.goal = goal + self._check_offset() + + self.y, self.dy, self.ddy = self.y0, self.dy0, self.ddy0 + + # set stiffness and damping coefficient (if not specified, make them critically damped, i.e. D=2\sqrt{K}) + self.D = torch.ones(self.num_dmps) * 25. if damping is None else damping + self.K = self.D**2 / 4. if stiffness is None else stiffness + + # set up the DMP system + self.prev_s = self.cs.init_phase + self.reset() + + # target forcing term (keep a copy) + self.f_target = None + + def __repr__(self): + return self.__class__.__name__ + + def __call__(self, *args, **kwargs): + return self.step(*args, **kwargs) + + ############## + # Properties # + ############## + + @property + def goal(self): + """Return the goal position.""" + return self._goal + + @goal.setter + def goal(self, goal): + """Set the goal position.""" + if isinstance(goal, torch.Tensor): + goal = goal.clone().float() + elif isinstance(goal, (int, float)): + goal = torch.ones(self.num_dmps) * goal + elif isinstance(goal, (list, tuple)): + goal = torch.tensor(goal).float() + elif isinstance(goal, np.ndarray): + goal = torch.from_numpy(goal).float() + else: + raise TypeError("Expecting the goal to be a list/tuple/np.array/torch.Tensor of float/int, instead got: " + "{}".format(type(goal))) + self._goal = goal + + @property + def y0(self): + """Return the initial position.""" + return self._y0 + + @y0.setter + def y0(self, y0): + """Set the initial position.""" + if isinstance(y0, torch.Tensor): + y0 = y0.clone().float() + elif isinstance(y0, (int, float)): + y0 = torch.ones(self.num_dmps) * y0 + elif isinstance(y0, (list, tuple)): + y0 = torch.tensor(y0).float() + elif isinstance(y0, np.ndarray): + y0 = torch.from_numpy(y0).float() + else: + raise TypeError("Expecting the initial position to be a list/tuple/np.array/torch.Tensor of float/int, " + "instead got: {}".format(type(y0))) + self._y0 = y0 + + @property + def input_size(self): + """Return the input size of the model.""" + return 1 # 1 canonical system + + @property + def output_size(self): + """Return the output size of the model.""" + return len(self.forces) + + @property + def input_shape(self): + """Return the input shape of the model.""" + return tuple([self.input_size]) + + @property + def output_shape(self): + """Return the output shape of the model.""" + return tuple([self.output_size]) + + @property + def input_dim(self): + """Return the input dimension of the model; i.e. len(input_shape).""" + return len(self.input_shape) + + @property + def output_dim(self): + """Return the output dimension of the model; i.e. len(output_shape).""" + return len(self.output_shape) + + @property + def num_parameters(self): + """Return the total number of parameters""" + return sum([force.num_parameters for force in self.forces]) + + ################## + # Static Methods # + ################## + + @staticmethod + def copy(other, deep=False): + if not isinstance(other, DMP): + raise TypeError("Trying to copy an object which is not a DMP") + if deep: + return copy.deepcopy(other) + return copy.copy(other) + + @staticmethod + def is_parametric(): + """Return True as a DMP has weights that need to be optimized.""" + return True + + @staticmethod + def is_linear(): + """Return True as a DMP is linear in terms of its weights (i.e. learnable parameters)""" + return True + + @staticmethod + def is_recurrent(): + """Return False as a DMP is not a recurrent model.""" + return False + + @staticmethod + def is_sequential(): + """Return True as a DMP is a (temporal) sequential model.""" + return True + + @staticmethod + def is_probabilistic(): + """The DMP is a deterministic model.""" + return False + + @staticmethod + def is_discriminative(): + """The DMP is a discriminative model which predicts the output :math:`y` given the input :math:`x`""" + return True + + @staticmethod + def is_generative(): + """The DMP is not a generative model.""" + return False + + ########### + # Methods # + ########### + + def parameters(self): + """Returns an iterator over the model parameters.""" + for force in self.forces: + yield force.weights + + def named_parameters(self): + """Returns an iterator over the model parameters, yielding both the name and the parameter itself""" + for force in self.forces: + yield str(force), force.weights + + def list_parameters(self): + """Return a list of parameters""" + return list(self.parameters()) + + def hyperparameters(self): + """Return an iterator over the hyper-parameters.""" + yield self.K + yield self.D + # yield basis_functions + + def named_hyperparameters(self): + """Return an iterator over the hyper-parameters, yielding both the name and the hyper-parameter itself.""" + yield "stiffness", self.K + yield "damping", self.D + + def list_hyperparameters(self): + """Return a list of hyper-parameters.""" + return list(self.hyperparameters()) + + def get_vectorized_parameters(self, to_numpy=True): + """Return a vectorized form (1 dimensional array) of the parameters.""" + parameters = self.parameters() + vector = torch.cat([parameter.view(-1) for parameter in parameters]) + # if to_numpy: + # return vector.detach().numpy() + return vector + + def set_vectorized_parameters(self, vector): + """Set the vector parameters.""" + # convert the vector to torch array + # if isinstance(vector, np.ndarray): + # vector = torch.from_numpy(vector).float() + + # set the parameters from the vectorized one + # idx = 0 + # for parameter in self.parameters(): + # size = parameter.nelement() + # parameter.data = vector[idx:idx+size].view(parameter.shape) + # idx += size + + # set the parameters from the vectorized one + idx = 0 + for force in self.forces: + size = force.weights.size + force.weights = vector[idx:idx+size].view(force.weights.shape) + idx += size + + def get_damping_ratio(self): + r""" + Return the damping ratio :math:`\zeta = D / D_c` where :math:`D_c = 2 \sqrt{K}`. + + * if :math:`\zeta` = 0, the system is undamped (i.e. no damping) + * if :math:`\zeta` < 1, the system is underdamped (i.e. there will be some oscillations) + * if :math:`\zeta` = 1, the system is critically damped (i.e. return to equilibrium as fast as possible + without oscillating). + * if :math:`\zeta` > 1, the system is overdamped (i.e. the system returns to equilibrium without oscillating + but might be slow depending on the damping value). + """ + return self.D / (2. * torch.sqrt(self.K)) + + def _check_offset(self): + """Check to see if the initial position and goal are the same. If that is the case, offset slightly so that + the forcing term is not 0. Otherwise, look at the `BioDMP` class. + """ + self.goal[self.y0 == self.goal] += 1e-4 + + def get_scaling_term(self, new_goal=None): + # this is overridden by the child classes + return torch.ones(self.num_dmps) + + def _generate_goal(self, y_des): + raise NotImplementedError() + + def reset(self): + """Reset the transformation and canonical systems""" + self.y = self.y0.clone() + self.dy = self.dy0.clone() # torch.zeros(self.num_dmps) + self.ddy = self.ddy0.clone() + self.prev_s = self.cs.reset() + + def step(self, s=None, tau=1.0, error=0.0, forces=None, new_goal=None, external_force=None, + rescale_force=True): + """Run the DMP transformation system for a single time step. + + Args: + s (None, float): the phase value. If None, it will use the canonical system. + tau (float): Increase tau to make the system slower, and decrease it to make it faster + error (float): optional system feedback + forces (float[M]): if given, it will replace the forcing term (where `M` = number of DMPs) + new_goal (float[M]): new goal (where `M` = number of DMPs) + rescale_force (bool): if the given forcing term should be rescaled. + """ + + # system feedback + error_coupling = 1.0 / (1.0 + error) + + # get phase from canonical system + if s is None: + s = self.cs.step(tau=tau, error_coupling=error_coupling) + elif not isinstance(s, (float, int)): + raise TypeError("Expecting the phase 's' to be a float or integer. Instead, I got {}".format(type(s))) + + # check if same phase as before + if s == self.prev_s: + return self.y, self.dy, self.ddy + + if new_goal is None: + new_goal = self.goal + + # save previous position and velocity + prev_y, prev_dy = self.y.clone(), self.dy.clone() + + # compute scaling factor for the forcing term + scaling = self.get_scaling_term(new_goal) + + # for each DMP, solve transformation system equation using Euler's method + for d in range(self.num_dmps): + + # compute forcing term + if forces is None: + f = self.forces[d](s) * scaling[d] + # f = self.f_gen(s) * scaling[d] + else: + if rescale_force: + f = forces[d] * scaling[d] + else: + f = forces[d] + + # DMP acceleration + self.ddy[d] = self.K[d]/(tau**2) * (new_goal[d] - self.y[d]) - self.D[d]/tau * self.dy[d] + f/(tau**2) + if external_force is not None: + self.ddy[d] += external_force[d] + self.dy[d] += self.ddy[d] / tau * self.dt * error_coupling + self.y[d] += self.dy[d] * self.dt * error_coupling + + # return self.y, self.dy, self.ddy + return prev_y, prev_dy, self.ddy + + def rollout(self, timesteps=None, tau=1.0, error=0.0, forces=None, new_goal=None, rescale_force=True, + **kwargs): + """Generate position, velocity, and acceleration trajectories, no feedback is incorporated. + + Args: + tau (float): Increase tau to make the system slower, and decrease it to make it faster + timesteps (None, int): the number of steps to perform + error (float): optional system feedback + forces (torch.Tensor): if given, it will replace the forcing term (shape [num_dmps, timesteps]) + new_goal (torch.Tensor): new goal (of shape [num_dmps,]) + + Returns: + float[M,T]: y (position) trajectories + float[M,T]: dy (velocity) trajectories + float[M,T]: ddy (acceleration) trajectories + """ + # reset the canonical and transformation systems + self.reset() + + if timesteps is None: + timesteps = int(self.timesteps * tau) + + # set up tracking vectors + y_track = torch.zeros(self.num_dmps, timesteps) + dy_track = torch.zeros(self.num_dmps, timesteps) + ddy_track = torch.zeros(self.num_dmps, timesteps) + + # for the other timesteps, solve DMP equation using Euler's method + for t in range(timesteps): + if forces is None: + y, dy, ddy = self.step(tau=tau, error=error, new_goal=new_goal, external_force=None) + else: + y, dy, ddy = self.step(tau=tau, error=error, forces=forces[:, t], new_goal=new_goal, + rescale_force=rescale_force) + + # record timestep + y_track[:, t] = y + dy_track[:, t] = dy + ddy_track[:, t] = ddy + + return y_track, dy_track, ddy_track + + def train(self, f_target): + """Train the forcing terms.""" + # train each forcing term + if f_target.shape[0] != len(self.forces): + raise ValueError("Mismatch between the number of forcing terms") + + # train each forcing term + for forces, target in zip(self.forces, f_target): + forces.train(target) + + def imitate(self, y_des, dy_des=None, ddy_des=None, interpolation='cubic', plot=False): + """Imitate a desired trajectory, and learn the parameters that best realizes it. + + Args: + y_des (torch.Tensor): the desired position trajectories of each DMP with shape [num_dmps, timesteps] + dy_des (torch.Tensor): the desired velocities with shape [num_dmps, timesteps] + ddy_des (torch.Tensor): the desired accelerations with shape [num_dmps, timesteps] + interpolation (str): how to interpolate the data. Select between 'linear', 'cubic', and 'hermite'. + """ + if isinstance(y_des, np.ndarray): + y_des = torch.from_numpy(y_des).float() + if isinstance(dy_des, np.ndarray): + dy_des = torch.from_numpy(dy_des).float() + if isinstance(ddy_des, np.ndarray): + ddy_des = torch.from_numpy(ddy_des).float() + + # set initial state and goal + if y_des.dim() == 1: + y_des = y_des.view(1, len(y_des)) + self._y0 = torch.clone(y_des[:, 0]) + self._goal = self._generate_goal(y_des) + self._check_offset() + + timesteps = y_des.shape[1] + + def interpolate(x, dt, period, timesteps, new_timesteps, interpolation=interpolation, return_gen=False): + # generate function to interpolate the desired trajectory + x = x.numpy() + t = np.linspace(0, period, timesteps) + if interpolation == 'linear': # use linear interpolation + path_gen = scipy.interpolate.interp1d(t, x, axis=-1) + elif interpolation == 'cubic': # use cubic spline interpolation + path_gen = scipy.interpolate.CubicSpline(t, x, axis=-1) + else: # TODO: implement hermite (see utils.interpolator.hermite) + raise ValueError("The requested interpolation has not been implemented. Select between 'linear' or " + "'cubic'") + if return_gen: + return path_gen + return torch.from_numpy(path_gen([t * self.dt for t in range(new_timesteps)])).float() + + y_des = interpolate(y_des, dt=self.dt, period=self.cs.T, timesteps=timesteps, + new_timesteps=self.timesteps, interpolation=interpolation) + + def diff(tensor): + if tensor.dim() == 1: + return y_des[1:] - y_des[:-1] + else: + return y_des[..., 1:] - y_des[..., :-1] + + # compute desired velocity if necessary + if dy_des is None: + # calculate velocity of y_des + dy_des = diff(y_des) / self.dt + # add zero to the beginning of every row + dy_des = torch.cat((torch.zeros(self.num_dmps, 1), dy_des), dim=1) + else: + if dy_des.dim() == 1: + dy_des = dy_des.view(1, len(dy_des)) + dy_des = interpolate(dy_des, self.dt, self.cs.T, dy_des.shape[1], self.timesteps, + interpolation=interpolation) + self.dy0 = torch.clone(dy_des[:, 0]) + + # compute desired acceleration if necessary + if ddy_des is None: + # calculate acceleration of y_des + ddy_des = diff(dy_des) / self.dt + # add zero to the beginning of every row + ddy_des = torch.cat((torch.zeros(self.num_dmps, 1), ddy_des), dim=1) + else: + if ddy_des.dim() == 1: + ddy_des = ddy_des.view(1, len(ddy_des)) + ddy_des = interpolate(ddy_des, self.dt, self.cs.T, ddy_des.shape[1], self.timesteps, + interpolation=interpolation) + self.ddy0 = torch.clone(ddy_des[:, 0]) + + # find the force required to move along this trajectory (with shape [num_dmps, timesteps]) + f_target = ddy_des - self.K.view(-1, 1) * (self.goal.view(-1, 1) - y_des) + self.D.view(-1, 1) * dy_des + + # plot + if plot: + import matplotlib.pyplot as plt + plt.figure() + plt.plot(y_des[0].numpy(), 'b', label='pos') + plt.plot(dy_des[0].numpy(), 'g', label='vel') + plt.plot(ddy_des[0].numpy(), 'r', label='acc') + plt.plot(f_target[0].numpy(), 'k', label='force') + plt.legend() + plt.show() + + # self.f_gen = interpolate(f_target, dt=self.dt, period=self.cs.T, timesteps=timesteps, + # new_timesteps=timesteps, interpolation=interpolation, return_gen=True) + + # efficiently generate weights to realize f_target + self.f_target = f_target + self.train(f_target) + + # reset the canonical and transformation systems + self.reset() + + return y_des + + def get_forces(self, s): + """ + Get the forcing terms based on the given phase value. + + Args: + s (float, float[T]): phase value(s) + + Returns: + float[M], float[M,T]: forcing terms + """ + return torch.tensor([self.forces[d](s) for d in range(self.num_dmps)]) + + def generate_goal(self, y0=None, dy0=None, ddy0=None, f0=None): + """ + Generate the goal from the initial positions, velocities, accelerations, and forces. + + Args: + y0 (float[M], None): initial positions. If None, it will take the default initial positions. + dy0 (float[M], None): initial velocities. If None, it will take the default initial velocities. + ddy0 (float[M], None): initial accelerations. If None, it will take the default initial accerelations. + f0 (float[M], None): initial forcing terms. If None, it will compute it based on the learned weights. + You can also give `dmp.f_target[:,0]` to get the correct goal. + + Returns: + float[M]: goal position for each DMP. + """ + if y0 is None: + y0 = self.y0 + if dy0 is None: + dy0 = self.dy0 + if ddy0 is None: + ddy0 = self.ddy0 + if f0 is None: + s0 = self.cs.init_phase + f0 = self.get_forces(s0) + + return 1/self.K * (ddy0 + self.D * dy0 + self.K * y0 - f0) + + def sequence(self, model, mode=0): + """ + Define how to sequence with another DMP model. + + Args: + model (DMP): DMP model + mode (int): specifies how to sequence the two DMP models. + + Returns: + DMP: the sequenced model + + References: + [1] "Action Sequencing using Dynamic Movement Primitives", Nemec et al., 2011 + """ + if not isinstance(model, DMP): + raise TypeError("The given model is not an instance of DMP.") + pass + + # def __rshift__(self, other): + # """ + # Sequence DMP model with another learning model. + # + # Ref: "Action Sequencing using Dynamic Movement Primitives", Nemec et al., 2011 + # + # :param other: another DMP model + # :return: + # """ + # # If we sequence two DMP models + # if isinstance(other, DMP): + # + # else: + # # if it is another model, call the parent's method which knows how to sequence different models + # super(DMP, self).__rshift__(other) + diff --git a/pyrobolearn/models/dmp/dmpytorch/forcing_terms.py b/pyrobolearn/models/dmp/dmpytorch/forcing_terms.py new file mode 100644 index 0000000..39ac422 --- /dev/null +++ b/pyrobolearn/models/dmp/dmpytorch/forcing_terms.py @@ -0,0 +1,356 @@ +#!/usr/bin/env python +"""Define the forcing terms used in dynamic movement primitives + +This file implements the forcing terms used for discrete and rhythmic dynamic movement primitives. +""" + +import matplotlib.pyplot as plt + +from pyrobolearn.models.dmp.dmpytorch.canonical_systems import * +from pyrobolearn.models.dmp.dmpytorch.basis_functions import * +from pyrobolearn.models.dmp.dmpytorch.weight import WeightModule + + +__author__ = "Brian Delhaisse" +__copyright__ = "Copyright 2018, PyRoboLearn" +__credits__ = ["Brian Delhaisse"] +__license__ = "MIT" +__version__ = "1.0.0" +__maintainer__ = "Brian Delhaisse" +__email__ = "briandelhaisse@gmail.com" +__status__ = "Development" + + +class ForcingTerm(torch.nn.Module): + r"""Forcing term used in DMPs + + This basically computes the unscaled forcing term, i.e. a weighted sum of basis functions, which is given by: + + .. math:: f(s) = \frac{ sum_{i} \psi_i(s) w_i }{ \sum_i \psi_i(s) } + + where :math:`w` are the learnable weight parameters, and :math:`\psi` are the basis functions evaluated at the + given input phase variable :math:`s`. + """ + + def __init__(self, weights, basis_functions): + """ + Initialize the forcing terms. + + Args: + weights (torch.nn.Module): weight module. + basis_functions (BF): basis functions. + """ + # check that the arguments have the same length + super(ForcingTerm, self).__init__() + if not isinstance(weights, torch.nn.Module): + raise TypeError("Expecting the weights to be an instance of `torch.nn.Module`, instead got: " + "{}".format(type(weights))) + self._w = weights + if not isinstance(basis_functions, torch.nn.Module): + raise TypeError("Expecting the basis_functions to be an instance of `torch.nn.Module`, instead got: " + "{}".format(type(basis_functions))) + self.psi = basis_functions + + @property + def weights(self): + """Return the inner weights which might be a torch.Tensor or torch.nn.Module.""" + if isinstance(self._w, WeightModule): + return self._w.weight + return self._w + + @property + def num_parameters(self): + """Return the number of parameters.""" + weights = self.weights + if isinstance(weights, torch.Tensor): + return weights.numel() + return sum(p.numel() for p in weights.parameters()) + + @staticmethod + def is_linear(): + return True + + @staticmethod + def is_parametric(): + return True + + @staticmethod + def is_recurrent(): + return False + + def forward(self, s): + """Compute the forcing term + + Compute the value of the forcing term :math:`f(s)` at the given phase value :math:`s`. + + Args: + s (float): phase value + + Returns: + float: value of the forcing term at the given phase value + """ + psi_track = self.psi(s) + if psi_track.dim() == 1: + return torch.dot(psi_track, self._w(s)) / torch.sum(psi_track) + return (torch.mm(psi_track, self._w(s).unsqueeze(1))).squeeze(1) / torch.sum(psi_track, dim=1) + + def weighted_basis(self, s): + """Generate weighted basis + + Returns: + np.array[T, M]: weighted basis + """ + return self.psi(s) * self._w(s) + + def normalized_weighted_basis(self, s): + """Generate normalized weighted basis + + Args: + s (float): phase value + + Returns: + np.array[T,M]: normalized weighted basis + """ + psi_track = self.psi(s) + return ((psi_track * self._w(s)).t() / torch.sum(psi_track, dim=1)).t() + + def __str__(self): + return self.__class__.__name__ + + # To override in child classes + def train(self, f_target): + raise NotImplementedError + + # alias + generate_weights = train + + +class DiscreteForcingTerm(ForcingTerm): + r"""Discrete Forcing Term + + .. math:: f(s) = \frac{ sum_{i} \psi_i(s) w_i }{ \sum_i \psi_i(s) } s + + where :math:`w` are the learnable weight parameters, and :math:`\psi` are the basis functions evaluated at the + given input phase variable :math:`s`. + + This forcing term has the property that as the phase converges to 0, it also converges to 0, allowing the + linear part of the DMP equation to converge to the goal. + """ + + def __init__(self, cs, num_basis): + """Initialize the discrete forcing term. + + Args: + cs (CS): discrete canonical system + num_basis (int): number of basis functions + """ + # set canonical system + if not isinstance(cs, DiscreteCS): + raise TypeError("Expecting 'cs' to be an instance of DiscreteCS") + self.cs = cs + + # set num_basis + self.num_basis = num_basis + + # create weights + weights = WeightModule(weight=torch.zeros(num_basis)) # default f=0 + + # desired activations throughout time + c = torch.linspace(0, cs.T, num_basis) + c = torch.exp(-cs.alpha_s * c) + + # set variance of basis functions (this was found by trial and error by DeWolf) + h = torch.ones(num_basis) * num_basis**1.5 / c / cs.alpha_s + + basis = EBF(center=c, h=h) + super(DiscreteForcingTerm, self).__init__(weights, basis) + + def forward(self, s): + # call parent compute + force = super(DiscreteForcingTerm, self).forward(s) + # scale with phase s + return force * s + + def train(self, f_target, plot=False): + """Train the weights to match the given target forcing term + + Generate a set of weights over the basis functions such that the target forcing term trajectory is matched. + + Args: + f_target (np.array): the desired forcing term trajectory + """ + + # calculate phase and basis functions + s_track = self.cs.rollout() + psi_track = self.psi(s_track) # shape=TxM + + # efficiently calculate BF weights using LWR (Locally Weighted (Linear) Regression) + # spatial scaling term + if isinstance(self.weights, torch.Tensor): + for b in range(self.num_basis): + numerator = torch.sum(s_track * psi_track[:, b] * f_target) + denominator = torch.sum(s_track**2 * psi_track[:, b]) + self.weights[b] = numerator / denominator + else: + raise NotImplementedError + + # set nan to 0 + if isinstance(self.weights, torch.Tensor): + self.weights[self.weights != self.weights] = 0. + # self.weights = np.nan_to_num(self.weights) + + if plot: + # plot the basis function activations + plt.figure() + plt.subplot(211) + plt.plot(psi_track.numpy()) + plt.title('basis functions') + + # plot the desired forcing function vs approx for the first dmp + plt.subplot(212) + plt.title('discrete force') + plt.plot(f_target.numpy(), label='f_target', linewidth=2.5) + plt.plot(self.forward(s_track).detach().numpy(), label='f_pred', linewidth=2.5) + + # weighted sum of basis functions + wps = self.weighted_basis(s_track) + plt.plot(wps.detach().numpy(), linewidth=0.5) + + plt.legend() + plt.tight_layout() + plt.show() + + +class RhythmicForcingTerm(ForcingTerm): + r"""Rhythmic Forcing Term + + .. math:: f(s) = \frac{ sum_{i} \psi_i(s) w_i }{ \sum_i \psi_i(s) } a + + where :math:`w` are the learnable weight parameters, :math:`\psi` are the basis functions evaluated at the + given input phase variable :math:`s`, and :math:`a` is the amplitude. + + When used with DMPs, it produces a limit cycle behavior. + """ + + def __init__(self, cs, num_basis, amplitude=1.): + """ + Initialize the rhythmic forcing term. + + Args: + cs (CS): rhythmic canonical system + num_basis (int): number of basis functions + amplitude (float): amplitude + """ + # set canonical system + if not isinstance(cs, RhythmicCS): + raise TypeError("Expecting 'cs' to be an instance of RhythmicCS") + self.cs = cs + + # set num_basis and amplitude + self.num_basis = num_basis + self.a = amplitude + + # create weights + weights = WeightModule(weight=torch.zeros(num_basis)) # default f=0 + + # set the centre of the Gaussian basis functions to be spaced evenly + c = torch.linspace(0, cs.T, num_basis + 1) # the '+1' is because it is rhythmic, c(0) = c(2pi) + c = c[:-1] + + # set concentration of basis function (this was found by trial and error by DeWolf) + h = torch.ones(num_basis) * num_basis + + # create basis functions + basis = CBF(center=c, h=h) + super(RhythmicForcingTerm, self).__init__(weights, basis) + + def forward(self, s): + # call parent compute + force = super(RhythmicForcingTerm, self).forward(s) + # scale with amplitude and return it + return force * self.a + + def train(self, f_target, plot=False): + """Train the weights to match the given target forcing term + + Generate a set of weights over the basis functions such that the target forcing term trajectory is matched. + + Args: + f_target (np.array): the desired forcing term trajectory + plot (bool): If True, it will plot. + """ + + # calculate phase and basis functions + s_track = self.cs.rollout() + psi_track = self.psi(s_track) # shape=TxM + + # efficiently calculate BF weights using LWR (Locally Weighted (Linear) Regression) + if isinstance(self.weights, torch.Tensor): + for b in range(self.num_basis): + self.weights[b] = (torch.dot(psi_track[:, b], f_target) / (torch.sum(psi_track[:, b]))) # + 1e-10)) + else: + raise NotImplementedError + + if plot: + # plot the basis function activations + plt.figure() + plt.subplot(211) + plt.plot(psi_track.numpy()) + plt.title('basis functions') + + # plot the desired forcing function vs approx for the first dmp + plt.subplot(212) + plt.title('rhythmic force') + plt.plot(f_target.numpy(), label='f_target', linewidth=2.5) + plt.plot(self.forward(s_track).detach().numpy(), label='f_pred', linewidth=2.5) + wps = self.weighted_basis(s_track) + plt.plot(wps.detach().numpy(), linewidth=0.5) + plt.legend() + plt.tight_layout() + plt.show() + + +# Tests +if __name__ == '__main__': + # tests canonical systems + discrete_cs = DiscreteCS() + rhythmic_cs = RhythmicCS() + + # plot canonical systems + plt.subplot(1, 2, 1) + plt.title('Discrete CS') + for tau in [1., 0.5, 2.]: + rollout = discrete_cs.rollout(tau=tau).numpy() + plt.plot(np.linspace(0, 1., len(rollout)), rollout, label='tau='+str(tau)) + plt.legend() + + plt.subplot(1, 2, 2) + plt.title('Rhythmic CS') + for tau in [1., 0.5, 2.]: + rollout = rhythmic_cs.rollout(tau=tau).numpy() + plt.plot(np.linspace(0, 1., len(rollout)), rollout, label='tau='+str(tau)) + plt.legend() + plt.show() + + # tests basis functions + num_basis = 20 + discrete_f = DiscreteForcingTerm(discrete_cs, num_basis) + rhythmic_f = RhythmicForcingTerm(rhythmic_cs, num_basis) + + plt.subplot(1, 2, 1) + rollout = discrete_cs.rollout() + plt.title('discrete basis fcts') + plt.plot(rollout.numpy(), discrete_f.psi(rollout).numpy()) + + plt.subplot(1, 2, 2) + rollout = rhythmic_cs.rollout() + plt.title('rhythmic basis fcts') + plt.plot(rollout.numpy(), rhythmic_f.psi(rollout).numpy()) + plt.show() + + # tests forcing terms + force = torch.sin(torch.linspace(0, 2*np.pi, 100)) + discrete_f.train(force, plot=True) + + force = torch.sin(torch.linspace(0, 2*np.pi, int(2*np.pi*100))) + rhythmic_f.train(force, plot=True) diff --git a/pyrobolearn/models/dmp/dmpytorch/rhythmic_dmp.py b/pyrobolearn/models/dmp/dmpytorch/rhythmic_dmp.py new file mode 100644 index 0000000..1bba9c7 --- /dev/null +++ b/pyrobolearn/models/dmp/dmpytorch/rhythmic_dmp.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python +"""Define the rhythmic dynamic movement primitive. +""" + +import numpy as np +import torch + +from pyrobolearn.models.dmp.dmpytorch.canonical_systems import RhythmicCS +from pyrobolearn.models.dmp.dmpytorch.forcing_terms import RhythmicForcingTerm +from pyrobolearn.models.dmp.dmpytorch.dmp import DMP + +__author__ = "Brian Delhaisse" +__copyright__ = "Copyright 2018, PyRoboLearn" +__credits__ = ["Brian Delhaisse"] +__license__ = "MIT" +__version__ = "1.0.0" +__maintainer__ = "Brian Delhaisse" +__email__ = "briandelhaisse@gmail.com" +__status__ = "Development" + + +class RhythmicDMP(DMP): + r"""Rhythmic Dynamic Movement Primitive + + Rhythmic DMPs have the same mathematical formulation as general DMPs, which is given by: + + .. math:: \tau^2 \ddot{y} = K (g - y) - D \tau \dot{y} + f(s) + + where :math:`\tau` is a scaling factor that allows to slow down or speed up the reproduced movement, :math:`K` + is the stiffness coefficient, :math:`D` is the damping coefficient, :math:`y, \dot{y}, \ddot{y}` are the position, + velocity, and acceleration of a DoF, and :math:`f(s)` is the non-linear forcing term. + + However, the forcing term in the case of rhythmic DMPs is given by: + + .. math:: f(s) = \frac{\sum_i \psi_i(s) w_i}{\sum_i \psi_i(s)} a + + where :math:`w` are the learnable weight parameters, and :math:`\psi` are the basis functions evaluated at the + given input phase variable :math:`s`, and :math:`a` is the amplitude. + + The basis functions (in the rhythmic case) are given by: + + .. math:: \psi_i(s) = \exp \left( - h_i (\cos(s - c_i) - 1) \right) + + where :math:`c_i` is the center of the basis, and :math:`h_i` is a measure of concentration. + + Also, the canonical system associated with this transformation system is given by: + + .. math:: \tau \dot{s} = 1 + + where :math:`\tau` is a scaling factor that allows to slow down or speed up the movement, and :math:`s` is the + phase variable that drives the DMP. + + All these differential equations are solved using Euler's method. + + References: + [1] "Dynamical movement primitives: Learning attractor models for motor behaviors", Ijspeert et al., 2013 + """ + + def __init__(self, num_dmps, num_basis, dt=0.01, y0=0, goal=1, + forces=None, stiffness=None, damping=None): + """Initialize the rhythmic DMP + + Args: + num_dmps (int): number of DMPs + num_basis (int): number of basis functions + dt (float): step integration for Euler's method + y0 (float, np.array): initial position(s) + goal (float, np.array): goal(s) + forces (list, ForcingTerm): the forcing terms (which can have different basis functions) + stiffness (float): stiffness coefficient + damping (float): damping coefficient + """ + + # create rhythmic canonical system + cs = RhythmicCS(dt=dt) + + # create forcing terms (each one contains the basis functions and learnable weights) + if forces is None: + if isinstance(num_basis, int): + forces = [RhythmicForcingTerm(cs, num_basis) for _ in range(num_dmps)] + else: + if not isinstance(num_basis, (np.ndarray, list, tuple, set)): + raise TypeError("Expecting 'num_basis' to be an int, list, tuple, np.array or set.") + if len(num_basis) != num_dmps: + raise ValueError("The length of th list of number of basis doesn't match the number of DMPs") + forces = [RhythmicForcingTerm(cs, n_basis) for n_basis in num_basis] + + # call super class constructor + super(RhythmicDMP, self).__init__(canonical_system=cs, forces=forces, y0=y0, goal=goal, + stiffness=stiffness, damping=damping) + + def get_scaling_term(self, new_goal=None): + """ + Return the scaling term for the forcing term. For rhythmic DMPs it's non-diminishing, so this function just + returns 1. + """ + return torch.ones(self.num_dmps) + + def _generate_goal(self, y_des): + """Generate the goal for path imitation. + + For rhythmic DMPs, the goal is the average of the desired trajectory. + + Args: + y_des (float[M,T]): the desired trajectory to follow (with shape [num_dmps, timesteps]) + + Returns: + float[M]: goal positions (one for each DMP) + """ + goal = np.zeros(self.num_dmps) + for n in range(self.num_dmps): + num_idx = ~torch.isnan(y_des[n]) # ignore nan's when calculating goal + goal[n] = .5 * (y_des[n, num_idx].min() + y_des[n, num_idx].max()) + return goal diff --git a/pyrobolearn/models/dmp/dmpytorch/weight.py b/pyrobolearn/models/dmp/dmpytorch/weight.py new file mode 100644 index 0000000..460f1ca --- /dev/null +++ b/pyrobolearn/models/dmp/dmpytorch/weight.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python +"""Define weights used in the forcing terms in dynamic movement primitives +""" + +import torch + +__author__ = "Brian Delhaisse" +__copyright__ = "Copyright 2018, PyRoboLearn" +__credits__ = ["Brian Delhaisse"] +__license__ = "MIT" +__version__ = "1.0.0" +__maintainer__ = "Brian Delhaisse" +__email__ = "briandelhaisse@gmail.com" +__status__ = "Development" + + +class WeightTensor(torch.nn.Module): + r"""Fixed weight vector + + Weights used with the basis functions in the forcing terms. + """ + + def __init__(self, weight=None, num_basis=None): + """ + Initialize the fixed weight vector. + + Args: + weight (torch.Tensor, None): weight vector. If None, it will create a weight vector using the provided + :attr:`num_basis`. + num_basis (int): number of basis functions. This is used if no weight vector was given. + """ + super(WeightTensor, self).__init__() + if weight is None: + weight = torch.zeros(num_basis) + elif not isinstance(weight, torch.Tensor): + raise TypeError("Expecting the weight vector to be None or an instance of `torch.Tensor`, instead got: " + "{}".format(weight)) + self._weight = weight + self._weight.requires_grad = True + + @property + def weight(self): + return self._weight + + def forward(self, *x): + """Return the weight vector.""" + return self.weight + + +class WeightModule(torch.nn.Module): + r"""Weight vector + + Weights used with the basis functions in the forcing terms. + """ + + def __init__(self, weight=None, num_basis=None): + """ + Initialize the weight vector. + + Args: + weight (torch.nn.Module, torch.Tensor, None): weight module or vector. If None, it will create a weight + tensor/module using the provided :attr:`num_basis`. + num_basis (int): number of basis functions. This is used if no weight vector was given. + """ + super(WeightModule, self).__init__() + if weight is None: + weight = WeightTensor(num_basis=num_basis) + elif isinstance(weight, torch.Tensor): + weight = WeightTensor(weight=weight) + elif not isinstance(weight, torch.nn.Module): + raise TypeError("Expecting the weight vector to be None, an instance of `torch.Tensor`, or " + "`torch.nn.Module`, instead got: {}".format(type(weight))) + self._weight = weight + + @property + def weight(self): + """Return a torch.Tensor or torch.nn.Module.""" + if isinstance(self._weight, WeightTensor): + return self._weight.weight + return self._weight + + def forward(self, x): + """Forward the input to the weight module.""" + return self._weight(x)