import dill
import haiku as hk
import jax
from jax.experimental import optix
import jax.numpy as jnp

from dataset import load_data


MINERL_ENV = 'MineRLTreechopVectorObf-v0'
PARAMS_FILENAME = 'bc_params_treechop.pkl'


class PovStack(hk.Module):
    """ PovStack is a module for processing the point-of-view image data that
    comes from the agent's viewport. This input is in NHWC format for a shape
    of (N, 64, 64, 3).

    This model is inspired from
    https://github.com/minerllabs/baselines/blob/master/general/chainerrl/baselines/behavioral_cloning.py
    """
    def __init__(self, name=None):
        super().__init__(name=name)
        conv_0 = hk.Conv2D(output_channels=32,
                           kernel_shape=(8, 8),
                           stride=4,
                           padding='SAME',
                           name='conv_0')
        layer_0 = (conv_0, jax.nn.relu)

        conv_1 = hk.Conv2D(output_channels=64,
                           kernel_shape=(4, 4),
                           stride=2,
                           padding='SAME',
                           name='conv_1')
        layer_1 = (conv_1, jax.nn.relu)

        conv_2 = hk.Conv2D(output_channels=64,
                           kernel_shape=(3, 3),
                           stride=1,
                           padding='SAME',
                           name='conv_2')
        layer_2 = (conv_2, jax.nn.relu)

        layer_3 = (hk.Flatten(),
                   hk.Linear(512, name='fc_0'), jax.nn.relu,
                   hk.Linear(128, name='fc_1'), jax.nn.relu)

        self.layers = layer_0 + layer_1 + layer_2 + layer_3

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x


class VectorStack(hk.Module):
    """ VectorStack is a module for processing the obfuscated "vector" data that
    is included in the agent's observation. This is a densely encoded form of
    the discrete information regarding the state of the agent other than the
    viewport, e.g. current inventory. The input is of shape (N, 64)
    """
    def __init__(self, name=None):
        super().__init__(name=name)
        layer_0 = (hk.Linear(32, name='fc_0'), jax.nn.relu)

        self.layers = layer_0

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x


def behavioral_cloning(batch):
    """ The full forward model definition """
    x_0 = PovStack(name='pov_stack')(batch[0])
    x_1 = VectorStack(name='vector_stack')(batch[1])
    x = jnp.concatenate((x_0, x_1), axis=1)
    return jnp.tanh(hk.Linear(64)(x))


@jax.jit
def mse_loss(logits, labels):
    """ Mean Squared Error loss """
    return jnp.mean(jnp.power(logits - labels, 2))


def main():
    net = hk.transform(behavioral_cloning)
    opt = optix.adam(0.001)

    @jax.jit
    def loss(params, batch):
        """ The loss criterion for our model """
        logits = net.apply(params, None, batch)
        return mse_loss(logits, batch[2])

    @jax.jit
    def update(opt_state, params, batch):
        grads = jax.grad(loss)(params, batch)
        updates, opt_state = opt.update(grads, opt_state)
        params = optix.apply_updates(params, updates)
        return params, opt_state

    @jax.jit
    def accuracy(params, batch):
        """ Simply report the loss for the current batch """
        logits = net.apply(params, None, batch)
        return mse_loss(logits, batch[2])

    train_dataset, val_dataset = load_data(MINERL_ENV,
                                           batch_size=32, epochs=100)

    rng = jax.random.PRNGKey(2020)
    batch = next(train_dataset)
    params = net.init(rng, batch)
    opt_state = opt.init(params)

    for i, batch in enumerate(train_dataset):
        params, opt_state = update(opt_state, params, batch)
        if i % 1000 == 0:
            print(accuracy(params, val_dataset))
        
        if i % 10000 == 0:
            with open(PARAMS_FILENAME, 'wb') as fh:
                dill.dump(params, fh)

    with open(PARAMS_FILENAME, 'wb') as fh:
        dill.dump(params, fh)



if __name__ == '__main__':
    main()

[((2536, 2571), 'jax.numpy.concatenate', 'jnp.concatenate', (['(x_0, x_1)'], {'axis': '(1)'}), '((x_0, x_1), axis=1)\n', (2551, 2571), True, 'import jax.numpy as jnp\n'), ((2762, 2794), 'haiku.transform', 'hk.transform', (['behavioral_cloning'], {}), '(behavioral_cloning)\n', (2774, 2794), True, 'import haiku as hk\n'), ((2805, 2822), 'jax.experimental.optix.adam', 'optix.adam', (['(0.001)'], {}), '(0.001)\n', (2815, 2822), False, 'from jax.experimental import optix\n'), ((3484, 3532), 'dataset.load_data', 'load_data', (['MINERL_ENV'], {'batch_size': '(32)', 'epochs': '(100)'}), '(MINERL_ENV, batch_size=32, epochs=100)\n', (3493, 3532), False, 'from dataset import load_data\n'), ((3587, 3611), 'jax.random.PRNGKey', 'jax.random.PRNGKey', (['(2020)'], {}), '(2020)\n', (3605, 3611), False, 'import jax\n'), ((662, 757), 'haiku.Conv2D', 'hk.Conv2D', ([], {'output_channels': '(32)', 'kernel_shape': '(8, 8)', 'stride': '(4)', 'padding': '"""SAME"""', 'name': '"""conv_0"""'}), "(output_channels=32, kernel_shape=(8, 8), stride=4, padding='SAME',\n    name='conv_0')\n", (671, 757), True, 'import haiku as hk\n'), ((920, 1015), 'haiku.Conv2D', 'hk.Conv2D', ([], {'output_channels': '(64)', 'kernel_shape': '(4, 4)', 'stride': '(2)', 'padding': '"""SAME"""', 'name': '"""conv_1"""'}), "(output_channels=64, kernel_shape=(4, 4), stride=2, padding='SAME',\n    name='conv_1')\n", (929, 1015), True, 'import haiku as hk\n'), ((1178, 1273), 'haiku.Conv2D', 'hk.Conv2D', ([], {'output_channels': '(64)', 'kernel_shape': '(3, 3)', 'stride': '(1)', 'padding': '"""SAME"""', 'name': '"""conv_2"""'}), "(output_channels=64, kernel_shape=(3, 3), stride=1, padding='SAME',\n    name='conv_2')\n", (1187, 1273), True, 'import haiku as hk\n'), ((2707, 2736), 'jax.numpy.power', 'jnp.power', (['(logits - labels)', '(2)'], {}), '(logits - labels, 2)\n', (2716, 2736), True, 'import jax.numpy as jnp\n'), ((3182, 3218), 'jax.experimental.optix.apply_updates', 'optix.apply_updates', (['params', 'updates'], {}), '(params, updates)\n', (3201, 3218), False, 'from jax.experimental import optix\n'), ((4073, 4094), 'dill.dump', 'dill.dump', (['params', 'fh'], {}), '(params, fh)\n', (4082, 4094), False, 'import dill\n'), ((1438, 1450), 'haiku.Flatten', 'hk.Flatten', ([], {}), '()\n', (1448, 1450), True, 'import haiku as hk\n'), ((1471, 1498), 'haiku.Linear', 'hk.Linear', (['(512)'], {'name': '"""fc_0"""'}), "(512, name='fc_0')\n", (1480, 1498), True, 'import haiku as hk\n'), ((1532, 1559), 'haiku.Linear', 'hk.Linear', (['(128)'], {'name': '"""fc_1"""'}), "(128, name='fc_1')\n", (1541, 1559), True, 'import haiku as hk\n'), ((2173, 2199), 'haiku.Linear', 'hk.Linear', (['(32)'], {'name': '"""fc_0"""'}), "(32, name='fc_0')\n", (2182, 2199), True, 'import haiku as hk\n'), ((2592, 2605), 'haiku.Linear', 'hk.Linear', (['(64)'], {}), '(64)\n', (2601, 2605), True, 'import haiku as hk\n'), ((3077, 3091), 'jax.grad', 'jax.grad', (['loss'], {}), '(loss)\n', (3085, 3091), False, 'import jax\n'), ((3998, 4019), 'dill.dump', 'dill.dump', (['params', 'fh'], {}), '(params, fh)\n', (4007, 4019), False, 'import dill\n')]