def plot_sine_figures(project_dir, time_tag, model_name):
# make the save folder
save_folder = project_dir / f'figures/{model_name}'
try:
os.makedirs(save_folder)
except FileExistsError:
pass
# load previously trained model
device = torch.device('cpu')
conf = OmegaConf.load(project_dir / f'runs/sine/{model_name}/conf.yaml')
h_sizes = OmegaConf.to_container(conf.hidden_sizes)
model = MLPModel(dim_y=1, dim_r=conf.dim_r, dim_z_prime=conf.dim_z_prime, dim_l=conf.dim_l,
hidden_sizes_encoder=h_sizes,
hidden_sizes_ode_net=h_sizes,
hidden_sizes_decoder=h_sizes, t0=-0.1, device=device)
model_weights = torch.load(project_dir / f'runs/sine/{model_name}/seed_0/model.pth',
map_location=device)
model.load_state_dict(model_weights)
# create a sine dataset
dataset = SineData()
dataloader = DataLoader(dataset, batch_size=conf.batch_size, drop_last=True)
t, y = next(iter(dataloader))
t_context, y_context, t_extra, y_extra, _, _ = get_split(t, y, test_context_size=conf.test_context_size)
with torch.no_grad():
plot_sine_img(model, 0, t, y, t_context, y_context, t_extra, save_folder, time_tag)
plot_sine_img(model, 1, t, y, t_context, y_context, t_extra, save_folder, time_tag)
plot_sine_img(model, 2, t, y, t_context, y_context, t_extra, save_folder, time_tag)
plot_sine_img(model, 3, t, y, t_context, y_context, t_extra, save_folder, time_tag)
plot_sine_img(model, 4, t, y, t_context, y_context, t_extra, save_folder, time_tag)
# make the save folder
save_folder = '../figures/'
try:
os.makedirs('./' + save_folder)
except FileExistsError:
pass
device = torch.device('cpu')
conf = OmegaConf.load("../runs/" + dataname + '/' + model_name + '/conf.yaml')
h_sizes = OmegaConf.to_container(conf.hidden_sizes)
model = MLPModel(dim_y=1,
dim_r=conf.dim_r,
dim_z_prime=conf.dim_z_prime,
dim_l=conf.dim_l,
hidden_sizes_encoder=h_sizes,
hidden_sizes_ode_net=h_sizes,
hidden_sizes_decoder=h_sizes,
t0=-3.2,
device=device)
model.load_state_dict(
torch.load('../runs' + '/' + dataname + '/' + model_name +
'/seed_0/model_ep10.pth'))
#import data
from datasets import NoisySineData
from math import pi
dataset = NoisySineData(sigma,
shift_range=(-0.00001, .00001),
freq_range=(1.9999999, 2.0),
num_samples=10)
self, data_batch: DataBatch) -> Tuple[Tensor, Tensor, Tensor]:
"""
Computes action logprobs, observation values and policy entropy for each of the (obs, action, hidden_state)
transitions. Preserves all the necessary gradients.
Args:
data_batch: data collected from a Collector for this agent
Returns:
action_logprobs: tensor of action logprobs (batch_size, )
values: tensor of observation values (batch_size, )
entropies: tensor of entropy values (batch_size, )
"""
obs_batch = data_batch['observations']
action_batch = data_batch['actions']
action_distribution, values = self.model(obs_batch)
action_logprobs = action_distribution.log_prob(action_batch).sum(1)
values = values.view(-1)
entropies = action_distribution.entropy().sum(1)
return action_logprobs, values, entropies
if __name__ == '__main__':
model = MLPModel({})
agent = Agent(model)
action = agent.compute_single_action(
np.random.randn(33).astype(np.float32))
def train(conf, project_dir: Path, run_dir: Path) -> torch.nn.Module:
writer = SummaryWriter(str(run_dir))
save_hparams(conf, writer)
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
if conf.dataset == 'sine':
# dataset of time-series
dataset_train = SineData()
dataloader_train = DataLoader(dataset_train,
batch_size=conf.batch_size,
drop_last=True)
dataset_test = SineData()
dataloader_test = DataLoader(dataset_test,
batch_size=conf.batch_size,
shuffle=False,
drop_last=True)
h_sizes = OmegaConf.to_container(
conf.hidden_sizes) # OmegaConf object to list
model = MLPModel(dim_y=1,
dim_r=conf.dim_r,
dim_z_prime=conf.dim_z_prime,
dim_l=conf.dim_l,
hidden_sizes_encoder=h_sizes,
hidden_sizes_ode_net=h_sizes,
hidden_sizes_decoder=h_sizes,
t0=dataset_train.t0,
device=device)
elif conf.dataset == 'sinefreq':
# dataset of frequency varying sinus time-series
dataset_train = FreqSineData(amplitude_range=(0.5, 1.),
shift_range=(-.5, .5),
freq_range=(1.0, 2.0),
num_samples=5000)
dataloader_train = DataLoader(dataset_train,
batch_size=conf.batch_size,
drop_last=True)
dataset_test = FreqSineData(amplitude_range=(0.5, 1.),
shift_range=(-.5, .5),
freq_range=(1.0, 2.0),
num_samples=1000)
dataloader_test = DataLoader(dataset_test,
batch_size=conf.batch_size,
shuffle=False,
drop_last=True)
h_sizes = OmegaConf.to_container(
conf.hidden_sizes) # OmegaConf object to list
model = MLPModel(dim_y=1,
dim_r=conf.dim_r,
dim_z_prime=conf.dim_z_prime,
dim_l=conf.dim_l,
hidden_sizes_encoder=h_sizes,
hidden_sizes_ode_net=h_sizes,
hidden_sizes_decoder=h_sizes,
t0=dataset_train.t0,
device=device)
elif conf.dataset == 'noisysine':
sigma = conf.sigma
# dataset of noisy sinus time-series
dataset_train = NoisySineData(sigma,
shift_range=(-0.1, .1),
freq_range=(1.9, 2.0),
num_samples=1000)
dataloader_train = DataLoader(dataset_train,
batch_size=conf.batch_size,
drop_last=True)
dataset_test = NoisySineData(sigma,
shift_range=(-0.1, .1),
freq_range=(1.9, 2.0),
num_samples=1000)
dataloader_test = DataLoader(dataset_test,
batch_size=conf.batch_size,
shuffle=False,
drop_last=True)
h_sizes = OmegaConf.to_container(
conf.hidden_sizes) # OmegaConf object to list
model = MLPModel(dim_y=1,
dim_r=conf.dim_r,
dim_z_prime=conf.dim_z_prime,
dim_l=conf.dim_l,
hidden_sizes_encoder=h_sizes,
hidden_sizes_ode_net=h_sizes,
hidden_sizes_decoder=h_sizes,
t0=dataset_train.t0,
device=device)
elif conf.dataset == 'rotnist':
# dataset of Rotating MNIST (in the literature)
dataset_mnist = RotNISTDataset(data_dir=str(project_dir / 'data'))
len_test = 10
dataset_train = dataset_mnist[:len(dataset_mnist) - len_test]
dataset_test = dataset_mnist[len(dataset_mnist) - len_test:]
dataloader_train = DataLoader(dataset_train,
batch_size=conf.batch_size,
drop_last=True)
dataloader_test = DataLoader(dataset_test,
batch_size=conf.batch_size,
shuffle=False,
drop_last=True)
h_sizes = OmegaConf.to_container(conf.hidden_sizes)
model = ConvNetModel(dim_r=conf.dim_r,
dim_z_prime=conf.dim_z_prime,
dim_l=conf.dim_l,
hidden_sizes_ode_net=h_sizes,
t0=dataset_mnist.t0,
device=device)
else:
raise ValueError(f'Dataset {conf.dataset} not recognized')
model = model.to(device)
optimizer = torch.optim.RMSprop(model.parameters(), lr=conf.lr)
context_range = OmegaConf.to_container(conf.context_range)
extra_target_range = OmegaConf.to_container(conf.extra_target_range)
global_train_step = 0
global_test_step = 0
for epoch in tqdm(range(conf.epochs)):
mse_train_list = []
mse_test_list = []
with torch.no_grad():
for step, (t, y) in enumerate(dataloader_test):
t, y = t.to(device), y.to(device)
t_context, y_context, t_extra, y_extra, _, _ = get_split(
t, y, test_context_size=conf.test_context_size)
p_y, _, _ = model(
t_context, y_context, t_extra
) # for testing, we only need predictions at t_extra
output = p_y.loc
mse_test = F.mse_loss(output, y_extra)
# log test results
writer.add_scalar('mse_test', mse_test.item(),
global_test_step)
mse_test_list.append(mse_test.item())
if step == 0 and epoch % 2 == 0:
if conf.dataset in ['sine', 'sinefreq', 'noisysine']:
log_sine_plot(writer, model, t, y, t_context,
y_context, t_extra, epoch)
elif conf.dataset == 'rotnist':
log_rotnist_plot2(writer, model, t, y, epoch, 'test')
global_test_step += 1
for (t, y) in dataloader_train:
t, y = t.to(device), y.to(device)
(t_context, y_context, t_extra, y_extra, t_target,
y_target) = get_split(t,
y,
context_range=context_range,
extra_target_range=extra_target_range)
p_y, q_z_T, q_z_C = model(t_context,
y_context,
t_target,
y_target=y_target)
log_p = p_y.log_prob(y_target).sum(dim=(1, 2)).mean(
dim=0) # mean on batch dim, sum on time dim/y dim
output = p_y.loc
mse_train = F.mse_loss(output, y_target)
# mean on batch dim, sum on z dim (equivalent to kl_div of the multivariate normal)
kl_div = kl_divergence(q_z_C, q_z_T).sum(dim=1).mean(dim=0)
loss = -log_p + kl_div
optimizer.zero_grad()
loss.backward()
optimizer.step()
# log training metrics
writer.add_scalar('kl_div', kl_div.item(), global_train_step)
writer.add_scalar('log_p', log_p.item(), global_train_step)
writer.add_scalar('train_loss', loss.item(), global_train_step)
writer.add_scalar('mse_train', mse_train.item(), global_train_step)
mse_train_list.append(mse_train.item())
global_train_step += 1
# log test/train mse epoch-wise to match the paper's figures
writer.add_scalar('mse_train_epoch', np.mean(mse_train_list), epoch)
writer.add_scalar('mse_test_epoch', np.mean(mse_test_list), epoch)
if epoch % conf.checkpoint_freq == 0 and epoch > 0:
torch.save(model.state_dict(), run_dir / f'model_ep{epoch}.pth')
torch.save(model.state_dict(), run_dir / f'model.pth')
return model
np.random.seed(0)
env = UnityEnvironment(file_name='Reacher.app')
# get the default brain
brain_name = env.brain_names[0]
brain = env.brains[brain_name]
model_config = {
"input_size": 33,
"num_actions": 4,
"activation": "relu",
"hidden_sizes": (512, 512),
"initializer": None, #"kaiming_uniform",
}
model = MLPModel(model_config)
agent = Agent(model)
trainer_config = {
"steps": 2049,
# Tensorboard settings
"tensorboard_name": "test", # str, set explicitly
# PPO
"ppo_config": {
# GD settings
"optimizer": "adam",
"optimizer_kwargs": {
"lr": 3e-4,
return params
if __name__ == '__main__':
params = arg_parser()
model_name = params['model_name']
if params['is_train']:
data = Data()
data.load_data()
ms = MeanShiftClustering(data.Y_train)
ms.fit()
clusters = ms.get_cluster_centroids()
if params['model_name'] == 'mlp':
mlp = MLPModel.MLPModel()
mlp.set_data(data)
mlp.set_clusters(clusters)
mlp.create_model()
mlp.fit()
mlp.save()
elif params['model_name'] == 'lstm':
lstm = RNNModel.RNNModel()
lstm.set_data(data)
lstm.set_clusters(clusters)
lstm.create_model()
lstm.fit()
lstm.save()
elif params['model_name'] == 'brnn':
blstm = BidirectionalRNN.BidirectionalRNN()
blstm.set_data(data)