add framework for training

824b68b7 · lvzhengyang · 96539d53 · 824b68b7
Commit 824b68b7 authored Sep 02, 2022 by lvzhengyang
Hide whitespace changes
Inline Side-by-side

Showing with 11 additions and 6 deletions

train.py
+11 -6

No files found.
--- a/train.py
+++ b/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)
+  for i in range(num_nodes):
-  # v/h is a tensor of shape [#batch_size] on model.device
+    if i == 0:
-  mask_visited, rewards, done, _ = envs.step(v, h)
+      u0 = model.first_act(nodes_xy)
-  if done:
+      mask_visited, rewards, done, _ = envs.step(u0)
-    model.learn(nodes_xy, rewards)
+      continue
-    nodes_xy, mask_visited = envs.reset()
+    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()