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
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() action, delta_state, last_state = data action, delta_state, last_state = action.float(),\ delta_state.float(), last_state.float() if use_cuda: action, delta_state, last_state = action.cuda(),\ delta_state.cuda(), last_state.cuda() init_graph_features(G1, graph_feat_size, node_feat_size, edge_feat_size, cuda=True, bs=200) load_graph_features(G1, action, last_state, delta_state, bs=200, noise=0.02, std=std) G_out = gn(in_normalizer.normalize(G1)) init_graph_features(G_target, graph_feat_size, node_feat_size, edge_feat_size,
if use_cuda: action, delta_state, last_state = action.cuda(), delta_state.cuda( ), last_state.cuda() state = last_state[1].unsqueeze(0) state_gt = last_state[1].unsqueeze(0).clone() for i in range(1, 50): print(i) action_i = action[i].unsqueeze(0) delta_state_i = delta_state[i].unsqueeze(0) last_state_i = last_state[i].unsqueeze(0) init_graph_features(G1, graph_feat_size, node_feat_size, edge_feat_size, cuda=True, bs=1) load_graph_features(G1, action_i, state, None, bs=1, noise=0) G_out = gn(in_normalizer.normalize(G1)) G_out = out_normalizer.inormalize(G_out) delta_state_pred = get_graph_features(G_out) state_gt += delta_state_i state += delta_state_pred img = draw_state(state_gt) img_pred = draw_state(state) writer.append_data(np.concatenate([img, img_pred], axis=1))