# Hypersolvers for Optimal Control - Direct Optimal Control


```python
import sys; sys.path.append(2*'../') # go n dirs back
from src import *
```


```python
import time
import torch
import torch.nn as nn
import matplotlib.pyplot as plt

from torchdyn.core import NeuralODE
from torchdyn.datasets import *
from torchdyn.numerics import odeint, Euler, HyperEuler
```

## Optimal Control

We want to control an inverted pendulum and stabilize it in the upright position. The equations in Hamiltonian form describing an inverted pendulum with a torsional spring are as following:

$$\begin{equation}
    \begin{bmatrix} \dot{q}\\ \dot{p}\\ \end{bmatrix} = 
    \begin{bmatrix}
    0& 1/m \\
    -k& -\beta/m\\
    \end{bmatrix}
    \begin{bmatrix} q\\ p\\ \end{bmatrix} -
    \begin{bmatrix}
    0\\
    mgl \sin{q}\\
    \end{bmatrix}+
    \begin{bmatrix}
    0\\
    1\\
    \end{bmatrix} u
\end{equation}$$


```python
class ControlledPendulum(nn.Module):
    """
    Inverted pendulum with torsional spring
    """
    def __init__(self, u, m=1., k=.5, l=1., qr=0., β=.01, g=9.81):
        super().__init__()
        self.u = u # controller (nn.Module)
        self.nfe = 0 # number of function evaluations
        self.cur_f = None # current function evaluation
        self.cur_u = None # current controller evaluation 
        self.m, self.k, self.l, self.qr, self.β, self.g = m, k, l, qr, β, g # physics
        
    def forward(self, t, x):
        self.nfe += 1
        q, p = x[..., :1], x[..., 1:]
        self.cur_u = self.u(t, x)
        dq = p/self.m
        dp = -self.k*(q - self.qr) - self.m*self.g*self.l*torch.sin(q) \
            -self.β*p/self.m + self.cur_u
        self.cur_f = torch.cat([dq, dp], -1)
        return self.cur_f
```

In order to control the pendulum, we have to define a proper _integral cost function_ which will be our loss to be minimized during training. In a general form, it can be defined as:

\begin{equation}
        J = x^\top(t_f)\mathbf{P} x(t_f) + \int_{t_0}^{t_f} \left[ x^\top(t) \mathbf{Q} x(t) + u^\top(t) \mathbf{R} u(t) \right] dt
\end{equation}

where $ x = \begin{bmatrix} q\\ p\\ \end{bmatrix}$ is the state and $\mathbf{u}$ is the controller and matrices $\mathbf{P},~\mathbf{Q}, ~ \mathbf{R}$ are weights for controlling the performance.


```python
class IntegralCost(nn.Module):
    '''Integral cost function
    Args:
        x_star: torch.tensor, target position
        P: float, terminal cost weights
        Q: float, state weights
        R: float, controller regulator weights
    '''
    def __init__(self, x_star, P=0, Q=1, R=0):
        super().__init__()
        self.x_star = x_star
        self.P, self.Q, self.R, = P, Q, R
        
    def forward(self, x, u=torch.Tensor([0.])):
        """
        x: trajectory
        u: control input
        """
        cost = self.P*torch.norm(x[-1] - self.x_star, p=2, dim=-1).mean()
        cost += self.Q*torch.norm(x - self.x_star, p=2, dim=-1).mean()
        cost += self.R*torch.norm(u - 0, p=2).mean()
        return cost
```


```python
# Change device according to your configuration
# device = torch.device('cuda:1') if torch.cuda.is_available() else torch.device('cpu')
device = torch.device('cpu') # feel free to change :)
```


```python
# The controller is a simple MLP with one hidden layer with bounded output
class NeuralController(nn.Module):
    def __init__(self, model, u_min=-20, u_max=20):
        super().__init__()
        self.model = model
        self.u_min, self.u_max = u_min, u_max
        
    def forward(self, t, x):
        x = self.model(x)
        return torch.clamp(x, self.u_min, self.u_max)

model = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 1)).to(device)
u = NeuralController(model) 
for p in u.model[-1].parameters(): torch.nn.init.zeros_(p)

# Controlled system
sys = ControlledPendulum(u).to(device)
```


```python
from math import pi as π

# Loss function declaration
x_star = torch.Tensor([0., 0.]).to(device)
cost_func = IntegralCost(x_star)

# Time span
dt = 0.2
t0, tf = 0, 3 # initial and final time for controlling the system
steps = int((tf - t0)/dt) + 1 # so we have a time step of 0.2s
t_span = torch.linspace(t0, tf, steps).to(device)

# Initial distribution
x0 = π # limit of the state distribution (in rads and rads/second)
init_dist = torch.distributions.Uniform(torch.Tensor([-x0, -x0]), torch.Tensor([x0, x0]))
```


```python
# We consider the controller fixed during each solver step
class RandConstController(nn.Module):
    def __init__(self):
        super().__init__()
        self.u0 = torch.Tensor(1024, 1).uniform_(-10,10).to(device)
        
    def forward(self, t, x):
        return self.u0
    
# Save previously learned controller
u_no_hypersolver = sys.u
sys.u = RandConstController() # modify controller for training
```


```python
class VanillaHyperNet(nn.Module):
    """Simple hypernetwork for controlled systems
    Input: current x, f and u from the controlled system
    Output: p-th order residuals"""
    def __init__(self, net):
        super().__init__()
        self.net = net
        
    def forward(self, t, x):
        xfu = torch.cat([x, sys.cur_f, sys.cur_u], -1)
        return self.net(xfu)
    
net = nn.Sequential(nn.Linear(5, 32), nn.Softplus(), nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 2))
hypersolver = HyperEuler(VanillaHyperNet(net))
# model = nn.DataParallel(hypersolver, device_ids=[1]) # feel free to change here according to your setup and GPU available.
# model = model.to(device)
```


```python
model = torch.load('saved_models/hs_torchdyn.pt').to(device)
hypersolver = model.module
```

## Optimal Control with Hypersolvers
We can use the trained hypernetwork to generate trajectories of the pendulum: then, we cast the same optimal control problem defined at the beginning of the notebook and train the neural controller.


```python
sys = ControlledPendulum(None)
```


```python
# Reinstantiate controller
# net = nn.Sequential(nn.Linear(2, 64), nn.Softplus(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, 1)).to(device)
# u = NeuralController(net) 
# for p in u.model[-1].parameters(): torch.nn.init.zeros_(p)
u = BoxConstrainedController(2, 1, constrained=True, num_layers=3, output_scaling=torch.Tensor([-5, 5]).to(device))
sys.u = u

# Time span
t0, tf = 0, 3 # initial and final time for controlling the system
steps = int((tf - t0)/dt) + 1 # so we have a time step of 0.2s
t_span = torch.linspace(t0, tf, steps).to(device)

# Initial distribution
x0 = π # limit of the state distribution (in rads and rads/second)
init_dist = torch.distributions.Uniform(torch.Tensor([-x0, -x0]), torch.Tensor([x0, x0]))
```


```python

u_min, u_max = -5, 5 # constraints for the controller
cost_func = IntegralCost(x_star, P=1, Q=1) # final position is more important

# Hyperparameters
lr = 3e-3
epochs = 1000
bs = 1024
```

### Euler


```python
# Reinstantiate controller
# sys = ControlledPendulum(u)
u = BoxConstrainedController(2, 1, constrained=True, num_layers=3, output_scaling=torch.Tensor([-5, 5]).to(device))
sys.u = u.to(device)
opt = torch.optim.Adam(u.parameters(), lr=lr)

# Training loop
t0 = time.time(); losses=[]
for e in range(epochs):
    x0 = init_dist.sample((bs,)).to(device)
    _, trajectory = odeint(sys, x0, t_span, solver='euler', atol=1e-7, rtol=1e-7)    
    loss = cost_func(trajectory); losses.append(loss.detach().cpu().item())
    loss.backward(); opt.step(); opt.zero_grad()
    print('Loss {:.4f} , epoch {}'.format(loss.item(), e), end='\r')
timing = time.time() - t0; print('\nTraining time: {:.4f} s'.format(timing))

u_euler = sys.u
```

    /home/botu/anaconda3/lib/python3.8/site-packages/torchdyn/numerics/odeint.py:82: UserWarning: Setting tolerances has no effect on fixed-step methods
      warn("Setting tolerances has no effect on fixed-step methods")


    Loss 5.2740 , epoch 999
    Training time: 28.8481 s


### HyperEuler


```python
# Reinstantiate controller
u = BoxConstrainedController(2, 1, constrained=True, num_layers=3, output_scaling=torch.Tensor([-5, 5]).to(device))
sys.u = u.to(device)
opt = torch.optim.Adam(u.parameters(), lr=lr)

# Training loop
# Time should be taken with gpu, neural network inference is slower on cpu
t0 = time.time(); losses=[]
for e in range(epochs):
    x0 = init_dist.sample((bs,)).to(device)
    _, trajectory = odeint(sys, x0, t_span, solver=hypersolver)    
    loss = cost_func(trajectory); losses.append(loss.detach().cpu().item())
    loss.backward(); opt.step(); opt.zero_grad()
    print('Loss {:.4f} , epoch {}'.format(loss.item(), e), end='\r')
timing = time.time() - t0; print('\nTraining time: {:.4f} s'.format(timing))

u_hyper = sys.u
```

    Loss 1.1993 , epoch 999
    Training time: 40.9267 s


### Midpoint


```python
# Reinstantiate controller
u = BoxConstrainedController(2, 1, constrained=True, num_layers=3, output_scaling=torch.Tensor([-5, 5]).to(device))
sys.u = u.to(device)
opt = torch.optim.Adam(u.parameters(), lr=lr)

# Training loop
t0 = time.time(); losses=[]
for e in range(epochs):
    x0 = init_dist.sample((bs,)).to(device)
    _, trajectory = odeint(sys, x0, t_span, solver='midpoint')    
    loss = cost_func(trajectory); losses.append(loss.detach().cpu().item())
    loss.backward(); opt.step(); opt.zero_grad()
    print('Loss {:.4f} , epoch {}'.format(loss.item(), e), end='\r')
timing = time.time() - t0; print('\nTraining time: {:.4f} s'.format(timing))

u_mp = sys.u
```

    Loss 1.1765 , epoch 999
    Training time: 54.4334 s


### RK4


```python
# Reinstantiate controller
u = BoxConstrainedController(2, 1, constrained=True, num_layers=3, output_scaling=torch.Tensor([-5, 5]).to(device))
sys.u = u.to(device)
opt = torch.optim.Adam(u.parameters(), lr=lr)

# Training loop
t0 = time.time(); losses=[]
for e in range(epochs):
    x0 = init_dist.sample((bs,)).to(device)
    _, trajectory = odeint(sys, x0, t_span, solver='rk4')    
    loss = cost_func(trajectory); losses.append(loss.detach().cpu().item())
    loss.backward(); opt.step(); opt.zero_grad()
    print('Loss {:.4f} , epoch {}'.format(loss.item(), e), end='\r')
timing = time.time() - t0; print('\nTraining time: {:.4f} s'.format(timing))

u_rk4 = sys.u
```

    Loss 1.0977 , epoch 999
    Training time: 109.5863 s



```python
# Saving learned controllers
torch.save(u_euler, 'saved_models/u_euler.pt')
torch.save(u_hyper, 'saved_models/u_hyper.pt')
torch.save(u_mp, 'saved_models/u_mp.pt')
torch.save(u_rk4, 'saved_models/u_rk4.pt')
```

bb446cc3697436fdfcd5abf6843c73fd94a5d96e
ipynb
Jupyter Notebook
hypersolvers-control/experiments/pendulum/03b_direct_control.ipynb
Juju-botu/diffeqml-research
aa796c87447e5299ec4f25a07fc4d032afb1f63e
Apache-2.0
hypersolvers-control/experiments/pendulum/03b_direct_control.ipynb
Juju-botu/diffeqml-research
aa796c87447e5299ec4f25a07fc4d032afb1f63e
Apache-2.0
hypersolvers-control/experiments/pendulum/03b_direct_control.ipynb
Juju-botu/diffeqml-research
aa796c87447e5299ec4f25a07fc4d032afb1f63e
Apache-2.0
Qwen/Qwen-72B
1. YES
2. YES