<jupyter_start><jupyter_text><jupyter_code>import numpy as np
import time
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
#import ipdb

import itertools
class Pde:
    def __init__(self, n_dim = 2):
        self.n_dim = n_dim    
        self.lam = 0.
        print('>>> n_dim: '+str(n_dim))
    drift = lambda self,s,a: a
    
    run_cost = lambda self,s,a: (
            self.n_dim + sum([s1**2 for s1 in s])*2.0 
            + sum([a1**2 for a1 in a])/2.0
            )
    
    term_cost = lambda self,s: - sum([s1**2 for s1 in s])
    exact_soln = lambda self,s: - sum([s1**2 for s1 in s])


class Mdp(Pde):
    def __init__(self, n_dim = 2, n_mesh = 8):
        super().__init__(n_dim)
        self.n_mesh= n_mesh  
        self.h_mesh = 1./self.n_mesh #mesh size
        self.v_shape = tuple([self.n_mesh + 1]*self.n_dim)
        print('>>> n_mesh: '+str(n_mesh))


    #input: list of index
    #return: physicial coordinate
    def i2s(self,ix): 
        return [x * self.h_mesh for x in ix]
    
    def is_interior(self,ix):
        return all(map(lambda a: 0<a<self.n_mesh, ix))
        
    #input: lists of index and action
    #return: discount rate, running cost, list of next index, list of probability
    def step(self, ix, a, fd='cfd'):
        #ix = list(ix)
        s = self.i2s(ix)
        b = Pde.drift(Pde, s, a)
        if fd=='cfd':
            lam = self.n_dim/(self.n_dim+self.lam*(self.h_mesh**2))
            run_cost_h = self.h_mesh**2*self.run_cost(s,a)/self.n_dim
            
            ix_next = []; pr_next= []
            #cfd
            if self.is_interior(ix):
                for i in range(self.n_dim):
                    ix1 = ix.copy(); ix1[i]+=1; ix_next += [ix1,]
                    pr1 = (1+2.*self.h_mesh*b[i])/(self.n_dim*2.0) 
                    pr_next += [pr1,]
                for i in range(self.n_dim):
                    ix1 = ix.copy(); ix1[i]-=1; ix_next += [ix1,]
                    pr1 = (1-2.*self.h_mesh*b[i])/(self.n_dim*2.0) 
                    pr_next += [pr1,]
        elif fd=='ufd':
            c = self.n_dim+sum([abs(b1) for b1 in b])*self.h_mesh
            b_plus = [(abs(b1)+b1)/2. for b1 in b]
            b_minus = [(abs(b1)-b1)/2. for b1 in b]
            lam = c/(c+self.h_mesh**2*self.lam)
            run_cost_h = self.h_mesh**2*self.run_cost(s,a)/c
            ix_next = []; pr_next= []
            #ufd
            if self.is_interior(ix):
                for i in range(self.n_dim):
                    ix1 = ix.copy(); ix1[i]+=1; ix_next += [ix1,]
                    pr1 = (1+2.*self.h_mesh*b_plus[i])/(c*2.0) 
                    pr_next += [pr1,]
                for i in range(self.n_dim):
                    ix1 = ix.copy(); ix1[i]-=1; ix_next += [ix1,]
                    pr1 = (1-2.*self.h_mesh*b_minus[i])/(c*2.0) 
                    pr_next += [pr1,]        
        return lam, run_cost_h, ix_next, pr_next
    

 
####################
import itertools

def deep_iter(*shape):
    iters = (range(i) for i in shape)
    return itertools.product(*iters)


#product of a list
def product(l):
    out = 1
    for x in l:
        out *= x
    return out
    


class solver_nn(Mdp):
    def __init__(self, n_dim = 1, n_mesh = 8, fd = 'ufd'):
        super().__init__(n_dim, n_mesh)
        self.fd = fd
        self.a_space = tuple([3*n_mesh+1,]*n_dim)

    
    def i2a(self, ix):
        ix = list(ix)
        return [ix1*1./self.n_mesh for ix1 in ix]
    
    #input:
        #list of index and action, method
    #return:
        #q_val assuming v is value
    def q_val(self, ix, a, fd, v):
        lam, run_cost_h, ix_next, pr_next = self.step(ix,a,fd)
        out = run_cost_h
        for ix1, pr1 in zip(ix_next, pr_next):
            out+=pr1*v(torch.FloatTensor(self.i2s(ix1)))
        out *= lam
        return out
    
    #minimum over action space
    def min_a(self, fun):
        out_ind = [0,]*self.n_dim; out_val = fun(self.i2a(out_ind))
        for ix in deep_iter(*self.a_space):
            if fun(self.i2a(ix))<out_val:
                #out_ind = ix;
                out_val = fun(self.i2a(ix))
        return out_val #,out_ind

    def bellman(self,ix, v):
        #start_time = time.time()
        s = self.i2s(ix)
        lhs = v(torch.FloatTensor(s)); rhs = 0.
        if self.is_interior(ix):            
            fun = lambda a: self.q_val(list(ix), a, self.fd,v)
            rhs = self.min_a(fun)
            res = (rhs - lhs)
        else:
            rhs = self.term_cost(self.i2s(ix))
            res = (rhs - lhs)
        #end_time = time.time()
        #print(end_time - start_time)
        return res

    def solver(self,n_epoch = 500):
        ######### nn for value
        # Linear regression model
        value = nn.Sequential(
            nn.Linear(self.n_dim, 2*self.n_dim+10),
            #nn.functional.gelu(),
            nn.ReLU(),
            nn.Linear(2*self.n_dim+10, 2*self.n_dim+10),
            nn.ReLU(),
            nn.Linear(2*self.n_dim+10, 1)
        )   
        print(value)
        #loss
        def tot_loss():
            #start_time = time.time()
            out = 0.
            boundary = 0.
            for ix in deep_iter(*self.v_shape):
              if self.is_interior(ix):
                out += self.bellman(ix,value)**2
              else:
                boundary += (self.term_cost(self.i2s(ix)) - value(torch.FloatTensor(self.i2s(ix))))**2
            #end_time = time.time()
            #print(end_time - start_time)
            return 300*out + 10*boundary#/mdp.v_size_
          
        print_n = 10
        epoch_per_print= int(n_epoch/print_n)
        
        start_time = time.time()
        for epoch in range(n_epoch):
            #ipdb.set_trace()
            loss = tot_loss() #forward pass
            #backward propogation
            # optimizer
            lr = max(1/((epoch+2000.)), .00045)
            optimizer = torch.optim.SGD(value.parameters(), lr, momentum = .8) 
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            if (epoch) % epoch_per_print == 0:
              print('Epoch [{}/{}], Loss: {:.4f}'.format(
                      epoch+1, n_epoch, loss.item()))
            if loss.item()<0.002:
                break
        end_time = time.time()
        print('>>>time elapsed is: ' + str(end_time - start_time))
        return value
vi = solver_nn(n_dim=2, n_mesh=8, fd='ufd')
value = vi.solver(100)

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
figure = plt.figure()
axes = Axes3D(figure)
x = np.zeros(9)
y = np.zeros(9)
z = np.zeros([9,9])
z2 = np.zeros([9,9])
for ix_s in deep_iter(*vi.v_shape):
    s1 = vi.i2s(ix_s)
    x[ix_s[0]] = s1[0]
    y[ix_s[1]] = s1[1]
    #v1 = v[ix_s]
    v1 =value(torch.FloatTensor(s1)).item()
    v2 = vi.term_cost(s1)
    z[ix_s] = v1
    z2[ix_s] = v2
X1, X2 = np.meshgrid(x, y)
axes.plot_surface(X1, X2, z,cmap='rainbow')
figure2 = plt.figure()
axes2 = Axes3D(figure2)
axes2.plot_surface(X1, X2, z2,cmap='rainbow')
print('>>>sup norm error is: ' + str(np.max(abs(z-z2))))<jupyter_output>>>>sup norm error is: 0.2677435278892517
{'fraction_non_alphanumeric': 0.1489006401335931, 'fraction_numerical': 0.028945171166156414, 'mean_word_length': 1.7674239507123604, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '59123', 'license': 'MIT', 'n_tokens_mistral': 2825, 'n_tokens_neox': 2690, 'n_words': 560, 'path': 'wenhao/hjb_mdp_nn.ipynb', 'repo_name': 'songqsh/capstone20-wenhaoqiu', 'size': '7186'}