<reponame>luigiberducci/gym-minigrid
from matplotlib import pyplot as plt

from algorithms.value_iteration import value_iteration
from gym_minigrid.envs import FullyObservableLavaGap
from gym_minigrid.minigrid import Wall, Lava, Goal
from gym_minigrid.wrappers import ReseedWrapper
import numpy as np

size = 10
n_directions = 4
n_actions = 3
seed = 0

state_space_shape = (size, size, n_directions)


def to_state(w, h, d):
    return np.ravel_multi_index([w, h, d], state_space_shape)


def to_coord(stateid):
    return np.unravel_index(stateid, state_space_shape)


def compute_probs(env):
    # prob distribution P: S x A x S -> [0, 1]
    ns = np.prod(state_space_shape)
    P = np.zeros(shape=(ns, n_actions, ns))
    # helpers
    terminal_mask = get_object_mask(env, [Lava, Wall, Goal])
    wall_mask = get_object_mask(env, [Wall])
    # the env is deterministic, then it is enough to iterate on s,a and compute next state s'. then set it to 1
    for state in range(ns):
        w, h, d = to_coord(state)
        if terminal_mask[w, h, d]:
            continue
        for a, action in enumerate(["left", "right", "forward"]):
            nw, nh, nd = w, h, d  # by default do nothing
            if action == "left":
                nd = (d - 1) % n_directions  # rotating left changes only agent's direction
            elif action == "right":
                nd = (d + 1) % n_directions  # rotating right changes only agent's direction
            elif action == "forward":
                nw = w + (d == 0) - (d == 2)  # increase w when pointing towards right (d=0), decr. when d==2
                nh = h + (d == 1) - (d == 3)  # increase h when pointing towards bottom (d=1), decr. when d==3
            # the action performed is not valid
            if not (0 <= nw < size) or not (0 <= nh <= size) or wall_mask[nw, nh, nd]:
                continue
            # convert state id
            newstate = to_state(nw, nh, nd)
            P[state, a, newstate] = 1.0  # deterministic next state
    return P


def get_object_mask(env, objects):
    mask = np.zeros(shape=state_space_shape).astype(bool)
    for w in range(size):
        for h in range(size):
            if any([isinstance(env.grid.get(w, h), obst) for obst in objects]):
                mask[w, h, :] = True
    return mask


def init_rewards(env, reward):
    potentials = np.zeros(shape=np.prod(state_space_shape))
    if reward == "default":
        rewards = -1 / 100 * 0.9 * np.ones(shape=np.prod(state_space_shape))
        for d in range(n_directions):
            gw, gh, gd = env.goal_pos[0], env.goal_pos[1], d
            state = to_state(gw, gh, gd)
            rewards[state] += 1
    elif reward == "sparse":
        rewards = np.zeros(shape=np.prod(state_space_shape))
        for d in range(n_directions):
            gw, gh, gd = env.goal_pos[0], env.goal_pos[1], d
            state = to_state(gw, gh, gd)
            rewards[state] += 1
    elif reward == "manattan":
        gw, gh = env.goal_pos
        obstacle_mask = get_object_mask(env, [Wall, Lava])
        rewards = np.zeros(shape=np.prod(state_space_shape))
        for state in range(np.prod(state_space_shape)):
            w, h, d = to_coord(state)
            rewards[state] = - (abs(w - gw) + abs(h - gh)) if not obstacle_mask[w, h, d] else -100
    elif reward == "sparse+manattan":
        rewards, _ = init_rewards(env, "sparse")
        potentials, _ = init_rewards(env, "manattan")
    return rewards, potentials


def rollout(policy, env, render=False):
    env.reset()
    state_dict = env.reset()
    w, h, d = state_dict["pos0"], state_dict["pos1"], state_dict["dir"]
    path_w, path_h = [w], [h]
    for i in range(100):
        stateid = to_state(w, h, d)
        action = policy[stateid]
        print(action)
        if action < 0:
            break
        # step
        next_state_dict, reward, done, info = env.step(action)
        if render:
            env.render()
        w, h, d = next_state_dict["pos0"], next_state_dict["pos1"], next_state_dict["dir"]
        path_w.append(w)
        path_h.append(h)
        if done:
            break
    return (path_w, path_h)


def vis_values(values, normalize=False):
    img_values = np.zeros(shape=(size, size))
    for state in range(np.prod(state_space_shape)):
        w, h, d = to_coord(state)
        img_values[w, h] += values[state]
    img_values /= n_directions  # average over the direction channel
    # normalize
    if normalize:
        minv, maxv = img_values.min(), img_values.max()
        img_values = (img_values - minv) / (maxv - minv)
    return img_values


def plot_values_and_path(values, path, title=""):
    img = vis_values(values)
    plt.imshow(img.T)
    plt.colorbar()
    plt.plot(path[0], path[1], c="red", marker="*")
    plt.title(title)


def main():
    # parameters
    gamma = 0.9
    epsilon = 1e-10
    max_iters = 1000
    # environment
    env = FullyObservableLavaGap(size=size, seed=seed)
    env = ReseedWrapper(env, seeds=[seed])
    env.reset()
    # value iterations
    probs = compute_probs(env)
    #
    policies = {}
    all_rewards = ["default", "sparse", "manattan", "sparse+manattan"]
    all_rewards = ["manattan"]
    for i, reward in enumerate(all_rewards):
        # prepare reward table
        rewards, potentials = init_rewards(env, reward)
        # run vi
        values, policy = value_iteration(probs, rewards, potentials, gamma, epsilon, max_iters)
        policies[reward] = policy
        # visualize
        plt.subplot(2, 2, i + 1)
        path = rollout(policy, env, render=False)
        plot_values_and_path(values, path, title=f"{reward}, gamma={gamma}")
    plt.show()
    # test
    for reward, policy in policies.items():
        _ = rollout(policy, env, render=True)

if __name__=="__main__":
    main()