POMDP Solution implementation for 2-state problem described in Probabilistic Robotics by Thrun et al.

pomdp.py

from sets import Set
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm

p_z_gx = np.array([[0.7, 0.3],
                   [0.3, 0.7]])
p_x_g_u_z = np.zeros((2,3,2))
p_x_g_u_z[:,2,:] = np.array([[0.2, 0.8],
                             [0.8, 0.2]])
gamma = 1.0

r = np.array([[-100, 100, -1],
              [100, -50, -1]])

class Y_struct():
    def __init__(self):
        self.uv = Set()

    def add_uv(self, uv):
        self.uv.add(uv)

    def prune(self):
        array = np.array(list(self.uv))
        end = array[:,2]
        start = array[:,1]
        steps = (1.0/1000.0)*(end-start)
        lines = steps[:,None] *np.arange(1000) + start[:,None]
        lines_to_keep = np.unique(np.argmax(lines, axis=0))
        self.uv = Set([tuple(i) for i in array[lines_to_keep,:].tolist()])

    def plot(self):
        plt.figure(1)
        plt.clf()
        for uv in self.uv:
            plt.plot([1, 0], [uv[1], uv[2]], '-b')
        plt.show()

    def __str__(self):
        string = ''
        for uv in self.uv:
            string += str(uv[1]) + ' ' + str(uv[2]) + '\n'
        return string

    def copy(self):
        Y = Y_struct()
        Y.uv = self.uv.copy()
        return Y


def run_POMDP(T):
    Y = Y_struct()
    Y.uv.add((0, 0, 0))  # u, v1, v2
    for tau in tqdm(range(T)):
        v = np.zeros((len(Y.uv), 3, 2, 2))
        Ypr = Y_struct()
        for k , v_k in enumerate(Y.uv):
            for u in range(3):
                for z in range(2):
                    for j in range(2):
                        for i in range(2):
                            v[k, u, z, j] += v_k[i+1]*p_z_gx[z,i]*p_x_g_u_z[i,u,j]
        for u in range(3):
            for k1 in range(len(Y.uv)):
                for k2 in range(len(Y.uv)):
                    new_constraint = [u]
                    for i in range(2):
                        vpr = gamma * r[i, u] + v[k1, u, 0, i] + v[k2, u, 1, i]
                        new_constraint.append(vpr)
                    Ypr.add_uv(tuple(new_constraint))
        Ypr.prune()
        # print Ypr
        Y = Ypr.copy()
    Ypr.plot()


if __name__ == '__main__':
    run_POMDP(3)