"""
VAE on the swirl task.

Basically, VAEs don't work. It's probably because the prior isn't very good
and/or because the learning signal is pretty weak when both the encoder and
decoder change quickly. However, I tried also alternating between the two,
and that didn't seem to help.
"""
from torch.distributions import Normal
from torch.optim import Adam
import torch
import numpy as np
import matplotlib.pyplot as plt
from torch import nn as nn
import railrl.torch.pytorch_util as ptu

SWIRL_RATE = 1
T = 10
BS = 128
N_BATCHES = 2000
N_VIS = 1000
HIDDEN_SIZE = 32
VERBOSE = False


def swirl_data(batch_size):
    t = np.random.uniform(size=batch_size, low=0, high=T)
    x = t * np.cos(t * SWIRL_RATE) / T
    y = t * np.sin(t * SWIRL_RATE) / T
    data = np.array([x, y]).T
    noise = np.random.randn(batch_size, 2) / (T * 2)
    return data + noise, t.reshape(-1, 1)


def swirl_t_to_data(t):
    x = t * np.cos(t * SWIRL_RATE) / T
    y = t * np.sin(t * SWIRL_RATE) / T
    return np.array([x, y]).T


def kl_to_prior(means, log_stds, stds):
    """
    KL between a Gaussian and a standard Gaussian.

    https://stats.stackexchange.com/questions/60680/kl-divergence-between-two-multivariate-gaussians
    """
    return 0.5 * (
            - 2 * log_stds  # log std_prior = 0
            - 1  # d = 1
            + stds ** 2
            + means ** 2
    )


class Encoder(nn.Sequential):
    def encode(self, x):
        return self.get_encoding_and_suff_stats(x)[0]

    def get_encoding_and_suff_stats(self, x):
        output = self(x)
        means, log_stds = (
            output[:, 0:1], output[:, 1:2]
        )
        stds = log_stds.exp()
        epsilon = ptu.Variable(torch.randn(*means.size()))
        latents = epsilon * stds + means
        latents = latents
        return latents, means, log_stds, stds


class Decoder(nn.Sequential):
    def decode(self, latents):
        output = self(latents)
        means, log_stds = output[:, 0:2], output[:, 2:4]
        distribution = Normal(means, log_stds.exp())
        return distribution.sample()


def t_to_xy(t):
    if len(t.shape) == 2:
        t = t[:, 0]
    x = t * np.cos(t * SWIRL_RATE) / T
    y = t * np.sin(t * SWIRL_RATE) / T
    return np.array([x, y]).T


def pretrain_encoder(encoder, opt):
    losses = []
    for _ in range(1000):
        x_np, y_np = swirl_data(BS)
        x = ptu.np_to_var(x_np)
        y = ptu.np_to_var(y_np)
        y_hat = encoder.encode(x)
        loss = ((y_hat - y) ** 2).mean()
        opt.zero_grad()
        loss.backward()
        opt.step()

        losses.append(loss.data.numpy())

    if VERBOSE:
        x_np, y_np = swirl_data(N_VIS)
        x = ptu.np_to_var(x_np)
        y_hat = encoder.encode(x)
        y_hat_np = y_hat.data.numpy()
        x_hat_np = t_to_xy(y_hat_np[:, 0])

        plt.subplot(2, 1, 1)
        plt.plot(np.array(losses))
        plt.title("Training Loss")

        plt.subplot(2, 1, 2)
        plt.plot(x_np[:, 0], x_np[:, 1], '.')
        plt.plot(x_hat_np[:, 0], x_hat_np[:, 1], '.')
        plt.title("Samples")
        plt.legend(["Samples", "Estimates"])
        plt.show()


def train_encoder(encoder, decoder, encoder_opt):
    batch, true_latents = swirl_data(BS)
    batch = ptu.np_to_var(batch)

    latents, means, log_stds, stds = encoder.get_encoding_and_suff_stats(
        batch
    )
    kl = kl_to_prior(means, log_stds, stds)

    latents = encoder.encode(batch)
    decoder_output = decoder(latents)
    decoder_means = decoder_output[:, 0:2]
    decoder_log_stds = decoder_output[:, 2:4]
    distribution = Normal(decoder_means, decoder_log_stds.exp())
    reconstruction_log_prob = distribution.log_prob(batch).sum(dim=1)

    # elbo = - kl + reconstruction_log_prob
    # loss = - elbo.mean()
    loss = - reconstruction_log_prob.mean()
    # This is the second place where we cheat:
    latent_loss = ((ptu.np_to_var(true_latents) - latents) ** 2).mean()
    loss = loss# + latent_loss
    encoder_opt.zero_grad()
    loss.backward()
    encoder_opt.step()
    return loss


def train_decoder(encoder, decoder, decoder_opt):
    batch, true_latents = swirl_data(BS)
    batch = ptu.np_to_var(batch)

    latents = encoder.encode(batch)
    decoder_output = decoder(latents)
    decoder_means = decoder_output[:, 0:2]
    decoder_log_stds = decoder_output[:, 2:4]
    distribution = Normal(decoder_means, decoder_log_stds.exp())
    reconstruction_log_prob = distribution.log_prob(batch).sum(dim=1)

    loss = - reconstruction_log_prob.mean()
    decoder_opt.zero_grad()
    loss.backward()
    decoder_opt.step()
    return loss


def train_alternating(*_):
    encoder = Encoder(
        nn.Linear(2, HIDDEN_SIZE),
        nn.ReLU(),
        nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE),
        nn.ReLU(),
        nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE),
        nn.ReLU(),
        nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE),
        nn.ReLU(),
        nn.Linear(HIDDEN_SIZE, 2),
    )
    encoder_opt = Adam(encoder.parameters())
    decoder = Decoder(
        nn.Linear(1, HIDDEN_SIZE),
        nn.ReLU(),
        nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE),
        nn.ReLU(),
        nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE),
        nn.ReLU(),
        nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE),
        nn.ReLU(),
        nn.Linear(HIDDEN_SIZE, 4),
    )
    decoder_opt = Adam(decoder.parameters())

    encoder_losses = []
    decoder_losses = []
    for _ in range(100):
        for _ in range(N_BATCHES):
            encoder_losses.append(
                train_encoder(encoder, decoder, encoder_opt).data.numpy()
            )
        for _ in range(N_BATCHES):
            decoder_losses.append(
                train_decoder(encoder, decoder, decoder_opt).data.numpy()
            )

    # Visualize
    vis_samples_np, true_latents_np = swirl_data(N_VIS)
    vis_samples = ptu.np_to_var(vis_samples_np)
    true_xy_mean_np = t_to_xy(true_latents_np)
    latents = encoder.encode(vis_samples)
    reconstructed_samples = decoder.decode(latents).data.numpy()
    generated_samples = decoder.decode(
        ptu.Variable(torch.randn(*latents.shape))
    ).data.numpy()

    plt.subplot(2, 2, 1)
    plt.plot(np.array(encoder_losses))
    plt.title("Encoder Loss")
    plt.subplot(2, 2, 2)
    plt.plot(np.array(decoder_losses))
    plt.title("Decoder Loss")

    plt.subplot(2, 3, 4)
    plt.plot(generated_samples[:, 0], generated_samples[:, 1], '.')
    plt.title("Generated Samples")
    plt.subplot(2, 3, 5)
    plt.plot(reconstructed_samples[:, 0], reconstructed_samples[:, 1], '.')
    estimated_means = t_to_xy(latents.data.numpy())
    # plt.plot(estimated_means[:, 0], estimated_means[:, 1], '.')
    plt.title("Reconstruction")
    # plt.legend(["Samples", "Projected Latents"])
    plt.subplot(2, 3, 6)
    plt.plot(vis_samples_np[:, 0], vis_samples_np[:, 1], '.')
    plt.plot(true_xy_mean_np[:, 0], true_xy_mean_np[:, 1], '.')
    plt.title("Original Samples")
    plt.legend(["Original", "True means"])
    plt.show()


def train():
    encoder = Encoder(
        nn.Linear(2, HIDDEN_SIZE),
        nn.ReLU(),
        nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE),
        nn.ReLU(),
        nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE),
        nn.ReLU(),
        nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE),
        nn.ReLU(),
        nn.Linear(HIDDEN_SIZE, 2),
    )
    encoder_opt = Adam(encoder.parameters())
    # This is the first place that we cheat. However, this pretraining isn't
    # needed if you just add the loss to the training (see below)
    # pretrain_encoder(encoder, encoder_opt)
    decoder = Decoder(
        nn.Linear(1, HIDDEN_SIZE),
        nn.ReLU(),
        nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE),
        nn.ReLU(),
        nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE),
        nn.ReLU(),
        nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE),
        nn.ReLU(),
        nn.Linear(HIDDEN_SIZE, 4),
    )
    decoder_opt = Adam(decoder.parameters())
    print("Done training encoder")

    losses = []
    kls = []
    log_probs = []
    for _ in range(N_BATCHES):
        batch, true_latents = swirl_data(BS)
        batch = ptu.np_to_var(batch)

        latents, means, log_stds, stds = encoder.get_encoding_and_suff_stats(
            batch
        )
        kl = kl_to_prior(means, log_stds, stds)


        latents = encoder.encode(batch)
        # decoder_output = decoder(latents.detach())
        decoder_output = decoder(latents)
        decoder_means = decoder_output[:, 0:2]
        decoder_log_stds = decoder_output[:, 2:4]
        distribution = Normal(decoder_means, decoder_log_stds.exp())
        reconstruction_log_prob = distribution.log_prob(batch).sum(dim=1)

        elbo = - kl + reconstruction_log_prob
        loss = - elbo.mean()
        # This is the second place where we cheat:
        latent_loss = ((ptu.np_to_var(true_latents) - latents) ** 2).mean()
        loss = loss + latent_loss
        decoder_opt.zero_grad()
        encoder_opt.zero_grad()
        loss.backward()
        decoder_opt.step()
        encoder_opt.step()

        losses.append(loss.data.numpy())
        kls.append(kl.mean().data.numpy())
        log_probs.append(reconstruction_log_prob.mean().data.numpy())

    # Visualize
    vis_samples_np, true_latents_np = swirl_data(N_VIS)
    vis_samples = ptu.np_to_var(vis_samples_np)
    true_xy_mean_np = t_to_xy(true_latents_np)
    latents = encoder.encode(vis_samples)
    reconstructed_samples = decoder.decode(latents).data.numpy()
    generated_samples = decoder.decode(
        ptu.Variable(torch.randn(*latents.shape))
    ).data.numpy()

    plt.subplot(2, 3, 1)
    plt.plot(np.array(losses))
    plt.title("Training Loss")
    plt.subplot(2, 3, 2)
    plt.plot(np.array(kls))
    plt.title("KLs")
    plt.subplot(2, 3, 3)
    plt.plot(np.array(log_probs))
    plt.title("Log Probs")

    plt.subplot(2, 3, 4)
    plt.plot(generated_samples[:, 0], generated_samples[:, 1], '.')
    plt.title("Generated Samples")
    plt.subplot(2, 3, 5)
    plt.plot(reconstructed_samples[:, 0], reconstructed_samples[:, 1], '.')
    estimated_means = t_to_xy(latents.data.numpy())
    plt.plot(estimated_means[:, 0], estimated_means[:, 1], '.')
    plt.title("Reconstruction")
    plt.subplot(2, 3, 6)
    plt.plot(vis_samples_np[:, 0], vis_samples_np[:, 1], '.')
    plt.plot(true_xy_mean_np[:, 0], true_xy_mean_np[:, 1], '.')
    plt.title("Original Samples")
    plt.legend(["Original", "True means"])
    plt.show()


if __name__ == '__main__':
    train_alternating()
    # train()

[((27, 8, 27, 57), 'numpy.random.uniform', 'np.random.uniform', (), '', True, 'import numpy as np\n'), ((122, 12, 122, 32), 'railrl.torch.pytorch_util.np_to_var', 'ptu.np_to_var', ({(122, 26, 122, 31): 'batch'}, {}), '(batch)', True, 'import railrl.torch.pytorch_util as ptu\n'), ((150, 12, 150, 32), 'railrl.torch.pytorch_util.np_to_var', 'ptu.np_to_var', ({(150, 26, 150, 31): 'batch'}, {}), '(batch)', True, 'import railrl.torch.pytorch_util as ptu\n'), ((206, 18, 206, 47), 'railrl.torch.pytorch_util.np_to_var', 'ptu.np_to_var', ({(206, 32, 206, 46): 'vis_samples_np'}, {}), '(vis_samples_np)', True, 'import railrl.torch.pytorch_util as ptu\n'), ((214, 4, 214, 24), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(214, 16, 214, 17): '(2)', (214, 19, 214, 20): '(2)', (214, 22, 214, 23): '(1)'}, {}), '(2, 2, 1)', True, 'import matplotlib.pyplot as plt\n'), ((216, 4, 216, 29), 'matplotlib.pyplot.title', 'plt.title', ({(216, 14, 216, 28): '"""Encoder Loss"""'}, {}), "('Encoder Loss')", True, 'import matplotlib.pyplot as plt\n'), ((217, 4, 217, 24), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(217, 16, 217, 17): '(2)', (217, 19, 217, 20): '(2)', (217, 22, 217, 23): '(2)'}, {}), '(2, 2, 2)', True, 'import matplotlib.pyplot as plt\n'), ((219, 4, 219, 29), 'matplotlib.pyplot.title', 'plt.title', ({(219, 14, 219, 28): '"""Decoder Loss"""'}, {}), "('Decoder Loss')", True, 'import matplotlib.pyplot as plt\n'), ((221, 4, 221, 24), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(221, 16, 221, 17): '(2)', (221, 19, 221, 20): '(3)', (221, 22, 221, 23): '(4)'}, {}), '(2, 3, 4)', True, 'import matplotlib.pyplot as plt\n'), ((222, 4, 222, 67), 'matplotlib.pyplot.plot', 'plt.plot', ({(222, 13, 222, 36): 'generated_samples[:, (0)]', (222, 38, 222, 61): 'generated_samples[:, (1)]', (222, 63, 222, 66): '"""."""'}, {}), "(generated_samples[:, (0)], generated_samples[:, (1)], '.')", True, 'import matplotlib.pyplot as plt\n'), ((223, 4, 223, 34), 'matplotlib.pyplot.title', 'plt.title', ({(223, 14, 223, 33): '"""Generated Samples"""'}, {}), "('Generated Samples')", True, 'import matplotlib.pyplot as plt\n'), ((224, 4, 224, 24), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(224, 16, 224, 17): '(2)', (224, 19, 224, 20): '(3)', (224, 22, 224, 23): '(5)'}, {}), '(2, 3, 5)', True, 'import matplotlib.pyplot as plt\n'), ((225, 4, 225, 75), 'matplotlib.pyplot.plot', 'plt.plot', ({(225, 13, 225, 40): 'reconstructed_samples[:, (0)]', (225, 42, 225, 69): 'reconstructed_samples[:, (1)]', (225, 71, 225, 74): '"""."""'}, {}), "(reconstructed_samples[:, (0)], reconstructed_samples[:, (1)], '.')", True, 'import matplotlib.pyplot as plt\n'), ((228, 4, 228, 31), 'matplotlib.pyplot.title', 'plt.title', ({(228, 14, 228, 30): '"""Reconstruction"""'}, {}), "('Reconstruction')", True, 'import matplotlib.pyplot as plt\n'), ((230, 4, 230, 24), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(230, 16, 230, 17): '(2)', (230, 19, 230, 20): '(3)', (230, 22, 230, 23): '(6)'}, {}), '(2, 3, 6)', True, 'import matplotlib.pyplot as plt\n'), ((231, 4, 231, 61), 'matplotlib.pyplot.plot', 'plt.plot', ({(231, 13, 231, 33): 'vis_samples_np[:, (0)]', (231, 35, 231, 55): 'vis_samples_np[:, (1)]', (231, 57, 231, 60): '"""."""'}, {}), "(vis_samples_np[:, (0)], vis_samples_np[:, (1)], '.')", True, 'import matplotlib.pyplot as plt\n'), ((232, 4, 232, 63), 'matplotlib.pyplot.plot', 'plt.plot', ({(232, 13, 232, 34): 'true_xy_mean_np[:, (0)]', (232, 36, 232, 57): 'true_xy_mean_np[:, (1)]', (232, 59, 232, 62): '"""."""'}, {}), "(true_xy_mean_np[:, (0)], true_xy_mean_np[:, (1)], '.')", True, 'import matplotlib.pyplot as plt\n'), ((233, 4, 233, 33), 'matplotlib.pyplot.title', 'plt.title', ({(233, 14, 233, 32): '"""Original Samples"""'}, {}), "('Original Samples')", True, 'import matplotlib.pyplot as plt\n'), ((234, 4, 234, 42), 'matplotlib.pyplot.legend', 'plt.legend', ({(234, 15, 234, 41): "['Original', 'True means']"}, {}), "(['Original', 'True means'])", True, 'import matplotlib.pyplot as plt\n'), ((235, 4, 235, 14), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((306, 18, 306, 47), 'railrl.torch.pytorch_util.np_to_var', 'ptu.np_to_var', ({(306, 32, 306, 46): 'vis_samples_np'}, {}), '(vis_samples_np)', True, 'import railrl.torch.pytorch_util as ptu\n'), ((314, 4, 314, 24), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(314, 16, 314, 17): '(2)', (314, 19, 314, 20): '(3)', (314, 22, 314, 23): '(1)'}, {}), '(2, 3, 1)', True, 'import matplotlib.pyplot as plt\n'), ((316, 4, 316, 30), 'matplotlib.pyplot.title', 'plt.title', ({(316, 14, 316, 29): '"""Training Loss"""'}, {}), "('Training Loss')", True, 'import matplotlib.pyplot as plt\n'), ((317, 4, 317, 24), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(317, 16, 317, 17): '(2)', (317, 19, 317, 20): '(3)', (317, 22, 317, 23): '(2)'}, {}), '(2, 3, 2)', True, 'import matplotlib.pyplot as plt\n'), ((319, 4, 319, 20), 'matplotlib.pyplot.title', 'plt.title', ({(319, 14, 319, 19): '"""KLs"""'}, {}), "('KLs')", True, 'import matplotlib.pyplot as plt\n'), ((320, 4, 320, 24), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(320, 16, 320, 17): '(2)', (320, 19, 320, 20): '(3)', (320, 22, 320, 23): '(3)'}, {}), '(2, 3, 3)', True, 'import matplotlib.pyplot as plt\n'), ((322, 4, 322, 26), 'matplotlib.pyplot.title', 'plt.title', ({(322, 14, 322, 25): '"""Log Probs"""'}, {}), "('Log Probs')", True, 'import matplotlib.pyplot as plt\n'), ((324, 4, 324, 24), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(324, 16, 324, 17): '(2)', (324, 19, 324, 20): '(3)', (324, 22, 324, 23): '(4)'}, {}), '(2, 3, 4)', True, 'import matplotlib.pyplot as plt\n'), ((325, 4, 325, 67), 'matplotlib.pyplot.plot', 'plt.plot', ({(325, 13, 325, 36): 'generated_samples[:, (0)]', (325, 38, 325, 61): 'generated_samples[:, (1)]', (325, 63, 325, 66): '"""."""'}, {}), "(generated_samples[:, (0)], generated_samples[:, (1)], '.')", True, 'import matplotlib.pyplot as plt\n'), ((326, 4, 326, 34), 'matplotlib.pyplot.title', 'plt.title', ({(326, 14, 326, 33): '"""Generated Samples"""'}, {}), "('Generated Samples')", True, 'import matplotlib.pyplot as plt\n'), ((327, 4, 327, 24), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(327, 16, 327, 17): '(2)', (327, 19, 327, 20): '(3)', (327, 22, 327, 23): '(5)'}, {}), '(2, 3, 5)', True, 'import matplotlib.pyplot as plt\n'), ((328, 4, 328, 75), 'matplotlib.pyplot.plot', 'plt.plot', ({(328, 13, 328, 40): 'reconstructed_samples[:, (0)]', (328, 42, 328, 69): 'reconstructed_samples[:, (1)]', (328, 71, 328, 74): '"""."""'}, {}), "(reconstructed_samples[:, (0)], reconstructed_samples[:, (1)], '.')", True, 'import matplotlib.pyplot as plt\n'), ((330, 4, 330, 63), 'matplotlib.pyplot.plot', 'plt.plot', ({(330, 13, 330, 34): 'estimated_means[:, (0)]', (330, 36, 330, 57): 'estimated_means[:, (1)]', (330, 59, 330, 62): '"""."""'}, {}), "(estimated_means[:, (0)], estimated_means[:, (1)], '.')", True, 'import matplotlib.pyplot as plt\n'), ((331, 4, 331, 31), 'matplotlib.pyplot.title', 'plt.title', ({(331, 14, 331, 30): '"""Reconstruction"""'}, {}), "('Reconstruction')", True, 'import matplotlib.pyplot as plt\n'), ((332, 4, 332, 24), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(332, 16, 332, 17): '(2)', (332, 19, 332, 20): '(3)', (332, 22, 332, 23): '(6)'}, {}), '(2, 3, 6)', True, 'import matplotlib.pyplot as plt\n'), ((333, 4, 333, 61), 'matplotlib.pyplot.plot', 'plt.plot', ({(333, 13, 333, 33): 'vis_samples_np[:, (0)]', (333, 35, 333, 55): 'vis_samples_np[:, (1)]', (333, 57, 333, 60): '"""."""'}, {}), "(vis_samples_np[:, (0)], vis_samples_np[:, (1)], '.')", True, 'import matplotlib.pyplot as plt\n'), ((334, 4, 334, 63), 'matplotlib.pyplot.plot', 'plt.plot', ({(334, 13, 334, 34): 'true_xy_mean_np[:, (0)]', (334, 36, 334, 57): 'true_xy_mean_np[:, (1)]', (334, 59, 334, 62): '"""."""'}, {}), "(true_xy_mean_np[:, (0)], true_xy_mean_np[:, (1)], '.')", True, 'import matplotlib.pyplot as plt\n'), ((335, 4, 335, 33), 'matplotlib.pyplot.title', 'plt.title', ({(335, 14, 335, 32): '"""Original Samples"""'}, {}), "('Original Samples')", True, 'import matplotlib.pyplot as plt\n'), ((336, 4, 336, 42), 'matplotlib.pyplot.legend', 'plt.legend', ({(336, 15, 336, 41): "['Original', 'True means']"}, {}), "(['Original', 'True means'])", True, 'import matplotlib.pyplot as plt\n'), ((337, 4, 337, 14), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((30, 11, 30, 27), 'numpy.array', 'np.array', ({(30, 20, 30, 26): '[x, y]'}, {}), '([x, y])', True, 'import numpy as np\n'), ((31, 12, 31, 42), 'numpy.random.randn', 'np.random.randn', ({(31, 28, 31, 38): 'batch_size', (31, 40, 31, 41): '(2)'}, {}), '(batch_size, 2)', True, 'import numpy as np\n'), ((38, 11, 38, 27), 'numpy.array', 'np.array', ({(38, 20, 38, 26): '[x, y]'}, {}), '([x, y])', True, 'import numpy as np\n'), ((84, 11, 84, 27), 'numpy.array', 'np.array', ({(84, 20, 84, 26): '[x, y]'}, {}), '([x, y])', True, 'import numpy as np\n'), ((91, 12, 91, 31), 'railrl.torch.pytorch_util.np_to_var', 'ptu.np_to_var', ({(91, 26, 91, 30): 'x_np'}, {}), '(x_np)', True, 'import railrl.torch.pytorch_util as ptu\n'), ((92, 12, 92, 31), 'railrl.torch.pytorch_util.np_to_var', 'ptu.np_to_var', ({(92, 26, 92, 30): 'y_np'}, {}), '(y_np)', True, 'import railrl.torch.pytorch_util as ptu\n'), ((103, 12, 103, 31), 'railrl.torch.pytorch_util.np_to_var', 'ptu.np_to_var', ({(103, 26, 103, 30): 'x_np'}, {}), '(x_np)', True, 'import railrl.torch.pytorch_util as ptu\n'), ((108, 8, 108, 28), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(108, 20, 108, 21): '(2)', (108, 23, 108, 24): '(1)', (108, 26, 108, 27): '(1)'}, {}), '(2, 1, 1)', True, 'import matplotlib.pyplot as plt\n'), ((110, 8, 110, 34), 'matplotlib.pyplot.title', 'plt.title', ({(110, 18, 110, 33): '"""Training Loss"""'}, {}), "('Training Loss')", True, 'import matplotlib.pyplot as plt\n'), ((112, 8, 112, 28), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(112, 20, 112, 21): '(2)', (112, 23, 112, 24): '(1)', (112, 26, 112, 27): '(2)'}, {}), '(2, 1, 2)', True, 'import matplotlib.pyplot as plt\n'), ((113, 8, 113, 45), 'matplotlib.pyplot.plot', 'plt.plot', ({(113, 17, 113, 27): 'x_np[:, (0)]', (113, 29, 113, 39): 'x_np[:, (1)]', (113, 41, 113, 44): '"""."""'}, {}), "(x_np[:, (0)], x_np[:, (1)], '.')", True, 'import matplotlib.pyplot as plt\n'), ((114, 8, 114, 53), 'matplotlib.pyplot.plot', 'plt.plot', ({(114, 17, 114, 31): 'x_hat_np[:, (0)]', (114, 33, 114, 47): 'x_hat_np[:, (1)]', (114, 49, 114, 52): '"""."""'}, {}), "(x_hat_np[:, (0)], x_hat_np[:, (1)], '.')", True, 'import matplotlib.pyplot as plt\n'), ((115, 8, 115, 28), 'matplotlib.pyplot.title', 'plt.title', ({(115, 18, 115, 27): '"""Samples"""'}, {}), "('Samples')", True, 'import matplotlib.pyplot as plt\n'), ((116, 8, 116, 44), 'matplotlib.pyplot.legend', 'plt.legend', ({(116, 19, 116, 43): "['Samples', 'Estimates']"}, {}), "(['Samples', 'Estimates'])", True, 'import matplotlib.pyplot as plt\n'), ((117, 8, 117, 18), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((168, 8, 168, 33), 'torch.nn.Linear', 'nn.Linear', ({(168, 18, 168, 19): '2', (168, 21, 168, 32): 'HIDDEN_SIZE'}, {}), '(2, HIDDEN_SIZE)', True, 'from torch import nn as nn\n'), ((169, 8, 169, 17), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'from torch import nn as nn\n'), ((170, 8, 170, 43), 'torch.nn.Linear', 'nn.Linear', ({(170, 18, 170, 29): 'HIDDEN_SIZE', (170, 31, 170, 42): 'HIDDEN_SIZE'}, {}), '(HIDDEN_SIZE, HIDDEN_SIZE)', True, 'from torch import nn as nn\n'), ((171, 8, 171, 17), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'from torch import nn as nn\n'), ((172, 8, 172, 43), 'torch.nn.Linear', 'nn.Linear', ({(172, 18, 172, 29): 'HIDDEN_SIZE', (172, 31, 172, 42): 'HIDDEN_SIZE'}, {}), '(HIDDEN_SIZE, HIDDEN_SIZE)', True, 'from torch import nn as nn\n'), ((173, 8, 173, 17), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'from torch import nn as nn\n'), ((174, 8, 174, 43), 'torch.nn.Linear', 'nn.Linear', ({(174, 18, 174, 29): 'HIDDEN_SIZE', (174, 31, 174, 42): 'HIDDEN_SIZE'}, {}), '(HIDDEN_SIZE, HIDDEN_SIZE)', True, 'from torch import nn as nn\n'), ((175, 8, 175, 17), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'from torch import nn as nn\n'), ((176, 8, 176, 33), 'torch.nn.Linear', 'nn.Linear', ({(176, 18, 176, 29): 'HIDDEN_SIZE', (176, 31, 176, 32): '2'}, {}), '(HIDDEN_SIZE, 2)', True, 'from torch import nn as nn\n'), ((180, 8, 180, 33), 'torch.nn.Linear', 'nn.Linear', ({(180, 18, 180, 19): '1', (180, 21, 180, 32): 'HIDDEN_SIZE'}, {}), '(1, HIDDEN_SIZE)', True, 'from torch import nn as nn\n'), ((181, 8, 181, 17), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'from torch import nn as nn\n'), ((182, 8, 182, 43), 'torch.nn.Linear', 'nn.Linear', ({(182, 18, 182, 29): 'HIDDEN_SIZE', (182, 31, 182, 42): 'HIDDEN_SIZE'}, {}), '(HIDDEN_SIZE, HIDDEN_SIZE)', True, 'from torch import nn as nn\n'), ((183, 8, 183, 17), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'from torch import nn as nn\n'), ((184, 8, 184, 43), 'torch.nn.Linear', 'nn.Linear', ({(184, 18, 184, 29): 'HIDDEN_SIZE', (184, 31, 184, 42): 'HIDDEN_SIZE'}, {}), '(HIDDEN_SIZE, HIDDEN_SIZE)', True, 'from torch import nn as nn\n'), ((185, 8, 185, 17), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'from torch import nn as nn\n'), ((186, 8, 186, 43), 'torch.nn.Linear', 'nn.Linear', ({(186, 18, 186, 29): 'HIDDEN_SIZE', (186, 31, 186, 42): 'HIDDEN_SIZE'}, {}), '(HIDDEN_SIZE, HIDDEN_SIZE)', True, 'from torch import nn as nn\n'), ((187, 8, 187, 17), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'from torch import nn as nn\n'), ((188, 8, 188, 33), 'torch.nn.Linear', 'nn.Linear', ({(188, 18, 188, 29): 'HIDDEN_SIZE', (188, 31, 188, 32): '4'}, {}), '(HIDDEN_SIZE, 4)', True, 'from torch import nn as nn\n'), ((215, 13, 215, 37), 'numpy.array', 'np.array', ({(215, 22, 215, 36): 'encoder_losses'}, {}), '(encoder_losses)', True, 'import numpy as np\n'), ((218, 13, 218, 37), 'numpy.array', 'np.array', ({(218, 22, 218, 36): 'decoder_losses'}, {}), '(decoder_losses)', True, 'import numpy as np\n'), ((240, 8, 240, 33), 'torch.nn.Linear', 'nn.Linear', ({(240, 18, 240, 19): '2', (240, 21, 240, 32): 'HIDDEN_SIZE'}, {}), '(2, HIDDEN_SIZE)', True, 'from torch import nn as nn\n'), ((241, 8, 241, 17), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'from torch import nn as nn\n'), ((242, 8, 242, 43), 'torch.nn.Linear', 'nn.Linear', ({(242, 18, 242, 29): 'HIDDEN_SIZE', (242, 31, 242, 42): 'HIDDEN_SIZE'}, {}), '(HIDDEN_SIZE, HIDDEN_SIZE)', True, 'from torch import nn as nn\n'), ((243, 8, 243, 17), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'from torch import nn as nn\n'), ((244, 8, 244, 43), 'torch.nn.Linear', 'nn.Linear', ({(244, 18, 244, 29): 'HIDDEN_SIZE', (244, 31, 244, 42): 'HIDDEN_SIZE'}, {}), '(HIDDEN_SIZE, HIDDEN_SIZE)', True, 'from torch import nn as nn\n'), ((245, 8, 245, 17), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'from torch import nn as nn\n'), ((246, 8, 246, 43), 'torch.nn.Linear', 'nn.Linear', ({(246, 18, 246, 29): 'HIDDEN_SIZE', (246, 31, 246, 42): 'HIDDEN_SIZE'}, {}), '(HIDDEN_SIZE, HIDDEN_SIZE)', True, 'from torch import nn as nn\n'), ((247, 8, 247, 17), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'from torch import nn as nn\n'), ((248, 8, 248, 33), 'torch.nn.Linear', 'nn.Linear', ({(248, 18, 248, 29): 'HIDDEN_SIZE', (248, 31, 248, 32): '2'}, {}), '(HIDDEN_SIZE, 2)', True, 'from torch import nn as nn\n'), ((255, 8, 255, 33), 'torch.nn.Linear', 'nn.Linear', ({(255, 18, 255, 19): '1', (255, 21, 255, 32): 'HIDDEN_SIZE'}, {}), '(1, HIDDEN_SIZE)', True, 'from torch import nn as nn\n'), ((256, 8, 256, 17), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'from torch import nn as nn\n'), ((257, 8, 257, 43), 'torch.nn.Linear', 'nn.Linear', ({(257, 18, 257, 29): 'HIDDEN_SIZE', (257, 31, 257, 42): 'HIDDEN_SIZE'}, {}), '(HIDDEN_SIZE, HIDDEN_SIZE)', True, 'from torch import nn as nn\n'), ((258, 8, 258, 17), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'from torch import nn as nn\n'), ((259, 8, 259, 43), 'torch.nn.Linear', 'nn.Linear', ({(259, 18, 259, 29): 'HIDDEN_SIZE', (259, 31, 259, 42): 'HIDDEN_SIZE'}, {}), '(HIDDEN_SIZE, HIDDEN_SIZE)', True, 'from torch import nn as nn\n'), ((260, 8, 260, 17), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'from torch import nn as nn\n'), ((261, 8, 261, 43), 'torch.nn.Linear', 'nn.Linear', ({(261, 18, 261, 29): 'HIDDEN_SIZE', (261, 31, 261, 42): 'HIDDEN_SIZE'}, {}), '(HIDDEN_SIZE, HIDDEN_SIZE)', True, 'from torch import nn as nn\n'), ((262, 8, 262, 17), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'from torch import nn as nn\n'), ((263, 8, 263, 33), 'torch.nn.Linear', 'nn.Linear', ({(263, 18, 263, 29): 'HIDDEN_SIZE', (263, 31, 263, 32): '4'}, {}), '(HIDDEN_SIZE, 4)', True, 'from torch import nn as nn\n'), ((273, 16, 273, 36), 'railrl.torch.pytorch_util.np_to_var', 'ptu.np_to_var', ({(273, 30, 273, 35): 'batch'}, {}), '(batch)', True, 'import railrl.torch.pytorch_util as ptu\n'), ((315, 13, 315, 29), 'numpy.array', 'np.array', ({(315, 22, 315, 28): 'losses'}, {}), '(losses)', True, 'import numpy as np\n'), ((318, 13, 318, 26), 'numpy.array', 'np.array', ({(318, 22, 318, 25): 'kls'}, {}), '(kls)', True, 'import numpy as np\n'), ((321, 13, 321, 32), 'numpy.array', 'np.array', ({(321, 22, 321, 31): 'log_probs'}, {}), '(log_probs)', True, 'import numpy as np\n'), ((28, 12, 28, 34), 'numpy.cos', 'np.cos', ({(28, 19, 28, 33): '(t * SWIRL_RATE)'}, {}), '(t * SWIRL_RATE)', True, 'import numpy as np\n'), ((29, 12, 29, 34), 'numpy.sin', 'np.sin', ({(29, 19, 29, 33): '(t * SWIRL_RATE)'}, {}), '(t * SWIRL_RATE)', True, 'import numpy as np\n'), ((36, 12, 36, 34), 'numpy.cos', 'np.cos', ({(36, 19, 36, 33): '(t * SWIRL_RATE)'}, {}), '(t * SWIRL_RATE)', True, 'import numpy as np\n'), ((37, 12, 37, 34), 'numpy.sin', 'np.sin', ({(37, 19, 37, 33): '(t * SWIRL_RATE)'}, {}), '(t * SWIRL_RATE)', True, 'import numpy as np\n'), ((82, 12, 82, 34), 'numpy.cos', 'np.cos', ({(82, 19, 82, 33): '(t * SWIRL_RATE)'}, {}), '(t * SWIRL_RATE)', True, 'import numpy as np\n'), ((83, 12, 83, 34), 'numpy.sin', 'np.sin', ({(83, 19, 83, 33): '(t * SWIRL_RATE)'}, {}), '(t * SWIRL_RATE)', True, 'import numpy as np\n'), ((109, 17, 109, 33), 'numpy.array', 'np.array', ({(109, 26, 109, 32): 'losses'}, {}), '(losses)', True, 'import numpy as np\n'), ((140, 20, 140, 47), 'railrl.torch.pytorch_util.np_to_var', 'ptu.np_to_var', ({(140, 34, 140, 46): 'true_latents'}, {}), '(true_latents)', True, 'import railrl.torch.pytorch_util as ptu\n'), ((211, 21, 211, 48), 'torch.randn', 'torch.randn', ({(211, 33, 211, 47): '*latents.shape'}, {}), '(*latents.shape)', False, 'import torch\n'), ((292, 24, 292, 51), 'railrl.torch.pytorch_util.np_to_var', 'ptu.np_to_var', ({(292, 38, 292, 50): 'true_latents'}, {}), '(true_latents)', True, 'import railrl.torch.pytorch_util as ptu\n'), ((311, 21, 311, 48), 'torch.randn', 'torch.randn', ({(311, 33, 311, 47): '*latents.shape'}, {}), '(*latents.shape)', False, 'import torch\n')]