Error in heat PDE implementation of Part3 L2C

[AliaNajwaMY]

Hello! I am trying to train the control (u) of my heat equation as below
$$y_t - ay_{xx} - by_x - cy =0, \quad (x,t) \in [0,1] \times [0,1]$$
$$y(0,t) = 0, \quad y(1,t)=u(t) $$
$$y(x,0)= \sin(\pi * x) $$
$$r(x,t) = 0.1 \sin(\pi x) \cos( \pi t) $$

using FDM, and then trying to solve each ODE for each x-direction steps. However, I keep encountering the error (RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn), which happens due to the fact that my output[loss] in the training process has properties requires_grad = False.

Setting requires_grad = True doesn't help since it makes my loss constant for each output.

I have also checked the computation graph to see any disconnected part, but it looks fine. What could be the issue here? My code is as below, please let me know if there's anything unclear:

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt


import neuromancer.psl as psl
from neuromancer.system import Node, System
from neuromancer.modules import blocks
from neuromancer.modules.activations import activations
from neuromancer.dataset import DictDataset
from neuromancer.constraint import variable
from neuromancer.loss import PenaltyLoss
from neuromancer.problem import Problem
from neuromancer.trainer import Trainer
from neuromancer.dynamics import integrators
from neuromancer.plot import pltCL, pltPhase
from neuromancer.constraint import variable
from abc import ABC, abstractmethod


# plotting imports
import matplotlib.pyplot as plt
#Set default dtype to float32
torch.set_default_dtype(torch.float)

#PyTorch random number generator
torch.manual_seed(1234)

# Random number generators in other libraries
np.random.seed(1234)

# Device configuration
if torch.backends.mps.is_available():
    device = torch.device('mps')
elif torch.cuda.is_available():
    device = torch.device('cuda')
else:
    device = torch.device('cpu')


class HeatEq:
    def __init__(self, a, b, c, Nx, Nt, L, T ):
        self.a = a
        self.b = b
        self.c = c
        self.Nx = Nx      # Number of spatial points
        self.Nt = Nt     # Number of time steps
        self.dx = L / self.Nx   # Spatial step size
        self.dt = T / self.Nt   # Time step size
        self.x = np.linspace(0, L, self.Nx + 1)  # Spatial points x
        self.t = np.linspace(0, T, self.Nt + 1)  # Spatial points time
        # self.y = np.zeros((self.Nx + 1, self.Nt + 1)) # Solution array
    
    def forward(self, t, x, u):
        u_t = torch.autograd.grad(u, t, create_graph=True)[0]
        u_x = torch.autograd.grad(u, x, create_graph=True)[0]
        u_xx = torch.autograd.grad(u_x, x, create_graph=True)[0]
        # return diff(u, t) - self.a * diff(u, x, order=2) - self.b* diff(u,x) - self.c * u
        return u_t - self.a * u_xx - self.b * u_x - self.c * u

   
class PDESystem(nn.Module, ABC):
    """
    Class for defining RHS of arbitrary ODE functions, 
    can be mix-and-matched according to expert knowledge.
    """

    def __init__(self, insize, outsize):
        super().__init__()
        self.in_features, self.out_features = insize, outsize
        self.nx = outsize
        self.nu = insize - outsize

    @abstractmethod
    def pde_equations(self, x, *args):
        pass

    def forward(self, x, *args):
        # assert len(x.shape) == Nx1
        return self.pde_equations(x, *args)


class HeatParam(PDESystem):

    def __init__(self,dx1, dt, Nt, Nx1, nu): # insize (y(x0), y(x1), ..., y(xNx1), u) (Nx1+1) and Nx1 output (y(x0), y(x1), ... y(xN))
        """

        :param insize:
        :param outsize:
        """
        super().__init__(insize=Nx1 + nu, outsize=Nx1)
        self.a = nn.Parameter(torch.tensor([0.1]), requires_grad=True)
        self.b = nn.Parameter(torch.tensor([0.1]), requires_grad=True)
        self.c = nn.Parameter(torch.tensor([0.1]), requires_grad=True)
        self.dx1 = dx1
        self.dt = dt
        self.Nx1 = Nx1
        self.Nt = Nt

    
    def g0(t):
        return 0
    
    def f1(x1): # initial condition u(x, 0)
        return np.sin(np.pi * x1)

 
    def pde_equations(self, x, u):
    
      
        dydt = torch.zeros_like(x) # check do we need Nx +1 and Nt+1 or if this is enough

        

    
        dydx = (x[:, 2:] - x[:, 1:-1]) / self.dx1
        dy2dxx = (x[:, 0:-2]+ x[:, 2:] - 2*x[:, 1:-1])/ (self.dx1 ** 2) 
        dydt[:, 1:-1] = self.a * dydx + self.b * dy2dxx + self.c * x[:, 1:-1]
            # dydt[n, -1] = (u[n+1]-u[n]) / self.dt # need to think about this

        return dydt
     # dydt is batchsize * Nx1
   


# Parameters
L = 1.0       # Length of the rod
T = 1.0       # Total time
a = 1
b = 1
c = 1
Nx1 = 50       # Number of spatial points
Nt = 1000     # Number of time steps
dx1 = L / Nx1  # Spatial step size
dt = T / Nt   # Time step size


# problem dimensions
nx = Nx1    # number of states 
nu = 1    # number of control inputs
nref = Nx1           # number of references


heateq = HeatEq(a, b, c, Nx1, Nt, L, T)


def reference_signal(x, t):
        """Returns the exact solution for a given x and t (for sinusoidal initial conditions).

        Parameters
        ----------
        x : np.ndarray
        t : np.ndarray
        """
        # return np.sin(np.pi * x + b * np.pi  * t) * np.exp((a * np.pi ** 2  + c) *t)
        return 0.1 * np.sin(np.pi * x ) * np.cos(np.pi *t)


nsteps = 5  # prediction horizon
n_samples = 1  # number of sampled scenarios

def ref_data(T, L, Nx1, Nt):
    # Create tensors:
    t = np.linspace(0, T, num=Nt).reshape(Nt, 1)
    print(t.shape)
    x = np.linspace(0, L, num=Nx1).reshape(Nx1, 1)
    print(x.shape)
    ref_sol = np.zeros((Nt, Nx1), dtype=np.float32).reshape(Nt, Nx1)

    # Obtain the value of the exact solution for each generated point:
    for i in range(Nx1): 
        for j in range(Nt):
            ref_sol[j][i] = reference_signal(x[i], t[j])
    
    np.savez("ref_data", x=x, t=t, ref_sol=ref_sol)

ref_data(T, L, Nx1, Nt)


ata = np.load('ref_data.npz')
ref_data = torch.from_numpy(data['ref_sol'].real)
# type(ref_data)
ref_data.shape

#  sampled references for training the policy
list_refs = ref_data[0:nsteps+1, :]
batched_ref = list_refs.reshape([n_samples, nsteps+1, nref])
batched_ref.shape

x1 = np.linspace(0, L, num=Nx1).reshape(Nx1, 1)
x1_init = torch.from_numpy(np.sin(x1 * np.pi)).reshape(n_samples, 1, nx)
x1_init.shape

# Training dataset
train_data = DictDataset({'x': torch.tensor(x1_init, dtype=torch.float32),   # sampled initial conditions of states
                          'r': batched_ref}, name='train')


# sampled references for development set
list_refs = ref_data[-(nsteps+1):, :] # last part of the horizon?
print(list_refs.shape)
batched_ref = list_refs.reshape([n_samples, nsteps+1, nref])
batched_ref.shape

# Development dataset
dev_data = DictDataset({'x': torch.rand(n_samples, 1, nx, dtype=torch.float32),    # sampled initial conditions of states
                        'r': batched_ref}, name='dev')

# torch dataloaders
batch_size = 1 # initially 200
train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size,
                                           collate_fn=train_data.collate_fn,
                                           shuffle=False)
dev_loader = torch.utils.data.DataLoader(dev_data, batch_size=batch_size,
                                         collate_fn=dev_data.collate_fn,
                                         shuffle=False)


# white-box PDE model with no-plant model mismatch

heat_pde = HeatParam(dx1=dx1, dt=dt, Nt=Nt, Nx1=Nx1, nu=nu)
heat_pde.a = nn.Parameter(torch.tensor(a), requires_grad=False)
heat_pde.b = nn.Parameter(torch.tensor(b), requires_grad=False)
heat_pde.c = nn.Parameter(torch.tensor(c), requires_grad=False)



# integrate continuous time PDE
integrator = integrators.Euler(heat_pde, h=torch.tensor(dt))   
# symbolic system model
model = Node(integrator, ['x', 'u'], ['x'], name='model')


# neural net control policy with hard control action bounds
net = blocks.MLP_bounds(insize=nx + nref, outsize=nu, hsizes=[32, 32],
                    nonlin=activations['gelu'] )
policy = Node(net, ['x', 'r'], ['u'], name='policy')

# closed-loop system model
cl_system = System([policy, model], nsteps=nsteps)
# cl_system.show()


# scaling factor for better convergence
scaling = 400

# variables

x = variable('x')
ref = variable('r')
# objectives
regulation_loss = (x == ref) ^ 2  # target posistion
regulation_loss.name = 'state_loss'

# list of constraints and objectives
objectives = [regulation_loss]
constraints = []


# data (x_k, r_k) -> parameters (xi_k) -> policy (u_k) -> dynamics (x_k+1)
components = [cl_system]
# create constrained optimization loss
loss = PenaltyLoss(objectives, constraints)
# construct constrained optimization problem
problem = Problem(components, loss)


problem.show()

optimizer = torch.optim.AdamW(problem.parameters(), lr=0.002)
#  Neuromancer trainer
trainer = Trainer(
    problem,
    train_loader, dev_loader,
    optimizer=optimizer,
    epochs=1000,
    train_metric='train_loss',
    eval_metric='dev_loss',
    warmup=50
)
# Train control policy
best_model = trainer.train()

[brunopjacob]

Hi @AliaNajwaMY! Thank you for using Neuromancer.

Regarding the runtime error, I believe the issue is in the DictDataset. Try adding the requires_grad=True there, i.e.,

train_data = DictDataset({'x': torch.tensor(x1_init, dtype=torch.float32, requires_grad=True),   # sampled initial conditions of states
                          'r': batched_ref}, name='train')

dev_data = DictDataset({'x': torch.rand(n_samples, 1, nx, dtype=torch.float32, requires_grad=True),    # sampled initial conditions of states
                        'r': batched_ref}, name='dev')

Regarding the loss not changing, I believe the cause is the lack of dependence of your control variable u in the pde_equations function. I.e., try adding this to the pde_equations function:

dydt[:, 0] = u.squeeze(-1) 

to enforce your boundary condition. This will cause u to have an effect on the predicted variables, and therefore the gradient will be non-zero, allowing for some backpropagation and consequently the loss should no longer be constant.

Please let me know if that helps!

Best wishes,
Bruno

[AliaNajwaMY]

Hello Bruno,Thanks for your reply. From what I’m understanding if I do that, I am enforcing the dy/dt = u at the boundary, which actually doesn’t satisfy the Dirichlet condition y= u at the boundary. When I looked at the documentation for neuromancer, it seems that the constraints are only defined as bounding conditions (ie max/min value of u or x or ref). How can I actually use the constraints to define my boundary conditions instead?Alia(She/ her) On 9 Jan 2025, at 4:19 PM, Bruno Jacob ***@***.***> wrote: Hi @AliaNajwaMY! Thank you for using Neuromancer. Regarding the runtime error, I believe the issue is in the DictDataset. Try adding the requires_grad=True there, i.e., train_data = DictDataset({'x': torch.tensor(x1_init, dtype=torch.float32, requires_grad=True), # sampled initial conditions of states 'r': batched_ref}, name='train') dev_data = DictDataset({'x': torch.rand(n_samples, 1, nx, dtype=torch.float32, requires_grad=True), # sampled initial conditions of states 'r': batched_ref}, name='dev') Regarding the loss not changing, I believe the cause is the lack of dependence of your control variable u in the pde_equations function. I.e., try adding this to the pde_equations function: dydt[:, 0] = u.squeeze(-1) to enforce your boundary condition. This will cause u to have an effect on the predicted variables, and therefore the gradient will be non-zero, allowing for some backpropagation and consequently the loss should no longer be constant. Please let me know if that helps! Best wishes, Bruno —Reply to this email directly, view it on GitHub, or unsubscribe.You are receiving this because you were mentioned.Message ID: ***@***.***>

[brunopjacob]

You are correct, good point. Perhaps you can try

x[:, 0] = 0.0
x[:, -1] = u.squeeze(dim=-1)

before the derivatives in the pde_equations method?