Infinite loss in one-dimensional HJB equation

[anthonybritto]

Thank you @lululxvi for your incredible work with this package. I'm trying to solve the following one-dimensional equation

over the domain $x \in [0, 1]$ subject to the boundary conditions $h(0) = 0$, $h'(1) = 0$.

My code and output is below (the constants module is also below). Apart from the ValueError, I cannot seem to get the "infinite" losses under control. I have tried modifying the NN as well as the solver options, loss function etc. but to no avail. Am I making an obvious mistake somewhere? Any insight is appreciated!

import deepxde as dde
# import numpy as np

import os
os.environ["TF_USE_LEGACY_KERAS"] = "1"

import constants as cons


def pde(x, h):
    
    # derivatives
    dx_h = dde.grad.jacobian(h, x)
    dxx_h = dde.grad.hessian(h, x)
    
    # equation terms
    t1 = (-1 + x) * (x + cons.z0 - x*cons.z0)**2 * cons.sigma_y**2 * (
        2*dx_h + (-1 + x)*dxx_h )
    
    t2 = 2 * cons.gamma**(cons.gamma/cons.epsilon) * cons.delta**(cons.delta/cons.epsilon) \
        * cons.epsilon * ((-1 + x)**2 * dx_h)**((-1+cons.epsilon)/cons.epsilon)
    
    t3 = 2 * (-1 + x) * (-1 + x + x*(-1 + cons.z0)*cons.Gamma1- cons.z0*cons.Gamma1) * dx_h
    
    t4 = ( (-1 + x) * (cons.mu_0 - cons.mu_s)**2 * dx_h**2 )\
    / ( cons.sigma_s**2 * ( 2*dx_h + (-1 + x)*dxx_h ))
    
    t5 = 2 * cons.Gamma0 * h
    
    # equation
    eq = t1 + t2 + t3 - t4 - t5
        
    return eq


def boundary_l(x, on_boundary):
    
    return on_boundary and dde.utils.isclose(x[0], 0)


def boundary_r(x, on_boundary):
    
    return on_boundary and dde.utils.isclose(x[0], 1)


geom = dde.geometry.Interval(0, 1)
bc_l = dde.icbc.DirichletBC(geom, lambda X: 0, boundary_l)
bc_r = dde.icbc.NeumannBC(geom, lambda X: 0, boundary_r)
data = dde.data.PDE(
    geometry=geom, pde=pde, bcs=[bc_l, bc_r],
    num_domain=16, num_boundary=2
)

layer_size = [1] + [50] * 3 + [1]
activation = "tanh"
initializer = "Glorot uniform"
net = dde.nn.FNN(layer_size, activation, initializer)

model = dde.Model(data, net)
model.compile("adam", lr=0.0000001, loss="mse")
losshistory, train_state = model.train(iterations=10000, display_every=10)

dde.saveplot(losshistory, train_state, issave=True, isplot=True)

Output:

Compiling model...
Building feed-forward neural network...
'build' took 0.103375 s

C:\Users\dt0955\AppData\Local\anaconda3\Lib\site-packages\deepxde\nn\tensorflow_compat_v1\fnn.py:116: UserWarning: `tf.layers.dense` is deprecated and will be removed in a future version. Please use `tf.keras.layers.Dense` instead.
  return tf.layers.dense(
'compile' took 2.730577 s

Training model...

Step      Train loss                        Test loss                         Test metric
0         [inf, 0.00e+00, 4.05e-01]         [inf, 0.00e+00, 4.05e-01]         []  
10        [nan, nan, nan]                   [nan, nan, nan]                   []  

Best model at step 0:
  train loss: inf
  test loss: inf
  test metric: 

'train' took 17.011475 s

Saving loss history to C:\Users\dt0955\OneDrive\Python\ee_consumption_investment\pde_nn\deepxde\loss.dat ...
Saving training data to C:\Users\dt0955\OneDrive\Python\ee_consumption_investment\pde_nn\deepxde\train.dat ...
Saving test data to C:\Users\dt0955\OneDrive\Python\ee_consumption_investment\pde_nn\deepxde\test.dat ...
Traceback (most recent call last):

  File ~\AppData\Local\anaconda3\Lib\site-packages\spyder_kernels\py3compat.py:356 in compat_exec
    exec(code, globals, locals)

  File c:\users\dt0955\onedrive\python\ee_consumption_investment\pde_nn\deepxde\terminal_gain_one_d.py:73
    dde.saveplot(losshistory, train_state, issave=True, isplot=True)

  File ~\AppData\Local\anaconda3\Lib\site-packages\deepxde\utils\external.py:175 in saveplot
    save_best_state(train_state, train_fname, test_fname)

  File ~\AppData\Local\anaconda3\Lib\site-packages\deepxde\utils\external.py:346 in save_best_state
    test = np.hstack((train_state.X_test, best_y))

  File <__array_function__ internals>:200 in hstack

  File ~\AppData\Roaming\Python\Python311\site-packages\numpy\core\shape_base.py:370 in hstack
    return _nx.concatenate(arrs, 1, dtype=dtype, casting=casting)

  File <__array_function__ internals>:200 in concatenate

ValueError: all the input arrays must have same number of dimensions, but the array at index 0 has 2 dimension(s) and the array at index 1 has 1 dimension(s)

For good measure, here is the constants module:

# discount rate
beta = 0.1

# utility parameters
gamma = 0.5
delta = 0.05

# efficiency parameter
eta = 0.6

# GBM drifts
mu_0 = 0.04
mu_p = 0.02
mu_q = 0.03
mu_s = 0.07
mu_y = 0.01

# GBM volatilities
sigma_p = 0.04
sigma_q = 0.05
sigma_s = 0.2
sigma_y = 0.02

# correlations
rho_pq = 0.3

# ode constants
z0 = -1 / (mu_0 - mu_y)
epsilon = 1 - gamma - delta
Gamma0 = 0.5 * (2*beta - gamma**2*(sigma_p**2 + sigma_y**2)\
    - delta * (-2*mu_q + 2*mu_y + (1+delta)*sigma_q**2 + (-1+delta)*sigma_y**2)\
    + gamma * (2*mu_p - 2*mu_y - sigma_p*(sigma_p + 2*delta*rho_pq*sigma_q)\
                + sigma_y**2 - 2*delta*sigma_y**2))
Gamma1 = mu_0 - mu_y + epsilon*sigma_y**2

[lululxvi]

There must exist some errors like divided by 0