def __init__(self,
directory,
file_limit=None,
silent=False,
max_episode_length=1000):
# gym env definition
super(ArchiveEnv, self).__init__()
self.action_space = spaces.Box(low=-1,
high=1,
shape=(3, ),
dtype=np.float32)
self.observation_space = spaces.Tuple((
spaces.Box(low=-np.inf,
high=np.inf,
shape=(1080, ),
dtype=np.float32),
spaces.Box(low=-np.inf, high=np.inf, shape=(5, ),
dtype=np.float32),
))
self.data = WorldModelDataset._load_data(None,
directory,
file_limit=file_limit)
if len(self.data["scans"]) == 0:
raise ValueError
self.current_iteration = None
self.data["dones"][-1] = 1
self.viewer = None
def vae1dlstm_worldmodel_error(vae1dlstm, test_dataset_folder, device):
sequence_size = 32
batch_size = 128
config = vae1dlstm.module.config if hasattr(vae1dlstm,
"module") else vae1dlstm.config
# load dataset
seq_loader = WorldModelDataset(test_dataset_folder,
sequence_size,
lidar_mode="scans",
channel_first=True,
as_torch_tensors=True,
file_limit=64)
batch_loader = DataLoader(seq_loader, shuffle=False, batch_size=batch_size)
# iterate over batches
batch_loader = tqdm(batch_loader, total=len(batch_loader))
n_batches = 0
sum_state_error = 0
sum_lidar_error = 0
for x, a, y, x_rs, y_rs, dones in batch_loader:
# initialize sequence rnn_state
this_batch_size = x.shape[0] # last batch can be smaller
h0 = torch.randn(config.lstm.n_layer, this_batch_size,
config.lstm.h_size)
c0 = torch.randn(config.lstm.n_layer, this_batch_size,
config.lstm.h_size)
# place data on the correct device
x = x.to(device)
x_rs = x_rs.to(device)
a = a.to(device)
y = y.to(device)
y_rs = y_rs.to(device)
dones = dones.to(device)
h0 = h0.to(device)
c0 = c0.to(device)
y_pred_rec, y_rs_pred, _ = vae1dlstm(x, x_rs, a, dones, [h0, c0])
y_pred_rec = y_pred_rec.detach().cpu().numpy()
y_rs_pred = y_rs_pred.detach().cpu().numpy()
sum_lidar_error += mse(
y_pred_rec,
y.cpu().numpy()) # because binary cross entropy is inf for 0
sum_state_error += mse(y_rs_pred,
y_rs.cpu().numpy()) # mean square error loss
n_batches += 1
lidar_error = sum_lidar_error / n_batches
state_error = sum_state_error / n_batches
return lidar_error, state_error
def gpt1d_worldmodel_error(gpt, test_dataset_folder, device):
sequence_size = gpt.module.block_size
batch_size = 128
# load dataset
seq_loader = WorldModelDataset(test_dataset_folder,
sequence_size,
lidar_mode="scans",
channel_first=True,
as_torch_tensors=True,
file_limit=64)
batch_loader = DataLoader(seq_loader, shuffle=False, batch_size=batch_size)
# iterate over batches
batch_loader = tqdm(batch_loader, total=len(batch_loader))
n_batches = 0
sum_state_error = 0
sum_lidar_error = 0
for x, a, y, x_rs, y_rs, dones in batch_loader:
# place data on the correct device
x = x.to(device)
x_rs = x_rs.to(device)
a = a.to(device)
y = y.to(device)
y_rs = y_rs.to(device)
dones = dones.to(device)
y_pred_rec, y_rs_pred, _ = gpt(x, x_rs, a, dones)
y_pred_rec = y_pred_rec.detach().cpu().numpy()
y_rs_pred = y_rs_pred.detach().cpu().numpy()
sum_lidar_error += mse(
y_pred_rec,
y.cpu().numpy()) # because binary cross entropy is inf for 0
sum_state_error += mse(y_rs_pred,
y_rs.cpu().numpy()) # mean square error loss
n_batches += 1
lidar_error = sum_lidar_error / n_batches
state_error = sum_state_error / n_batches
return lidar_error, state_error
def vae1d_rnn_worldmodel_error(rnn, test_dataset_folder, vae):
sequence_size = rnn.hps.max_seq_len - 1
batch_size = rnn.hps.batch_size
# load dataset
seq_loader = WorldModelDataset(test_dataset_folder,
sequence_size,
lidar_mode="scans",
channel_first=False,
as_torch_tensors=False,
file_limit=64)
batch_loader = DataLoader(seq_loader, shuffle=False, batch_size=batch_size)
# iterate over batches
pbar = tqdm(range(len(batch_loader) * batch_size),
total=len(batch_loader) * batch_size).__iter__()
n_batches = 0
sum_state_error = 0
sum_lidar_error = 0
for x, a, y, x_rs, y_rs, dones in batch_loader:
# for legacy reasons rnn likes x_rs to be normalized
x_rs = x_rs / MAX_GOAL_DIST
# convert img to z
batch_state = rnn.sess.run(rnn.initial_state)
# doing the whole batch at once is too mem intensive
z = np.zeros((batch_size, sequence_size, vae.z_size)) # (b, t, 32)
for i, single_x in enumerate(x):
pbar.__next__()
z[i] = vae.encode(single_x)
batch_z_rs = np.concatenate([z, x_rs], axis=-1) # (b, t, 34)
# add 1 to sequence
batch_z_rs = np.concatenate(
[batch_z_rs,
np.zeros((batch_size, 1, rnn.hps.seq_width))], axis=1)
a = np.concatenate([a, np.zeros((batch_size, 1, a.shape[-1]))], axis=1)
dones = np.concatenate([dones, np.zeros((batch_size, 1))], axis=1)
# predict
feed = {
rnn.batch_z_rs: batch_z_rs,
rnn.batch_action: a,
rnn.batch_restart: dones,
rnn.initial_state: batch_state,
}
(logmix, out_mean) = rnn.sess.run([rnn.out_logmix, rnn.out_mean], feed)
# sample deterministic from mean of MDN
z_rs_pred = out_mean[(np.arange(len(logmix)), np.argmax(logmix,
axis=-1))]
z_rs_pred = z_rs_pred.reshape(batch_size, sequence_size,
rnn.hps.seq_width)
z_pred = z_rs_pred[:, :, :_Z]
y_rs_pred = z_rs_pred[:, :, _Z:]
y_pred_rec = np.zeros((batch_size, sequence_size, 1080, 1))
for i, single_z_pred in enumerate(z_pred):
y_pred_rec[i] = vae.decode(single_z_pred)
# cancel the effect of x_rs normalization
y_rs_pred = y_rs_pred * MAX_GOAL_DIST
sum_lidar_error += mse(
y_pred_rec,
y.numpy()) # because binary cross entropy is inf for 0 predictions
sum_state_error += mse(y_rs_pred,
y_rs.numpy()) # mean square error loss
n_batches += 1
for _ in pbar:
pass
lidar_error = sum_lidar_error / n_batches
state_error = sum_state_error / n_batches
return lidar_error, state_error
# make deterministic
set_seed(42)
# set up logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
datefmt="%m/%d/%Y %H:%M:%S",
level=logging.INFO,
)
logger = logging.getLogger(__name__)
mconf = VAELSTMConfig(_Z, _H)
train_dataset = WorldModelDataset(
dataset_dir,
_S,
pre_convert_obs=True,
regen=data_regen,
)
# training params
# optimization parameters
max_steps = args.n
if max_steps is None:
max_steps = 222222
max_epochs = max_steps # don't stop based on epoch
batch_size = 128
learning_rate = 6e-4
betas = (0.9, 0.95)
grad_norm_clip = 1.0
lr_decay = True # learning rate decay params: linear warmup followed by cosine decay to 10% of original
weight_decay = 0.1 # only applied on matmul weights
# make deterministic
set_seed(42)
# set up logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
datefmt="%m/%d/%Y %H:%M:%S",
level=logging.INFO,
)
logger = logging.getLogger(__name__)
mconf = GPTConfig(_S, _H)
train_dataset = WorldModelDataset(
dataset_dir, _S,
lidar_mode="scans",
pre_convert_obs=True,
regen=data_regen,
)
# training params
# optimization parameters
max_steps = args.n
if max_steps is None:
max_steps = 222222
max_epochs = max_steps # don't stop based on epoch
batch_size = 128
learning_rate = 6e-4
betas = (0.9, 0.95)
grad_norm_clip = 1.0
lr_decay = True # learning rate decay params: linear warmup followed by cosine decay to 10% of original
weight_decay = 0.1 # only applied on matmul weights