title	README
author	Aidan Copinga

Lyapunov Trajectory Learning

This collection of examples serves to showcase a deep learning method for learning an energy function associated with a stable dynamical system, following 2309.08849 on the ArXiv.

Background

1. An autonomous dynamical system is characterized by $\dot{x} = f(x)$ where $f(x_0)=0$. Without loss of generality, $x_0=0$ with a translation $\tilde{x}= x-x_0$.
2. A stable system is one where $\lim_{t\to\infty} x(t) = x_0$.

Use Cases

Learning a dynamical system allows for more seamless trajectory learning and control for repeated tasks, such as robot writing or manufacturing, where the ending position is consistent.

1. Repeated Task -> Learn Energy -> Offline position/velocity from arbitrary starting point.
2. Couple lyapunov function $V$ with controller $u$ assuming control-affine systems $\dot{x} = f(x) + g(x)u$.

Example

Initializing

Getting started with training and using this small library is simple. Given some initial pytorch tensor data, $x$, with shape $B \times n_x$ and targets, $\dot{x}$, with shape $B \times n_x$, the forward of net_lyapunov produces an estimate for $\dot{x}$.

n_x = 2
n_y = 2 # dimension of the energy function, which should almost always be equal to n_x

n_hidden = 300 # optional
n_layers # hidden layers = n_layers - 2, optional
beta = 1e-4 # threshold for energy function time derivative learner
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

agent = net_lyapunov(n_x, n_y, n_hidden, n_layers, beta, device=device)

x_0 = torch.tensor([[1, 0]])

xdot_0 = agent(x_0)

Training Example

After the network is initialized, training is simple. The above paper uses a RMSE loss for velocity and a swept error area loss for the total difference in trajectory, as the lengths of the trajectories may not match, but for the sake of simplicity, the notebooks in this repository just use a MSE loss for velocity.

criterion = torch.nn.MSELoss()
optimizer = torch.optim.AdamW(agent.parameters(), lr=1e-5, weight_decay=0.99)

# in training loop
batch_x, batch_v = get_batch(data) # arbitrary data collection

optimizer.zero_grad(set_to_none=True)
agent.zero_grad(set_to_none=True)

one_step_velocity = agent(batch_x) # Computes the one step velocities associated with each x value in the batch.
# Alternative methods would be to simulate trajectories using an initial value and compare velocities or position data of simulation.

loss = criterion(one_step_velocity, batch_v)

Results

There were two examples presented in this series of experiments:

1. Using a known dynamical system ($\dot{x} = Ax + x \odot \sin(Ax)$, where $A$ is a known stable linear system, and $\sin$ of a vector is understood to mean $\sin(x) = [\sin(x_1), \sin(x_2)]^T$).
2. Using the PyLASA dataset to learn trajectories at different starting points based on demos.

Both systems used the following parameters:

1. Batch size of 64
2. AdamW Optimizer with learning rate 1e-5 and weight decay 0.99.

All results can be reproduced in the provided notebooks.

Known Dynamical System

Data Simulation

1. Redone every 128 epochs to decrease overfitting
2. Performed by forward Euler on the provided dynamical system with dt=0.01 over 1000 steps.
3. Simulated rough environment with small amounts of slippage on position sensors.

Training Loss Over Time

Trajectory Results

After training, the energy landscape, vector field, and random trajectories are plotted.

LASA Dataset

Data Collection

1. Training:Evalulation ratio is 6:1.
2. Training trajectories divided into 990 sub-trajectories of length 10. I.e. $[x_i, v_i, t_i], [x_{i+1}, v_{i+1}, t_{i+1}],\dots, [x_{i+9}, v_{i+9}, t_{i+9}]$.
3. Data simulation performed with forward Euler for 10 steps each batch.

Training Loss Over Time

Trajectory Results

After training, the energy landscape, vector field, and random trajectories are plotted.

Citation

Zhang et al., Learning a Stable Dynamic System with a Lyapunov Energy Function for Demonstratives Using Neural Networks, 2024, arXiv:2309.08849