class DatasetRealSmolBulletV2NoSim(Dataset):
    def __init__(self,
                 path="~/data/sim2real/data-realigned-v3-{}-bullet.npz",
                 train=True):
        super().__init__()

        ds_name = "train"
        if not train:
            ds_name = "test"
        path_ext = path.format(ds_name)

        print("using dataset:", path_ext)

        self.ds = DatasetProduction()
        self.ds.load(path_ext)

    def __len__(self):
        return len(self.ds.current_real)

    def format_data(self, idx):
        return (np.hstack((self.ds.current_real[idx], self.ds.action[idx])),
                self.ds.next_real[idx] - self.ds.current_real[idx])

    def __getitem__(self, idx):
        x, y = self.format_data(idx)
        return {"x": x, "y": y}
    def __init__(self,
                 path="~/data/sim2real/data-realigned-v3-{}-bullet.npz",
                 train=True):
        super().__init__()

        ds_name = "train"
        if not train:
            ds_name = "test"
        path_ext = path.format(ds_name)

        print("using dataset:", path_ext)

        self.ds = DatasetProduction()
        self.ds.load(path_ext)
import time
import numpy as np
import os
import matplotlib.pyplot as plt
import gym
import gym_ergojr
import torch
from gym_ergojr.sim.single_robot import SingleRobot
from hyperdash import Experiment
from s2rr.movements.dataset import DatasetProduction
from torch.autograd import Variable
from simple_joints_lstm.lstm_net_real_v3 import LstmNetRealv3
import torch.nn.functional as F

ds = DatasetProduction()
ds.load("~/data/sim2real/data-realigned-v3-{}-bullet.npz".format("train"))

net = LstmNetRealv3(nodes=128, layers=3)
if torch.cuda.is_available():
    net = net.cuda()

HIDDEN_NODES = 128
LSTM_LAYERS = 3
EXPERIMENT = 1
EPOCHS = 5
MODEL_PATH = "./trained_models/lstm_real_vX5_exp{}_l{}_n{}.pt".format(
    EXPERIMENT, LSTM_LAYERS, HIDDEN_NODES)


def double_unsqueeze(data):
    return torch.unsqueeze(torch.unsqueeze(data, dim=0), dim=0)