def is_valid(self, state):
if has_nonterminals(state):
return False
robot = build_normalized_robot(state)
_, has_self_collision = presimulate(robot)
return not has_self_collision
def sample_design(args, task_id, seed, env, V, eps, results_queue, time_queue,
done_event):
tt0 = time.time()
random.seed(seed)
valid = False
samples = []
while not valid:
state = make_initial_graph()
rule_seq = []
no_action_flag = False
for _ in range(args.depth):
action, step_type = select_action(env, V, state, eps)
if action is None:
no_action_flag = True
break
rule_seq.append(action)
next_state = env.transite(state, action)
state = next_state
if not has_nonterminals(state):
break
valid = env.is_valid(state)
if not valid:
# update the invalid sample's count
if no_action_flag:
info = 'no_action'
elif has_nonterminals(state):
info = 'step_exceeded'
else:
info = 'self_collision'
samples.append(Sample(task_id, rule_seq, -2.0, info))
else:
samples.append(
Sample(task_id, rule_seq, predict(V, state), info='valid'))
tt = time.time() - tt0
time_queue.put(tt)
results_queue.put(samples)
done_event.wait()
def compute_Vhat(robot_graph, env, V):
if has_nonterminals(robot_graph):
available_actions = env.get_available_actions(robot_graph)
if len(available_actions) == 0:
return -5.0
next_states = []
for action in available_actions:
next_states.append(env.transite(robot_graph, action))
values = predict_batch(V, next_states)
return np.max(values)
else:
return env.get_reward(robot_graph)[1]
def step(self, action):
next_state = self.transite(self.state, action)
self.rule_seq.append(action)
if has_nonterminals(next_state):
reward = 0.
done = False
else:
input_sequence, reward = self.get_reward(next_state)
done = True
self.state = next_state
return self.state, reward, done
def search_algo(args):
# iniailize random seed
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.set_num_threads(1)
# initialize/load
task_class = getattr(tasks, args.task)
if args.no_noise:
task = task_class(force_std=0.0, torque_std=0.0)
else:
task = task_class()
graphs = rd.load_graphs(args.grammar_file)
rules = [rd.create_rule_from_graph(g) for g in graphs]
# initialize preprocessor
# Find all possible link labels, so they can be one-hot encoded
all_labels = set()
for rule in rules:
for node in rule.lhs.nodes:
all_labels.add(node.attrs.require_label)
all_labels = sorted(list(all_labels))
# TODO: use 80 to fit the input of trained MPC GNN, use args.depth * 3 later for real mpc
max_nodes = args.depth * 3
global preprocessor
# preprocessor = Preprocessor(max_nodes = max_nodes, all_labels = all_labels)
preprocessor = Preprocessor(all_labels=all_labels)
# initialize the env
env = RobotGrammarEnv(task,
rules,
seed=args.seed,
mpc_num_processes=args.mpc_num_processes)
# initialize Value function
device = 'cpu'
state = env.reset()
sample_adj_matrix, sample_features, sample_masks = preprocessor.preprocess(
state)
num_features = sample_features.shape[1]
V = Net(max_nodes=max_nodes, num_channels=num_features,
num_outputs=1).to(device)
# load pretrained V function
if args.load_V_path is not None:
V.load_state_dict(torch.load(args.load_V_path))
print_info('Loaded pretrained V function from {}'.format(
args.load_V_path))
# initialize target V_hat look up table
V_hat = dict()
# load pretrained V_hat
if args.load_Vhat_path is not None:
V_hat_fp = open(args.load_Vhat_path, 'rb')
V_hat = pickle.load(V_hat_fp)
V_hat_fp.close()
print_info('Loaded pretrained Vhat from {}'.format(
args.load_Vhat_path))
# initialize invalid_his
invalid_his = dict()
num_invalid_samples, num_valid_samples = 0, 0
repeated_cnt = 0
# initialize the seen states pool
states_pool = StatesPool(capacity=args.states_pool_capacity)
states_set = set()
# explored designs
designs = []
design_rewards = []
design_opt_seeds = []
# record prediction error
prediction_error_sum = 0.0
if not args.test:
# initialize save folders and files
fp_log = open(os.path.join(args.save_dir, 'log.txt'), 'w')
fp_log.close()
fp_eval = open(os.path.join(args.save_dir, 'eval.txt'), 'w')
fp_eval.close()
design_csv_path = os.path.join(args.save_dir, 'designs.csv')
fp_csv = open(design_csv_path, 'w')
fieldnames = ['rule_seq', 'reward', 'opt_seed']
writer = csv.DictWriter(fp_csv, fieldnames=fieldnames)
writer.writeheader()
fp_csv.close()
# initialize the optimizer
global optimizer
optimizer = torch.optim.Adam(V.parameters(), lr=args.lr)
# initialize best design rule sequence
best_design, best_reward = None, -np.inf
# reward history
epoch_rew_his = []
last_checkpoint = -1
# recording time
t_sample_sum = 0.
# record the count for invalid samples
no_action_samples, step_exceeded_samples, self_collision_samples = 0, 0, 0
for epoch in range(args.num_iterations):
t_start = time.time()
V.eval()
# update eps and eps_sample
if args.eps_schedule == 'linear-decay':
eps = args.eps_start + epoch / args.num_iterations * (
args.eps_end - args.eps_start)
elif args.eps_schedule == 'exp-decay':
eps = args.eps_end + (args.eps_start - args.eps_end) * np.exp(
-1.0 * epoch / args.num_iterations / args.eps_decay)
if args.eps_sample_schedule == 'linear-decay':
eps_sample = args.eps_sample_start + epoch / args.num_iterations * (
args.eps_sample_end - args.eps_sample_start)
elif args.eps_sample_schedule == 'exp-decay':
eps_sample = args.eps_sample_end + (
args.eps_sample_start - args.eps_sample_end) * np.exp(
-1.0 * epoch / args.num_iterations /
args.eps_sample_decay)
t_sample, t_update, t_mpc, t_opt = 0, 0, 0, 0
selected_design, selected_reward = None, -np.inf
selected_state_seq, selected_rule_seq = None, None
p = random.random()
if p < eps_sample:
num_samples = 1
else:
num_samples = args.num_samples
# use e-greedy to sample a design within maximum #steps.
for _ in range(num_samples):
valid = False
while not valid:
t0 = time.time()
state = env.reset()
rule_seq = []
state_seq = [state]
no_action_flag = False
for _ in range(args.depth):
action, step_type = select_action(env, V, state, eps)
if action is None:
no_action_flag = True
break
rule_seq.append(action)
next_state = env.transite(state, action)
state_seq.append(next_state)
state = next_state
if not has_nonterminals(state):
break
valid = env.is_valid(state)
t_sample += time.time() - t0
t0 = time.time()
if not valid:
# update the invalid sample's count
if no_action_flag:
no_action_samples += 1
elif has_nonterminals(state):
step_exceeded_samples += 1
else:
self_collision_samples += 1
# update the Vhat for invalid designs
update_Vhat(args,
V_hat,
state_seq,
-2.0,
invalid=True,
invalid_cnt=invalid_his)
# update states pool
update_states_pool(states_pool, state_seq, states_set,
V_hat)
num_invalid_samples += 1
else:
num_valid_samples += 1
t_update += time.time() - t0
predicted_value = predict(V, state)
if predicted_value > selected_reward:
selected_design, selected_reward = state, predicted_value
selected_rule_seq, selected_state_seq = rule_seq, state_seq
t0 = time.time()
repeated = False
if (hash(selected_design)
in V_hat) and (V_hat[hash(selected_design)] > -2.0 + 1e-3):
repeated = True
repeated_cnt += 1
reward, best_seed = -np.inf, None
for _ in range(args.num_eval):
_, rew = env.get_reward(selected_design)
if rew > reward:
reward, best_seed = rew, env.last_opt_seed
t_mpc += time.time() - t0
# save the design and the reward in the list
designs.append(selected_rule_seq)
design_rewards.append(reward)
design_opt_seeds.append(best_seed)
# update best design
if reward > best_reward:
best_design, best_reward = selected_rule_seq, reward
print_info(
'new best: reward = {:.4f}, predicted reward = {:.4f}, num_samples = {}'
.format(reward, selected_reward, num_samples))
t0 = time.time()
# update V_hat for the valid design
update_Vhat(args, V_hat, selected_state_seq, reward)
# update states pool for the valid design
update_states_pool(states_pool, selected_state_seq, states_set,
V_hat)
t_update += time.time() - t0
t0 = time.time()
# optimize
V.train()
total_loss = 0.0
for _ in range(args.opt_iter):
minibatch = states_pool.sample(
min(len(states_pool), args.batch_size))
train_adj_matrix, train_features, train_masks, train_reward = [], [], [], []
max_nodes = 0
for robot_graph in minibatch:
hash_key = hash(robot_graph)
target_reward = V_hat[hash_key]
# adj_matrix, features, masks = preprocessor.preprocess(robot_graph)
adj_matrix, features, _ = preprocessor.preprocess(
robot_graph)
max_nodes = max(max_nodes, len(features))
train_adj_matrix.append(adj_matrix)
train_features.append(features)
# train_masks.append(masks)
train_reward.append(target_reward)
for i in range(len(minibatch)):
train_adj_matrix[i], train_features[i], masks = \
preprocessor.pad_graph(train_adj_matrix[i], train_features[i], max_nodes)
train_masks.append(masks)
train_adj_matrix_torch = torch.tensor(train_adj_matrix)
train_features_torch = torch.tensor(train_features)
train_masks_torch = torch.tensor(train_masks)
train_reward_torch = torch.tensor(train_reward)
optimizer.zero_grad()
output, loss_link, loss_entropy = V(train_features_torch,
train_adj_matrix_torch,
train_masks_torch)
loss = F.mse_loss(output[:, 0], train_reward_torch)
loss.backward()
total_loss += loss.item()
optimizer.step()
t_opt += time.time() - t0
t_end = time.time()
t_sample_sum += t_sample
# logging
if (epoch + 1
) % args.log_interval == 0 or epoch + 1 == args.num_iterations:
iter_save_dir = os.path.join(args.save_dir,
'{}'.format(epoch + 1))
os.makedirs(os.path.join(iter_save_dir), exist_ok=True)
# save model
save_path = os.path.join(iter_save_dir, 'V_model.pt')
torch.save(V.state_dict(), save_path)
# save V_hat
save_path = os.path.join(iter_save_dir, 'V_hat')
fp = open(save_path, 'wb')
pickle.dump(V_hat, fp)
fp.close()
# save explored design and its reward
fp_csv = open(design_csv_path, 'a')
fieldnames = ['rule_seq', 'reward', 'opt_seed']
writer = csv.DictWriter(fp_csv, fieldnames=fieldnames)
for i in range(last_checkpoint + 1, len(designs)):
writer.writerow({
'rule_seq': str(designs[i]),
'reward': design_rewards[i],
'opt_seed': design_opt_seeds[i]
})
last_checkpoint = len(designs) - 1
fp_csv.close()
epoch_rew_his.append(reward)
avg_loss = total_loss / args.opt_iter
len_his = min(len(epoch_rew_his), 30)
avg_reward = np.sum(epoch_rew_his[-len_his:]) / len_his
prediction_error_sum += (selected_reward - reward)**2
avg_prediction_error = prediction_error_sum / (epoch + 1)
if repeated:
print_white('Epoch {:4}: T_sample = {:5.2f}, T_update = {:5.2f}, T_mpc = {:5.2f}, T_opt = {:5.2f}, eps = {:5.3f}, eps_sample = {:5.3f}, #samples = {:2}, training loss = {:7.4f}, pred_error = {:6.4f}, predicted_reward = {:6.4f}, reward = {:6.4f}, last 30 epoch reward = {:6.4f}, best reward = {:6.4f}'.format(\
epoch, t_sample, t_update, t_mpc, t_opt, eps, eps_sample, num_samples, \
avg_loss, avg_prediction_error, selected_reward, reward, avg_reward, best_reward))
else:
print_warning('Epoch {:4}: T_sample = {:5.2f}, T_update = {:5.2f}, T_mpc = {:5.2f}, T_opt = {:5.2f}, eps = {:5.3f}, eps_sample = {:5.3f}, #samples = {:2}, training loss = {:7.4f}, pred_error = {:6.4f}, predicted_reward = {:6.4f}, reward = {:6.4f}, last 30 epoch reward = {:6.4f}, best reward = {:6.4f}'.format(\
epoch, t_sample, t_update, t_mpc, t_opt, eps, eps_sample, num_samples, \
avg_loss, avg_prediction_error, selected_reward, reward, avg_reward, best_reward))
fp_log = open(os.path.join(args.save_dir, 'log.txt'), 'a')
fp_log.write('eps = {:.4f}, eps_sample = {:.4f}, num_samples = {}, T_sample = {:4f}, T_update = {:4f}, T_mpc = {:.4f}, T_opt = {:.4f}, loss = {:.4f}, predicted_reward = {:.4f}, reward = {:.4f}, avg_reward = {:.4f}\n'.format(\
eps, eps_sample, num_samples, t_sample, t_update, t_mpc, t_opt, avg_loss, selected_reward, reward, avg_reward))
fp_log.close()
if (epoch + 1) % args.log_interval == 0:
print_info(
'Avg sampling time for last {} epoch: {:.4f} second'.
format(args.log_interval,
t_sample_sum / args.log_interval))
t_sample_sum = 0.
print_info('size of states_pool = {}'.format(len(states_pool)))
print_info(
'#valid samples = {}, #invalid samples = {}, #valid / #invalid = {}'
.format(
num_valid_samples, num_invalid_samples,
num_valid_samples / num_invalid_samples
if num_invalid_samples > 0 else 10000.0))
print_info(
'Invalid samples: #no_action_samples = {}, #step_exceeded_samples = {}, #self_collision_samples = {}'
.format(no_action_samples, step_exceeded_samples,
self_collision_samples))
max_trials, cnt = 0, 0
for key in invalid_his.keys():
if invalid_his[key] > max_trials:
if key not in V_hat:
max_trials = invalid_his[key]
elif V_hat[key] < -2.0 + 1e-3:
max_trials = invalid_his[key]
if invalid_his[key] >= args.max_trials:
if V_hat[key] < -2.0 + 1e-3:
cnt += 1
print_info(
'max invalid_trials = {}, #failed nodes = {}'.format(
max_trials, cnt))
print_info('repeated rate = {}'.format(repeated_cnt /
(epoch + 1)))
save_path = os.path.join(args.save_dir, 'model_state_dict_final.pt')
torch.save(V.state_dict(), save_path)
else:
import IPython
IPython.embed()
# test
V.eval()
print('Start testing')
test_epoch = 30
y0 = []
y1 = []
x = []
for ii in range(0, 11):
eps = 1.0 - 0.1 * ii
print('------------------------------------------')
print('eps = ', eps)
reward_sum = 0.
best_reward = -np.inf
for epoch in range(test_epoch):
t0 = time.time()
# use e-greedy to sample a design within maximum #steps.
vaild = False
while not valid:
state = env.reset()
rule_seq = []
state_seq = [state]
for _ in range(args.depth):
action, step_type = select_action(env, V, state, eps)
if action is None:
break
rule_seq.append(action)
next_state = env.transite(state, action)
state_seq.append(next_state)
if not has_nonterminals(next_state):
valid = True
break
state = next_state
_, reward = env.get_reward(state)
reward_sum += reward
best_reward = max(best_reward, reward)
print(
f'design {epoch}: reward = {reward}, time = {time.time() - t0}'
)
print('test avg reward = ', reward_sum / test_epoch)
print('best reward found = ', best_reward)
x.append(eps)
y0.append(reward_sum / test_epoch)
y1.append(best_reward)
import matplotlib.pyplot as plt
fig, ax = plt.subplots(1, 2, figsize=(10, 5))
ax[0].plot(x, y0)
ax[0].set_title('Avg Reward')
ax[0].set_xlabel('eps')
ax[0].set_ylabel('reward')
ax[1].plot(x, y1)
ax[0].set_title('Best Reward')
ax[0].set_xlabel('eps')
ax[0].set_ylabel('reward')
plt.show()