add framework for training

train.py

@@ -29,10 +29,15 @@ model = Policy(actor=actor_net,
 
 nodes_xy, mask_visited = envs.reset()
 for episode in range(max_episode_num):
-  v, h = model.act(nodes_xy, mask_visited=mask_visited)
-  # v/h is a tensor of shape [#batch_size] on model.device
-  mask_visited, rewards, done, _ = envs.step(v, h)
-  if done:
-    model.learn(nodes_xy, rewards)
-    nodes_xy, mask_visited = envs.reset()
+  for i in range(num_nodes):
+    if i == 0:
+      u0 = model.first_act(nodes_xy)
+      mask_visited, rewards, done, _ = envs.step(u0)
+      continue
+    v, h = model.act(nodes_xy, mask_visited=mask_visited)
+    # v/h is a tensor of shape [#batch_size] on model.device
+    mask_visited, rewards, done, _ = envs.step(v, h)
+    if done:
+      model.learn(nodes_xy, rewards)
+      nodes_xy, mask_visited = envs.reset()