def train(args):
net = HypothesisNet(args)
if not args.no_log:
log = log_this(net.args, 'logs/sim', args.name, checkpoints=False)
simulator = net.simulator
if args.no_reservoir:
layer = net.W_ro
else:
layer = net.reservoir
batch_size = 10
criterion = nn.MSELoss()
train_params = simulator.parameters()
optimizer = optim.Adam(train_params, lr=args.lr)
for i in range(args.iters):
if not args.no_reservoir:
layer.reset('random')
optimizer.zero_grad()
prop = torch.Tensor(
np.random.normal(0, 5, size=(batch_size, net.args.D)))
state = torch.Tensor(
np.random.normal(0, 10, size=(batch_size, net.args.L)))
sim_out = simulator(state, prop)
# run reservoir 10 steps, so predict 10 steps in future
outs = []
for j in range(args.forward_steps):
outs.append(layer(prop))
actions = sum(outs)
# get state output
layer_out = actions + state
# validation makes sure performance is poor if we use someone else's output
layer_out_val = actions.roll(1, 0) + state
# calculate euclidean loss
loss = criterion(layer_out, sim_out)
loss_val = criterion(layer_out_val, sim_out)
loss.backward()
optimizer.step()
if i % 50 == 0 and i != 0:
print(f'iteration: {i} | loss {loss} | loss_val {loss_val}')
if not args.no_log:
save_model_path = os.path.join(log.run_dir, f'model_{log.run_id}.pth')
save_sim_path = os.path.join(log.run_dir, f'sim_{log.run_id}.pth')
torch.save(net.state_dict(), save_model_path)
torch.save(simulator.state_dict(), save_sim_path)
print(f'saved model to {save_model_path}, sim to {save_sim_path}')
def train(args):
net = HypothesisNet(args)
if not args.no_log:
log = log_this(net.args, 'logs/hyp', args.name, checkpoints=False)
simulator = net.simulator
hypothesizer = net.hypothesizer
batch_size = 10
criterion = nn.MSELoss()
train_params = hypothesizer.parameters()
optimizer = optim.Adam(train_params, lr=1e-3)
for i in range(args.iters):
optimizer.zero_grad()
state = torch.Tensor(
np.random.normal(0, 10, size=(batch_size, net.args.L)))
task = torch.Tensor(
np.random.normal(0, 10, size=(batch_size, net.args.L)))
prop = hypothesizer(state, task)
sim_out = simulator(state, prop)
# run reservoir 10 steps, so predict 10 steps in future
outs = []
for j in range(10):
outs.append(layer(prop))
actions = sum(outs)
# get state output
layer_out = actions + state
# validation makes sure performance is poor if we use someone else's output
layer_out_val = actions.roll(1, 0) + state
# calculate euclidean loss
diff = torch.norm(layer_out - sim_out, dim=1)
loss = criterion(diff, torch.zeros_like(diff))
diff_val = torch.norm(layer_out_val - sim_out, dim=1)
loss_val = criterion(diff_val, torch.zeros_like(diff_val))
loss.backward()
optimizer.step()
if i % 50 == 0 and i != 0:
print(f'iteration: {i} | loss {loss} | loss_val {loss_val}')
if not args.no_log:
save_model_path = os.path.join(log.run_dir, f'model_{log.run_id}.pth')
torch.save(net.state_dict(), save_model_path)
print(f'saved model to {save_model_path}')
def load_model_path(path, config):
if type(config) is dict:
config = Bunch(**config)
config.model_path = path
if config.net == 'basic':
net = BasicNetwork(config)
elif config.net == 'state':
net = StateNet(config)
elif config.net == 'hypothesis':
net = HypothesisNet(config)
else:
raise NotImplementedError
net.eval()
return net
def test(args):
net = HypothesisNet(args)
simulator = net.simulator
reservoir = net.reservoir
W_ro = net.W_ro
batch_size = 50
criterion = nn.MSELoss()
train_params = simulator.parameters()
optimizer = optim.Adam(train_params, lr=1e-3)
reservoir.reset()
optimizer.zero_grad()
prop = torch.Tensor(np.random.normal(size=(batch_size, net.args.D)))
state = torch.Tensor(np.random.normal(size=(batch_size, net.args.L)))
sim_out = simulator(state, prop)
print(reservoir.J.weight)
# run reservoir 10 steps, so predict 10 steps in future
outs = []
for j in range(10):
outs.append(W_ro(reservoir(prop)))
actions = sum(outs)
# get state output
layer_out = actions + state
sim_out_val = sim_out.roll(1, 0) + state
# calculate euclidean loss
diff = torch.norm(layer_out - sim_out, dim=1)
loss = criterion(diff, torch.zeros_like(diff))
diff2 = torch.norm(layer_out - sim_out_val, dim=1)
loss2 = criterion(diff2, torch.zeros_like(diff))
print(f'loss {loss}, loss2 {loss2}')
def test(args):
net = HypothesisNet(args)
simulator = net.simulator
reservoir = net.reservoir
W_ro = net.W_ro
net.eval()
batch_size = 50
criterion = nn.MSELoss()
train_params = simulator.parameters()
reservoir.reset('random')
prop = torch.Tensor(np.random.normal(size=(batch_size, net.args.D)))
state = torch.Tensor(np.random.normal(size=(batch_size, net.args.L)))
sim_out = simulator(state, prop)
print(reservoir.J.weight)
# run reservoir 10 steps, so predict 10 steps in future
outs = []
for j in range(args.forward_steps):
outs.append(W_ro(reservoir(prop)))
actions = sum(outs)
# get state output
layer_out = actions + state
sim_out_val = sim_out.roll(1, 0) + state
# calculate euclidean loss
loss = criterion(layer_out, sim_out)
loss2 = criterion(layer_out, sim_out_val)
print(f'loss {loss}, loss2 {loss2}')
def main(args):
batch_size = 12
n_steps = 10
# args.res_init_std = 1.5
# args.D = 1
net = HypothesisNet(args)
reservoir = net.reservoir
# reservoir.reset(np.random.normal(0, 1, (1, net.args.N)))
reservoir.reset(np.random.normal(0, 1, (batch_size, net.args.N)))
# reservoir.reset('zero')
# reservoir.W_ro.bias.data = reservoir.W_ro.bias.data
# torch.manual_seed(0)
# np.random.seed(0)
# prop = torch.Tensor(np.tile(np.random.normal(0, 10, size=(1, net.args.D)), (batch_size, 1))) # same input to all 12
prop = torch.Tensor(np.tile(np.random.normal(0, 1, size=(batch_size, net.args.D)), (1, 1))) * 0 # distinct input to all 12
prop2 = torch.Tensor(np.tile(np.random.normal(0, 10, size=(1, net.args.D)), (batch_size, 1))) * 0
init_state = torch.Tensor(np.zeros((batch_size, net.args.L)))
states = [init_state.numpy()]
with torch.no_grad():
for i in range(n_steps):
action = reservoir(prop, extras=False)
next_state = states[-1] + action.numpy()
states.append(next_state)
for i in range(n_steps):
action = reservoir(prop2, extras=False)
next_state = states[-1] + action.numpy()
states.append(next_state)
states = np.transpose(np.array(states), (1, 2, 0))
fig, ax = plt.subplots(3,4,sharex=True, sharey=True, figsize=(12,7))
for i, ax in enumerate(fig.axes):
ax.axvline(x=0, color='dimgray', alpha = 1)
ax.axhline(y=0, color='dimgray', alpha = 1)
ax.grid(True, which='major', lw=1, color='lightgray', alpha=0.4)
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['left'].set_visible(False)
ax.spines['bottom'].set_visible(False)
ax.scatter(states[i][0], states[i][1], color='coral', alpha=0.5, lw=1)
# ax.plot(states[i][, outs[i], color='cornflowerblue', alpha=1, lw=1.5, label='out')
lim = 10
ax.set_xlim([-lim,lim])
ax.set_ylim([-lim,lim])
ax.tick_params(axis='both', color='white')
plt.show()
def __init__(self, args):
super().__init__()
self.args = args
if self.args.net == 'basic':
self.net = BasicNetwork(self.args)
elif self.args.net == 'state':
self.net = StateNet(self.args)
elif self.args.net == 'hypothesis':
self.net = HypothesisNet(self.args)
# picks which parameters to train and which not to train
self.n_params = {}
self.train_params = []
self.not_train_params = []
logging.info('Training the following parameters:')
for k, v in self.net.named_parameters():
# k is name, v is weight
found = False
# filtering just for the parts that will be trained
for part in self.args.train_parts:
if part in k:
logging.info(f' {k}')
self.n_params[k] = (v.shape, v.numel())
self.train_params.append(v)
found = True
break
if not found:
self.not_train_params.append(k)
logging.info('Not training:')
for k in self.not_train_params:
logging.info(f' {k}')
self.criterion = get_criterion(self.args)
self.optimizer = get_optimizer(self.args, self.train_params)
self.dset = load_rb(self.args.dataset)
self.potential = get_potential(self.args)
# if using separate training and test sets, separate them out
if not self.args.same_test:
np.random.shuffle(self.dset)
cutoff = round(.9 * len(self.dset))
self.train_set = self.dset[:cutoff]
self.test_set = self.dset[cutoff:]
logging.info(
f'Using separate training ({cutoff}) and test ({len(self.dset) - cutoff}) sets.'
)
else:
self.train_set = self.dset
self.test_set = self.dset
self.log_interval = self.args.log_interval
if not self.args.no_log:
self.log = self.args.log
self.run_id = self.args.log.run_id
self.vis_samples = []
self.csv_path = open(
os.path.join(self.log.run_dir, f'losses_{self.run_id}.csv'),
'a')
self.writer = csv.writer(self.csv_path,
delimiter=',',
quotechar='|',
quoting=csv.QUOTE_MINIMAL)
self.writer.writerow(['ix', 'avg_loss'])
self.plot_checkpoint_path = os.path.join(
self.log.run_dir, f'checkpoints_{self.run_id}.pkl')
self.save_model_path = os.path.join(self.log.run_dir,
f'model_{self.run_id}.pth')