Merge remote-tracking branch 'origin/feature/decmcts'

src/test_decmcts.py

@@ -4,51 +4,11 @@ Test Functions in decmcts.py
 from decmcts import *
 
 
-SAFE_TIME = 2
-
-
-def test_reward(contour_tracker, n_contours=None, n_arms=2, gui=False):
+def test_reward(contour_tracker, combined_schedule, real_plan):
     """
     test _reward function
     """
-    # Parse inputs
-    if n_contours is None:
-        #n_contours = len(contour_tracker.all_contours)
-        n_contours = len(contour_tracker.available_contours)
-
-    env = tu.SimEnv(gui=gui, home=home) # pyBullet env for collision checking
-    #plan = tu.Plan(env, [env.home, env.home]) # planner for 2 arms
-
-    #plan for these contours
-    contours = list(range(n_contours))
-
-    c_traj = {}
-    for arm in range(n_arms):
-        for cid in contours:
-
-            if cid not in c_traj:
-                c_traj[cid] = [None]*n_arms
-
-            c_traj[cid][arm] = tu.contour_trajs(
-                    contour_tracker.all_contours[cid],
-                    env,
-                    arm
-                    )
-
-    # random order of trajectories
-    np.random.seed(1)
-    arm1_plan = deque([c_traj[c][0][0] for c in contours if len(c_traj[c][0])>0])
-    np.random.shuffle(arm1_plan)
-    arm2_plan = deque([c_traj[c][1][0] for c in contours if len(c_traj[c][1])>0])
-    np.random.shuffle(arm2_plan)
-
-    data = {"env": env,
-            "collisions": {}}
-    s = time.time()
-    p,r = createCombinedSchedule(data, {0:arm1_plan, 1:arm2_plan})
-    e = time.time()
-    print(e-s)
-
+    p = combined_schedule
 
     print("orange horizonatl line means looking for start time above" + \
           ", black line means looking for a start time below\n")
@@ -59,8 +19,10 @@ def test_reward(contour_tracker, n_contours=None, n_arms=2, gui=False):
         ct_r0, ct_r1 = p.trajs[0][i-1], p.trajs[1][j-1]
         t_wait = p.start_times[1][j] - p.start_times[0][i]
 
+        print("({},{})".format(i,j))
+
         tu.vis_col_region(ct_r0, ct_r1, env, res=1, chosen_time=t_wait, safe_time = SAFE_TIME)
-       
+
         # increment ct we're up to
         if p.end_times[0][i] <= p.end_times[1][j]:
             i +=1
@@ -69,10 +31,14 @@ def test_reward(contour_tracker, n_contours=None, n_arms=2, gui=False):
             j +=1
             robot = 1
 
-        if input('enter for next, or "q" for quit')=='q':
-              break
+        user_input = input('enter for next, "p" for show plan and next, or "q" for quit: ')
+        if user_input =='q':
+            break
+        elif user_input == 'p':
+            visualiseCombinedSchedule(p)
+
+    return
 
-    return data, r, p, c_traj
 
 if __name__ == "__main__":
 
@@ -80,5 +46,18 @@ if __name__ == "__main__":
     model = "flexirex_big"
     tracker = tu.get_tracker(model)
 
-    # run test test_reward 
-    data, r, p, c_traj = test_reward(tracker, n_contours=50, n_arms=2, gui=False)
+    # 2arms real calibration
+    env_desc2 = tu.read_env('calibrations/r1_tforms.yaml',
+                            'calibrations/r2_tforms.yaml')
+    env = tu.SimEnv(env_desc=env_desc2, gui=False, home=home)
+
+    # decmcts plan first n contours
+    c_traj_file = "gcode/" + model + ".c_traj"
+    cs, plan = decmcts(tracker,
+                   env,
+                   n_layers=1,
+                   mcts_nodes=1000
+                   )
+
+    # run test test_reward
+    test_reward(tracker, cs, plan)