Develop Reference

convergence issue encountered while implementing policy iteration from scratch

I am trying to implement policy iteration from scratch. I have a 2D grid world environment named GridWorld that returns the successor state and reward from a given action, and it also has a function that returns the transition probability. Below is my code for policy iteration:
import matplotlib
matplotlib.use('Agg')
import random
import numpy as np
import matplotlib.pyplot as plt
import gridworld
from tqdm import tqdm
class PolicyIteration:
def __init__(self, env, gamma):
self.env = env
self.num_states = self.env.num_states
self.num_actions = self.env.num_actions
self.max_num_steps = self.env.max_num_steps
self.gamma = gamma #discount factor
self.values = np.zeros(self.num_states) #Initialize `values` as zeros
self.policy = np.random.randint(0, self.num_actions, self.num_states)
def one_policy_evaluation(self):
"""
Runs one iteration of policy evaluation and updates the value function.
:return: the maximum change in value function
"""
delta = 0
for s in range(self.num_states):
v = self.values[s]
a = self.policy[s]
(s_new, r, _) = self.env.step(a)
p = self.env.p(s_new, s, a)
""" update V(s)"""
self.values[s] = np.sum(p * (r + self.gamma * self.values[s_new]))
delta = max(delta, abs(v - self.values[s]))
return delta
def run_policy_evaluation(self, tol = 1e-3):
"""
Runs policy evaluation until convergence.
:param tol: the tolerance level for convergence
:return: the number of iterations of policy evaluation until convergence
"""
delta = self.one_policy_evaluation()
delta_history = [delta]
while delta > tol:
delta = self.one_policy_evaluation()
delta_history.append(delta)
return len(delta_history)
def run_policy_improvement(self):
update_policy_count = 0
for s in range(self.num_states):
temp = self.policy[s]
v_list = np.zeros(self.num_actions)
for a in range(self.num_actions):
(s_new, r, _) = self.env.step(a)
p = self.env.p(s_new, s, a)
v_list[a] = np.sum(p * (r + self.gamma * self.values[s_new]))
self.policy[s] = np.argmax(v_list)
if temp != self.policy[s]:
update_policy_count += 1
return update_policy_count
def train(self, tol=1e-3, max_iters=100, plot=True):
eval_count = self.run_policy_evaluation(tol)
eval_count_history = [eval_count]
policy_change = self.run_policy_improvement()
policy_change_history = [policy_change]
epoch = 0
val_history= []
for i in tqdm(range(max_iters)):
epoch += 1
new_eval_count = self.run_policy_evaluation(tol)
new_policy_change = self.run_policy_improvement()
eval_count_history.append(new_eval_count)
policy_change_history.append(new_policy_change)
val_history.append(np.mean(self.values))
if new_policy_change == 0:
break
print(f'# epoch: {len(policy_change_history)}')
print(f'eval count = {eval_count_history}')
print(f'policy change = {policy_change_history}')
if plot is True:
plt.figure(dpi=200)
plt.plot(val_history)
plt.tight_layout()
plt.savefig('policy_iteration.png')
plt.show()
def main():
env = gridworld.GridWorld(hard_version=False)
agent = PolicyIteration(env, gamma=0.95)
agent.train()
if __name__ == '__main__':
main()
However, based on the figure generated, the sequence of values is oscillating up and down and never converges. I followed the algorithm in the Sutton book step by step, and can't find any issues with my code yet:
Any help is greatly appreciated!

JAX Tridiagonal Jacobians

What is the most efficient implementation of a scalable autonomous tridiagonal system using JAX?
import functools as ft
import jax as jx
import jax.numpy as jnp
import jax.random as jrn
import jax.lax as jlx
def make_T(m):
# Create a psuedo-random tridiagonal Jacobian and store band
T = jnp.zeros((3,m), dtype='f8')
T = T.at[0, 1: ].set(jrn.normal(jrn.PRNGKey(0), shape=(m-1,)))
T = T.at[1, : ].set(jrn.normal(jrn.PRNGKey(1), shape=(m ,)))
T = T.at[2, :-1].set(jrn.normal(jrn.PRNGKey(2), shape=(m-1,)))
return T
def make_y(m):
# Create a pseudo-random state array
y = jrn.normal(jrn.PRNGKey(3), shape=(m ,))
return y
def calc_f_base(y, T):
# Calculate the rate given the current state
f = T[1,:]*y
f = f.at[ 1: ].set(f[ 1: ]+T[0, 1: ]*y[ :-1])
f = f.at[ :-1].set(f[ :-1]+T[2, :-1]*y[ 1: ])
return f
m = 2**22 # potentially exhausts resources
T = make_T(m)
y = make_y(m)
calc_f = ft.partial(calc_f_base, T=T)
Using jax.jacrev or jax.jacfwd will generate a full Jacobian which limits the size of the system.
One attempt to overcome this limitation is as follows
#ft.partial(jx.jit, static_argnums=(0,))
def calc_jacfwd_trid(calc_f, y):
# Determine the Jacobian (forward-mode) tridiagonal band
def scan_body(carry, i):
t, T = carry
t = t.at[i ].set(1.0)
f, dfy = jx.jvp(calc_f, (y,), (t,))
T = T.at[2,i-1].set(dfy[i-1])
T = T.at[1,i ].set(dfy[i ])
T = T.at[0,i+1].set(dfy[i+1])
t = t.at[i-1].set(0.0)
return (t, T), None
# Initialise
m = y.size
t = jnp.zeros_like(y)
T = jnp.zeros((3,m), dtype=y.dtype)
# Differentiate wrt y[0]
t = t.at[0].set(1.0)
f, dfy = jx.jvp(calc_f, (y,), (t,))
idxs = jnp.array([1,0]), jnp.array([0,1])
T = T.at[idxs].set(dfy[0:2])
# Differentiate wrt y[1:-1]
(t, T), empty = jlx.scan(scan_body, (t,T), jnp.arange(1,m-1))
# Differentiate wrt y[-1]
t = t.at[m-2:].set(jnp.array([0.0,1.0]))
f, dfy = jx.jvp(calc_f, (y,), (t,))
idxs = jnp.array([2,1]), jnp.array([m-2,m-1])
T = T.at[idxs].set(dfy[-2:])
return T
which permits
T = jacfwd_trid(calc_f, y)
df = jrn.normal(jrn.PRNGKey(4), shape=y.shape)
dx = jlx.linalg.tridiagonal_solve(*T,df[:,None]).flatten()
Is there a better approach and/or can the time complexity of calc_jacfwd_trid be reduced further?
EDIT
The following implementation is more compact, but run times are slightly slower
#ft.partial(jx.jit, static_argnums=(0,))
def calc_jacfwd_trid_map(calc_f, y):
# Determine the Jacobian (forward-mode) tridiagonal band with lax map
def map_body(i, t):
t = t.at[i-1].set(0.0)
f, dfy = jx.jvp(calc_f, (y,), (t,))
im1 = jnp.where(i > 0, i-1, 0)
Ti = jlx.dynamic_slice(dfy, (im1,), (3,))
Ti = jnp.where(i > 0, Ti, jnp.roll(Ti, shift=+1))
Ti = jnp.where(i < m-1, Ti, jnp.roll(Ti, shift=-1))
t = t.at[i ].set(1.0)
return Ti
# Initialise
m = y.size
t = jnp.zeros_like(y)
# Differentiate wrt y[:]
T = jlx.map(lambda i : map_body(i, t=t), jnp.arange(m))
# Correct the orientation of T
T = T.transpose()
T = jnp.flip(T, axis=0)
T = T.at[0,:].set(jnp.roll(T[0,:], shift=+1))
T = T.at[2,:].set(jnp.roll(T[2,:], shift=-1))
return T

Need help understanding what this data analysis code (for optics lab) is doing

I am working in an undergraduate optics lab and we were given a few codes for data analysis. I need help understanding what two of them are doing exactly. To give some background on the experiment: light from a pulsed laser is injected into the back of a PTFE diffuser. Across from the diffuser is a photomultiplier tube that detects the optical signal. During the experiment, the diffuser is rotated and data is collected about the pulse waveform for every three degrees. I am just not sure what the code does to interpret this data, I hope that makes sense. Thank you in advance.
Here is the first code:
import numpy as np
DEFAULT_FMAX = 0.2
class Waveform(object):
def __init__(self, X, Y, fmax=None, peaknorm=False, areanorm=False):
self.X = np.copy(X)
self.Y = np.copy(Y)
if fmax is not None:
self.lowpass_filter(fmax)
if peaknorm:
self.normalise_peak()
if areanorm:
self.normalise_area()
def lowpass_filter(self, fmax):
# Fourier transform
freq_Y = np.fft.rfft(self.Y)
freq_X = np.fft.rfftfreq(self.Y.size, d=((max(self.X)-min(self.X))/self.Y.size))
# Apply max frequency cut
freq_Y[freq_X>fmax] = 0.0
# Inverse fourier transform
self.Y = np.fft.irfft(freq_Y)
def normalise_peak(self):
self.Y /= abs(np.min(self.Y))
def normalise_area(self):
self.Y /= np.sum(self.Y)
#property
def fft(self):
freq_Y = np.fft.rfft(self.Y)
freq_X = np.fft.rfftfreq(self.Y.size, d=((max(self.X)-min(self.X))/self.Y.size))
return freq_X, freq_Y
#property
def pedestal(self):
return np.mean(self.Y[0:100])
#property
def amplitude(self):
return abs(min(self.Y) - self.pedestal)
#property
def peak_index(self):
return np.argmin(self.Y)
#property
def peak_time(self):
return self.X[self.peak_index]
def _find_half_maxima_indices(self):
peak = self.peak_index
Y = np.copy(self.Y - self.pedestal)
Y += (self.amplitude / 2.0)
# find first half maximum
try:
first = np.argmax(Y[:peak] < 0)
except ValueError:
# can fail eg if peak == 0
first = peak
try:
second = peak + np.argmax(Y[peak:] > 0)
except ValueError:
# can fail eg if peak == last bin
second = peak
return first, second
#property
def fwhm(self):
return self.rise + self.fall
#property
def fall(self):
_, second = self._find_half_maxima_indices()
fall = self.X[second] - self.peak_time
return fall
#property
def rise(self):
first, _ = self._find_half_maxima_indices()
rise = self.peak_time - self.X[first]
return rise
#property
def integralindex_high(self):
pedestal = self.pedestal
#integrate falling edge
for ii in range(self.peak_index, self.Y.shape[0] - 1):
if (self.Y[ii] - pedestal) > 0.0:
break
return ii
#property
def integralindex_low(self):
pedestal = self.pedestal
#integrate falling edge
ii = 0
for ii in reversed(range(0, self.peak_index)):
if (self.Y[ii] - pedestal) > 0.0:
break
return ii
#property
def area(self):
area = 0.0
pedestal = self.pedestal
#integrate falling edge
for ii in range(self.peak_index, self.integralindex_high):
area += (self.Y[ii] - pedestal) * (self.X[ii + 1] - self.X[ii])
#integrate rising edge
for ii in reversed(range(self.integralindex_low, self.peak_index)):
area += (self.Y[ii] - pedestal) * (self.X[ii + 1] - self.X[ii])
return abs(area)
And the second:
import ROOT
import numpy as np
import datetime
from collections import namedtuple
from collections import defaultdict
from .waveform import Waveform
class ScanPoint:
def __init__(self, time_epoch, coord_angle, coord_y, lab_temp, lab_humid, dt, samples_PMT,
samples_PD, num_entry):
self.time_epoch = time_epoch
self.coord_angle = coord_angle
self.coord_y = coord_y
self.lab_temp = lab_temp
self.lab_humid = lab_humid
self.samples_PMT = np.array(list(samples_PMT), dtype=float)
self.samples_PD = np.array(list(samples_PD), dtype=float)
self.axis_time = np.arange(len(samples_PMT), dtype=float) * dt * 1.e9
self.num_entry = num_entry
def getWaveformPMT(self, fmax=None, peaknorm=False, areanorm=False):
return Waveform(self.axis_time, self.samples_PMT, fmax, peaknorm, areanorm)
def getWaveformPD(self, fmax=None, peaknorm=False, areanorm=False):
return Waveform(self.axis_time, self.samples_PD, fmax, peaknorm, areanorm)
def getDatetime(self):
return datetime.datetime.fromtimestamp(self.time_epoch)
# Tuple containing metadata and scanpoints for a full scan
DiffuserScan = namedtuple("DiffuserScan", ["ID_diffuser",
"ID_PMT",
"ID_PD",
"ID_lightsource",
"ID_experimentalist",
"notes",
"pulse_rate",
"pulse_N",
"scanpoints"])
def readscan(filename, treename):
file = ROOT.TFile.Open(filename, "READ")
tree = getattr(file, treename)
tree.GetEntry(0)
formatversion = (
int(tree.version_major),
int(tree.version_minor),
int(tree.version_patch),
)
# Check that input file format is implemented
if formatversion[:2] == (1, 0):
diffuserscan = DiffuserScan(
str(tree.ID_diffuser),
str(tree.ID_PMT),
str(tree.ID_PD),
str(tree.ID_lightsource),
str(tree.ID_experimentalist),
str(tree.notes),
float(tree.pulse_rate),
int(tree.pulse_N),
[],
)
for i, entry in enumerate(tree):
scanpoint = ScanPoint(entry.time_epoch,
entry.coord_angle,
entry.coord_y,
entry.lab_temp,
entry.lab_humid,
entry.dt,
entry.waveform_PMT,
entry.waveform_PD,
i)
diffuserscan.scanpoints.append(scanpoint)
else:
raise NotImplementedError("Input format version not implemented")
return diffuserscan
def combinefiles(scans):
combined_scan = scans[0]
compare_vars = ["ID_PMT", "ID_PD", "ID_lightsource", "pulse_rate", "pulse_N"]
for scan in scans[1:]:
combined_scan.scanpoints.extend(scan.scanpoints)
if scan.ID_diffuser != combined_scan.ID_diffuser:
combined_scan.ID_diffuser = "multi"
for var in compare_vars:
if getattr(scan, var) != getattr(combined_scan, var):
print('WARNING: Variable "{}" does not match between files! Should we be comparing
them?'.format(var))
return combined_scan

Multithreading a numpy nditerator

For an MCMC implementation, I want to calculate the covariance tensor C in numpy.
Working Single-Threaded Code
The distance between two elements is based on the distance between their indices. For reference, here is the working single threaded code (with an example distance):
import numpy as np
#set size, dimensions, etc
size = 20
ndim = 2
shape = (size,)*ndim*2
#initialize tensor
C = np.zeros(shape)
#example distance
dist = lambda x, y: np.sqrt(np.sum((x-y)**2))
#this runs as a class method, so please forgive my sloppy coding here
def update_tensor():
it = np.nditer(C, flags=['multi_index'], op_flags=['readwrite'])
while not it.finished:
idx = np.array(it.multi_index)
it[0] = dist(idx[:idx.shape[0]//2], idx[idx.shape[0]//2:])
it.iternext()
update_tensor()
Solution Attempt
Now the issue is, that while applying C to a matrix x is a multithreaded operation:
x = np.random.standard_normal((size,)*ndim)
result = np.tensordot(C, x, axes=ndim)
caculating the entries of C is not. My idea was, to split C after initialization along its first axis and iterate over the chunks separately:
import multiprocessing
def _calc_distances(C):
'Calculate distances of submatrices'
it = np.nditer(C, flags=['multi_index'], op_flags=['readwrite'])
while not it.finished:
idx = np.array(it.multi_index)
it[0] = dist(idx[:idx.shape[0]//2], idx[idx.shape[0]//2:])
it.iternext()
return C
def update_tensor(C):
'Updates Covariance Operator'
#Multicore Processing
n_processes = multiprocessing.cpu_count()
Chunks = [
C[i*C.shape[0]//n_processes:(i+1)*C.shape[0]//n_processes] for i in range(0, n_processes-1)
]
Chunks.append(C[C.shape[0]//n_processes*(n_processes-1):])
with multiprocessing.Pool(n_processes+1) as p:
#map and stitch together
C = np.concatenate(
p.map(_calc_distances, Chunks)
)
But this fails, because the indeces of the submatrices change.
Question
Is there a nicer solution to this? How do I fix the index issue? Probably the nicest way would be to just iterate over parts of the array with threads sharing the data of C. Is that possible?
Q/A
Q: Do you have to use a numpy iterator?
A: No, it’s nice, but I can give up on that.

Worked like this. Just going to post the class here.
Benchmarks
CPU: Intel Core i5-6300U#2.5GHz, boosting to ~2.9GHz
Windows 10 64-bit, Python 3.7.4, Numpy 1.17
Pro: Less compute time
Con: Uses a little more RAM; somewhat complicated code.
Working Multi-Threaded Code
import multiprocessing
import numpy as np
class CovOp(object):
'F[0,1]^ndim->C[0,1]^ndim'
def f(self, r):
return np.exp(-r/self.ro)#(1 + np.sqrt(3)*r / self.ro) * np.exp(-np.sqrt(3) * r / self.ro)
def dist(self, x,y):
return np.sum((x-y)**2)
def __init__(self, ndim, size, sigma=1, ro=1):
self.tensor_cached = False
self.inverse_cached = False
self.ndim = ndim
self.size = size
self.shape = (size,)*ndim*2
self.C = np.zeros(self.shape)
self.Inv = np.zeros(self.shape)
self.ro = ro * size
self.sigma = sigma
def __call__(self, x):
if not self.tensor_cached:
self.update_tensor
if self.ndim == 0:
return self.sigma * self.C * x
elif self.ndim == 1:
return self.sigma * np.dot(self.C, x)
return self.sigma * np.tensordot(self.C, x, axes=self.ndim)
def _calc_distances(self, Chunk:tuple):
'Calculate distances of submatrices'
C, offset = Chunk
it = np.nditer(C, flags=['multi_index'], op_flags=['readwrite'])
while not it.finished:
idx = np.array(it.multi_index)
idx[0]+=offset
d = self.dist(idx[:idx.shape[0]//2], idx[idx.shape[0]//2:])
it[0] = self.f(d)
it.iternext()
return C
def update_tensor(self):
'Updates Covariance Operator'
#Multicore Processing
n_processes = multiprocessing.cpu_count()
Chunks = [
(
self.C[i*self.C.shape[0]//n_processes:(i+1)*self.C.shape[0]//n_processes],
i*self.C.shape[0]//n_processes) for i in range(0, n_processes-1)
]
Chunks.append((
self.C[self.C.shape[0]//n_processes*(n_processes-1):],
self.C.shape[0]//n_processes*(n_processes-1)
)
)
with multiprocessing.Pool(n_processes+1) as p:
self.C = np.concatenate(
p.map(self._calc_distances, Chunks)
)
self.tensor_cached = True
#missing cholesky decomposition
def update_inverse(self):
if self.ndim==1:
self.Inv = np.linalg.inv(self.C)
elif self.ndim>1:
self.Inv = np.linalg.tensorinv(self.C)
else:
self.Inv = 1/self.C
self.inverse_cached = True
def inv(self, x):
if self.ndim == 0:
return self.Inv * x / self.sigma
elif self.ndim == 1:
return np.dot(self.Inv, x) / self.sigma
return np.tensordot(self.Inv, x) / self.sigma
if __name__=='__main__':
size = 30
ndim = 2
depth = 1
Cov = CovOp(ndim, size, 1, .2)
import time
n_tests = 5
t_start = time.perf_counter()
for i in range(n_tests):
Cov.update_tensor()
t_stop = time.perf_counter()
dt_new = t_stop - t_start
print(
'''Benchmark; NDim: %s, Size: %s NTests: %s
Mean time per test:
Multithreaded %ss'''%(ndim, size, n_tests, dt_new/n_tests)
)

Calculating value of x for which f(x) is given

This may be a simple question, but I can't wrap my head around it. Say,
nb = 100
tb = 25
ns = 90
ts = 15
A0 = 1
S_norm = 0.4
R = tb/ts
y_meas = (ns-nb/R)/A0
sigma_meas = np.sqrt(ns+(nb+1)/R**2)/A0
def likelihood(y):
def func_likelihood(x):
return np.exp(-0.5*(((x/A0)/S_norm)**2 + ((y-y_meas*A0/np.exp(x))/sigma_meas)**2))
return (scipy.integrate.quad(func_likelihood,-10,10))[0]
Is there an easy way to determine value of y such that likelihood(y) is 0.025?

Have you tried scipy.optimize.fsolve?
from scipy.optimize import fsolve
def func(y):
return likelihood(y) - 0.025
print(fsolve(func, 0))
Result:
[-7.3920919]