dl = DataLoader(dset, batch_size=200, num_workers=0, drop_last=True)
dl_eval = DataLoader(dset_eval,
batch_size=200,
num_workers=0,
drop_last=True)
G1 = nx.path_graph(6).to_directed()
G_target = nx.path_graph(6).to_directed()
nx.draw(G1)
plt.show()
node_feat_size = 6
edge_feat_size = 3
graph_feat_size = 10
gn = FFGN(graph_feat_size, node_feat_size, edge_feat_size).cuda()
if opt.model != '':
gn.load_state_dict(torch.load(opt.model))
optimizer = optim.Adam(gn.parameters(), lr=1e-4)
schedular = optim.lr_scheduler.StepLR(optimizer, 5e4, gamma=0.975)
savedir = os.path.join('./logs', 'runs',
datetime.now().strftime('%B%d_%H:%M:%S'))
writer = SummaryWriter(savedir)
step = 0
normalizers = torch.load('normalize.pth')
in_normalizer = normalizers['in_normalizer']
out_normalizer = normalizers['out_normalizer']
std = in_normalizer.get_std()
for epoch in range(300):
for i, data in tqdm(enumerate(dl), total=len(dset) / 200 + 1):
optimizer.zero_grad()
class Acrobot_GN(Acrobot):
'''
'''
STATE_THETA_1, STATE_THETA_2, STATE_V_1, STATE_V_2 = 0, 1, 2, 3
MIN_V_1, MAX_V_1 = -6., 6.
MIN_V_2, MAX_V_2 = -6., 6.
MIN_TORQUE, MAX_TORQUE = -4., 4.
MIN_ANGLE, MAX_ANGLE = -np.pi, np.pi
LENGTH = 20.
m = 1.0
lc = 0.5
lc2 = 0.25
l2 = 1.
I1 = 0.2
I2 = 1.0
g = 9.81
def __init__(self):
super(Acrobot, self).__init__()
self.G1 = nx.path_graph(2).to_directed()
self.node_feat_size = 2
self.edge_feat_size = 3
self.graph_feat_size = 10
self.gn = FFGN(self.graph_feat_size, self.node_feat_size,
self.edge_feat_size).cuda()
self.gn.load_state_dict(torch.load('model0.05.pth'))
normalizers = torch.load('normalized/acrobot0.05.pth')
self.in_normalizer = normalizers['in_normalizer']
self.out_normalizer = normalizers['out_normalizer']
def propagate(self, start_state, control, num_steps, integration_step):
'''
Integrate system dynamics
:param start_state: numpy array with the start state for the integration
:param control: numpy array with constant controls to be applied during integration
:param num_steps: number of steps to integrate
:param integration_step: dt of integration
:return: new state of the system
'''
action = control
delta_state = torch.zeros_like(start_state)
last_state = start_state.clone()
for i in range(num_steps):
init_graph_features(self.G1,
self.graph_feat_size,
self.node_feat_size,
self.edge_feat_size,
cuda=True,
bs=1)
load_graph_features(self.G1,
action,
last_state.clone(),
None,
bs=1,
noise=0,
std=None)
self.gn.eval()
# with torch.no_grad():
G_out = self.gn(self.in_normalizer.normalize(self.G1))
for node in G_out.nodes():
delta_state[0, [node, node +
2]] = G_out.nodes[node]['feat'][0].cpu()
last_state += delta_state
return last_state
def visualize_point(self, state):
x2 = self.LENGTH * np.cos(state[self.STATE_THETA_1] -
np.pi / 2) + self.LENGTH * np.cos(
state[self.STATE_THETA_1] +
state[self.STATE_THETA_2] - np.pi / 2)
y2 = self.LENGTH * np.sin(state[self.STATE_THETA_1] -
np.pi / 2) + self.LENGTH * np.sin(
state[self.STATE_THETA_1] +
state[self.STATE_THETA_2] - np.pi / 2)
x1 = self.LENGTH * np.cos(state[self.STATE_THETA_1] - np.pi / 2)
y1 = self.LENGTH * np.sin(state[self.STATE_THETA_1] - np.pi / 2)
# x2 = (x + 2 * self.LENGTH) / (4 * self.LENGTH)
# y2 = (y + 2 * self.LENGTH) / (4 * self.LENGTH)
return x1, y1, x2, y2
def get_state_bounds(self):
return [(self.MIN_ANGLE, self.MAX_ANGLE),
(self.MIN_ANGLE, self.MAX_ANGLE), (self.MIN_V_1, self.MAX_V_1),
(self.MIN_V_2, self.MAX_V_2)]
def get_control_bounds(self):
return [(self.MIN_TORQUE, self.MAX_TORQUE)]
plt.axis('equal')
img = fig2img(fig)
plt.close()
return img
if __name__ == '__main__':
model_fn = '/home/fei/Development/physics_predmodel/gn/logs/runs/October01_14:59:16/model_1240000.pth'
dset = SwimmerDataset('swimmer_test.npy')
use_cuda = True
dl = DataLoader(dset, batch_size=200, num_workers=0, drop_last=True)
node_feat_size = 6
edge_feat_size = 3
graph_feat_size = 10
gn = FFGN(graph_feat_size, node_feat_size, edge_feat_size).cuda()
gn.load_state_dict(torch.load(model_fn))
normalizers = torch.load('normalize.pth')
in_normalizer = normalizers['in_normalizer']
out_normalizer = normalizers['out_normalizer']
G1 = nx.path_graph(6).to_directed()
dl_e = enumerate(dl)
data = dset.__get_episode__(353)
data = [torch.from_numpy(item) for item in data]
writer = imageio.get_writer('test_pred.mp4', fps=6)
action, delta_state, last_state = data
action, delta_state, last_state = action.float(), delta_state.float(
), last_state.float()
from gn_models import init_graph_features, FFGN
from utils import load_graph_features_sst_acrobot as load_graph_features
import networkx as nx
from sparse_rrt.systems import Acrobot
import numpy as np
import torch
from matplotlib import pyplot as plt
G1 = nx.path_graph(2).to_directed()
node_feat_size = 2
edge_feat_size = 3
graph_feat_size = 10
gn = FFGN(graph_feat_size, node_feat_size, edge_feat_size).cuda()
gn.load_state_dict(torch.load('model0.05.pth'))
normalizers = torch.load('normalized/acrobot0.05.pth')
in_normalizer = normalizers['in_normalizer']
out_normalizer = normalizers['out_normalizer']
model = Acrobot()
action = [0.]
start_state = np.array([
np.random.uniform(low=model.MIN_ANGLE, high=model.MAX_ANGLE),
np.random.uniform(low=model.MIN_ANGLE, high=model.MAX_ANGLE),
np.random.uniform(low=model.MIN_V_1, high=model.MAX_V_1),
np.random.uniform(low=model.MIN_V_2, high=model.MAX_V_2)
])
maxframe = 15
baseline = np.zeros((maxframe + 1, 4))