decmcts.py 21.1 KB
Newer Older
1
import gcode2contour as gc
2
from pydecmcts import Tree
3
import utils as tu
4
from collections import deque
5

6
import matplotlib.pyplot as plt
7
from math import pi
8
import time
9
10
import numpy as np
import pickle
11
import copy
12
import warnings
13
import random
14
15


16
home = [0, -pi/2, pi/4, -pi/4, -pi/2, -pi/4]
17
SAFE_TIME = 2
18
J2TIME_C = 2.7
19
20


21
def avail_contours(data, state, robot_id):
22
23
24
    done_contours = [ct.contour for ct in state]
    return [ct for ct in data['actions'][robot_id] \
            if ct.contour not in done_contours]
25
26
27
28
29
30


def state_storer(data, parent_state, action, robot_id):
    """
    state at a node is just a list of ctrajs planned so far
    """
31
    # root node catch
32
    if parent_state is None:
33
34
35
36
        return deque([])
    new_state = copy.copy(parent_state)
    new_state.append(action)
    return new_state
37
38


39
def joint_diff(js1,js2):
40
    return sum([(j1-j2)**2 for j1,j2 in zip(js1,js2)])
41
42


43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def sim_select(data, options, sim_state, method='greedy'):
    """
    How to select option during simulation
    """
    if method == 'greedy':
        j  = sim_state[-1].positions[-1]
        costs = [joint_diff(j,o.positions[0]) for o in options]
        return options[np.argmin(costs)]
    elif method == 'random':
        return random.choice(options)
    elif method == 'weighted':
        j  = sim_state[-1].positions[-1]
        weights = [1/joint_diff(j,o.positions[0]) for o in options]
        return random.choices(options, weights=weights)[0]
    else:
        print("ERROR: Invalid sim_select method: '{}'".format(method))
59
60


61
class CombinedSchedule:
62
    """
63
    combine decmcts rollouts
64
65
    """

66
    def __init__(self, data, n_robots, joint2time_mul=J2TIME_C):
67
68
        self.data = data
        self.col_table = data["collisions"]
69
70
        self.n = n_robots
        self.trajs = [[] for i in range(self.n)]
71
72
        self.start_times = [[0] for i in range(self.n)]
        self.end_times = [[0] for i in range(self.n)]
73
        self.t_mul = joint2time_mul
74
        self.safe_time = SAFE_TIME # safety buffer for Collisions
75
        self.start_next_contour = [0]*self.n
76
77
78


    def planLenOrder(self):
79
80
81
82
83
        """
        orders robots by shortest plan first
        """
        plan_times = [self.end_times[r][-1] for r in range(self.n)]
        return np.argsort(plan_times)
84
85


86
87
88
89
    def lastJointPos(self, robot):
        """
        last joint position of a robot
        """
90
91
        if len(self.trajs[robot])==0:
            return self.data["env"].home
92
        return self.trajs[robot][-1].positions[-1]
93
94


95
    def len(self):
96
97
98
99
        """
        total time of plan is max time of robots
        """
        return max([self.end_times[r][-1] for r in range(self.n)])
100
101


102
    def append(self, ct, robot, force=False):
103
        """
104
        add a contour trajectory to the CombinedSchedule
105
        """
106

107
108
109
110
111
        # catch for n = 1
        if self.n == 1:
            self.trajs[robot].append(ct)
            self.start_times[robot].append(self.end_times[robot][-1])
            self.end_times[robot].append(self.start_times[robot][-1] + ct.time[-1])
112
            return True
113

114
115
        # travel time based on joint_diff
        trav_time = self.t_mul * joint_diff(ct.positions[-1], self.lastJointPos(robot))
116
        t_0 = self.start_next_contour[robot] + trav_time
117

118
        # for every other robot, get collision free time offsets
119
120
        allowed_offsets = []
        for r in range(self.n):
121
122
            if r == robot:
                continue
123

124
125
126
127
128
129
            t_diff = t_0 - self.start_times[r][-1]

            # if no trajectories in other robot, can start immidiately
            if len(self.trajs[r])==0:
                ao = [] # empty allowed because (0 to 1mils will be added below)
            else:
130

131
132
                # get last trajectory on other robot
                ct_r = self.trajs[r][-1]
133

134
135
136
137
138
139
140
141
                # got collision free offsets for that trajectory
                if (ct, ct_r) in self.col_table.keys():
                    ao = copy.deepcopy(self.col_table[ct, ct_r])
                elif (ct_r, ct) in self.col_table.keys():
                    ao = [(-b,-a) for a,b in
                            reversed(self.col_table[ct_r, ct])]
                else:
                    self.col_table[ct, ct_r] = \
142
                            tu.col_free_offsets(ct, ct_r, self.data["env"], safe_time=self.safe_time)
143
                    ao = copy.deepcopy(self.col_table[ct, ct_r])
144

145
            # apply t_diff and add allowed offset for when other robot is done
146
            allowed_offsets.append([(-1e6, t_0 - self.end_times[r][-1]-self.safe_time)] + \
147
                [(a+t_diff, b+t_diff) for a,b in ao if a+t_diff<=0])
148

149
150
        # easy for n = 2
        if self.n == 2:
151
152
153
154
155
156
157
158
159
160
161
            t_wait = max(trav_time, -min(0, allowed_offsets[0][-1][1]))

            #if this contour starts after other arm is finished
            # and it isnt the first contour of the plan
            # and it isn't being added regardless
            if t_0 + t_wait > self.end_times[1-robot][-1] \
                    and not (self.end_times[robot][-1] == 0 and self.end_times[1-robot][-1] == 0) \
                    and not force:
                # reject contour if it cannot be parallelised with current plan
                self.start_next_contour[robot] = self.end_times[1-robot][-1] + self.safe_time
                return False
162
            self.trajs[robot].append(ct)
163
            self.start_times[robot].append(t_0 + t_wait)
164
            self.end_times[robot].append(self.start_times[robot][-1] + ct.time[-1])
165
166
            self.start_next_contour[robot] = self.end_times[robot][-1]
            return True
167
168

        # for n > 2
169
        t_wait_last = -1
170
        while True:
171
            t_wait = max(trav_time,-min([0] + [o[-1][1] for o in allowed_offsets]))
172
173
174
            if t_wait_last == t_wait:
                Exception("Logic error, this should never happen")
            for ro in allowed_offsets:
175
                while ro[-1][0] > -t_wait:
176
                    ro.pop()
177
            valid = [ro[-1][0] <= -t_wait <= ro[-1][1] for ro in allowed_offsets]
178
179
180
181
            if all(valid):
                break
            t_wait_last = t_wait

182
183
184
185
        self.trajs[robot].append(ct)
        self.start_times[robot].append(self.end_times[robot][-1] + t_wait)
        self.end_times[robot].append(self.start_times[robot][-1] + ct.time[-1])
        return t_wait
186
187


188
def visualiseCombinedSchedule(combined_schedule):
189
    """
190
191
192
193
194
195
196
197
198
199
200
201
202
203
    Plot combined scedule so we can see density of results
    """
    print("green vertical line mark start, red vertical linesmark end of contour")
    fig, ax = plt.subplots()
    for r in range(combined_schedule.n):
        for t in range(1,len(combined_schedule.start_times[r])):
            ax.plot([combined_schedule.start_times[r][t],
                     combined_schedule.end_times[r][t]],
                    [r,r],
                    'bo-',
                    linewidth=2)

            ax.plot([combined_schedule.start_times[r][t],
                     combined_schedule.start_times[r][t]],
204
                    [0,combined_schedule.n-1],
205
206
207
208
209
                    'g--',
                    linewidth=0.5)

            ax.plot([combined_schedule.end_times[r][t],
                     combined_schedule.end_times[r][t]],
210
                    [0,combined_schedule.n-1],
211
212
213
214
215
216
217
218
219
220
221
                    'r--',
                    linewidth=0.5)
    plt.xlabel("Time")
    plt.ylabel("Robots")
    plt.show()
    return


def createCombinedSchedule(data, states):
    """
    given states, create a combined schedule
222
    """
223
    plan = CombinedSchedule(data, len(states))
224
    cts_left = [copy.copy(s) for s in states.values()]
225

226
227
228
    if len(cts_left)>2:
        raise Exception("New CombinedScheduler does not support >2 arms")

229
    done_contours = []
230
    t_used = 0
231

232
233
234
    robot = 0
    force=False
    assert(len(cts_left[0])>0)
235
236
    while sum([len(r) for r in cts_left])>0:
        # get next cts whose contour hasn't been done
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
        i = 0
        while i < len(cts_left[robot]):
            ct = cts_left[robot][i]

            if ct.contour in done_contours:
                del cts_left[robot][i]
                if len(cts_left[robot]) == 0:
                    robot = 1-robot
                    break
                continue
            else:
                # try adding
                if plan.append(ct, robot, force=force):

                    # log done contours and time
                    done_contours.append(ct.contour)
                    t_used += ct.time[-1]
                    force = False

                    # get next arm to add to
                    rs = plan.planLenOrder()
                    for r in rs:
                        if len(cts_left[r]) != 0:
                            robot = r
                            break
                    break

                # if tried everything, switch robots
                elif i == len(cts_left[robot])-1:
                    # TODO work with more than two robots
                    robot = 1-robot
                    force=True
                    break
                else:
                    # try next contour in list
                    i+=1

274
    return plan, t_used
275

276
277
278
279
280
281

def reward(data, states):
    """
    Joint Scheduler + reward function on outputplan
    """
    plan, t_used = createCombinedSchedule(data, states)
282
    # return t_used/(plan.len()*len(states)), plan # FOR TESTING
283
    return t_used/plan.len()
284
285


286
287
288
def decmcts_layer(
            contours,
            contour_tracker,
289
            env,
290
291
292
            plan=None,
            finish_home=True,
            col_print=False,
293
            mcts_nodes=100, # no. of nodes to grow on each tree
294
            comm_n=10, # no. of action dists to communicate
295
            joint2time_mul=J2TIME_C,
296
            default_wait = 5 # time spent waiting when traj add fails
297
298
299
300
            ):
    """
    decmcts planner
    """
301
302
303
304
305
    # update collision box for printed material
    if col_print:
        bound_mins, bound_maxs = tu.get_bounding_box(contour_tracker)
        bound_maxs[2] = tracker[contours[0]].pos[0][2]/1000 - 0.001
        print_block = env.add_block(bound_mins, bound_maxs)
306

307
308
    if plan is None:
        plan = tu.Plan(env)
309
310
311

    # save calculated trajectories for contours between loops
    c_traj = {}
312
    n_arms = env.numRobots
313
314
315
316
317
318
319
320
321
    arm_done = [False]*n_arms

    # DecMCTS variables
    data = {'env': env, # data used by reward function
            'collisions': {},
            'actions': [None]*n_arms}

    # DecMCTS main loop

322
    s_layer = time.time()
323

324
325
326
    # Update Data for trees to include latest contours
    for cid in contours:
        for arm in range(n_arms):
327

328
329
            if cid not in c_traj:
                c_traj[cid] = [None]*n_arms
330

331
332
333
334
335
336
            if c_traj[cid][arm] is None:
                multiple_trajs = tu.contour_trajs(
                        contour_tracker[cid],
                        env,
                        arm
                        )
337

338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
                # only one ik per contour
                if len(multiple_trajs)>0:
                    c_traj[cid][arm] = [tu.get_closest_traj(multiple_trajs,home)]
                else:
                    c_traj[cid][arm] = []

        # is this contour impossible?
        if c_traj[cid] == [[]]*n_arms:
            warnings.warn("Contour {} impossible for both arms, SKIPPING".format(cid))
            #contour_tracker.do_contour(cid)

    # Clear old actions and unroll new ones into flat list
    for r in range(n_arms):
        data['actions'][r] = [ct
                              for c in contours
                              for ct in c_traj[c][r]
                              if ct is not None]

    # Create DecMCTS trees
    trees = []
    for i in range(n_arms):
        trees.append(Tree(data,
                          reward,
                          avail_contours,
                          state_storer,
                          sim_select,
                          avail_contours,
                          comm_n,
                          i))

    # grow mcts_nodes nodes
    s = time.time()
    for i in range(mcts_nodes):
        # grow the trees
372
        for tree in trees:
373
            tree.grow(nsims=1, depth=1e6)
374
        if i%1==0 and i!=0:
375
376
            e = time.time()
            print("Completed {} nodes in {:.2g}s".format(i, e-s))
377

378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
        # send update to other trees
        for tree in trees:
            for other_tree in trees:
                if other_tree is not tree:
                    other_tree.receive_comms(tree.send_comms(), tree.id)

    # Best rollouts
    top_rolls = {}
    for tree in trees:
        top_rolls[tree.id] = tree.my_act_dist.best_action()

    # Get the CombinedSchedule of these rollouts
    combined_schedule, _ = createCombinedSchedule(data, top_rolls)

    # Add contours from combined schedule to real plan
    contours_left = sum([len(l) for l in combined_schedule.trajs])
    inds = [0] * n_arms
    while contours_left>0:

        # select trajectory to add next (as per plan not combined schedule)
        arms = np.argsort([plan.cumul_time[arm][-1] for arm in range(plan.n)])
        arm = None
        for i in arms:
            if inds[i] < len(combined_schedule.trajs[i]):
                arm = i
                break
404

405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
        assert arm != None, "something's gone wrong here - all arms complete but total is not"

        # Trajectory to add and time to wait for it
        traj = combined_schedule.trajs[arm][inds[arm]]

        # TODO THIS VAIRABLE SHOULD BE RELATIVE TO OTHER ARM
        decmcts_wait_time = combined_schedule.start_times[arm][inds[arm]] - \
                combined_schedule.end_times[arm][inds[arm]-1]

        # Estimated travel times
        time_last2next = joint2time_mul * \
                joint_diff(plan.trajs[arm][-1].positions[-1], traj.positions[0])
                # this should be minimum value of decmcts_wait_time
        time_last2home = joint2time_mul * \
            joint_diff(plan.trajs[arm][-1].positions[-1], home)
        time_home2next = joint2time_mul * \
            joint_diff(home, traj.positions[0])

        # if there's enough time to get home in the wait period, do that
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
        #athome = list(plan.getLastJoint(arm)) == home
        #additional_time_spent = 0
        #if time_last2home + time_home2next <= decmcts_wait_time:

        #    # try going to home pos
        #    while not athome:
        #        gohome = plan.plan_to(env.home, arm)

        #        # if no travel found, stay at last pos
        #        if gohome is not None and plan.appendTrajectory([gohome], arm)==0:
        #            print(arm, "returning to home to wait")
        #            athome = True
        #        else:
        #            print("There's enough time to get home but couldn't find a path - Unfortunate")
        #            print("Staying put for {}s".format(default_wait))
        #            last_js = plan.trajs[arm][-1].positions[-1]
        #            stayput = tu.JTrajectory([last_js, last_js],[0,default_wait], arm)
        #            if plan.appendTrajectory([stayput], arm)!=0:
        #                print("cannot go home or stay put Logically possible but unlikely - TODO")
        #                raise Exception("cannot go home or stay put Logically possible but unlikely - TODO")
        #            additional_time_spent += default_wait

        #            if time_last2home + additional_time_spent + default_wait + time_home2next > decmcts_wait_time:
        #                # can't go home anymore
        #                break
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475

        # add to plan
        traj_added = False
        retry_counter = 0
        while not traj_added:

            #print("---")
            # create plan to next traj
            travel_traj = plan.plan_to(traj.positions[0], arm)
            # try adding traj + contour to planner
            if travel_traj is not None and plan.appendTrajectory([travel_traj, traj], arm) == 0:
                # if success, break, else loop again
                print("contour", traj.contour, "added to real plan")
                #contour_tracker.do_contour(traj.contour)
                traj_added = True
                break
            else:
                #print("deviating from decmcts plan")
                # If the cumul time of current arm is greater than that of other arms, send other arms home
                if plan.cumul_time[arm][-1] == plan.len():
                    print("current arm has the longest plan but failing, take other arms home to clear")
                    # send other arms home
                    for other_arm in plan.trajs:
                        if other_arm == arm:
                            continue
                        # if already at home, stay at home
                        if all(plan.getLastJoint(other_arm) == home):
476
                            print("Other arm is already home, this should not occur !!!")
477
478
479
480
481
482
483
484
                            raise Exception(" this should not occur")

                        # Otherwise, go home
                        gohome = plan.plan_to(env.home, other_arm)
                        if gohome is not None and plan.appendTrajectory([gohome], other_arm)==0:
                            print("current arm is", arm, ", taking other arm home")
                            print("Arm", other_arm, "returning to home")
                            athome = True
485
                        else:
486
                            # other arm cannot go home, send current arm home to make room
487
488
489
                            if all(plan.getLastJoint(arm) == home):
                                print("other arm cannot go home despite current arm already being home - this should not happen")
                                raise Exception("other arm cannot go home despite current arm already being home - this should not happen")
490
491
                            gohome = plan.plan_to(env.home, arm)
                            if gohome is not None and plan.appendTrajectory([gohome], arm)==0:
492
                                print("other arm {} cannot go home, taking arm {} home since it cannot stay put".format(other_arm, arm))
493
                            else:
494
                                print("Neither arm can go home!")
495
                                raise Exception("Neither arm can go home!")
496

497
498
499
500
501
502
503
504
                last_js = plan.trajs[arm][-1].positions[-1]
                stayput = tu.JTrajectory([last_js, last_js],[0,default_wait], arm)
                if plan.appendTrajectory([stayput], arm)!=0:
                    print("cannot stay at last position, need to move out of way, going home")
                    gohome = plan.plan_to(env.home, arm)
                    if gohome is not None and plan.appendTrajectory([gohome], arm)==0:
                        print("taking arm {} home since it cannot stay put".format(arm))
                    else:
505
                        print("Cannot stay put or go home!")
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
                        raise Exception("Cannot stay put or go home!")


        contours_left -= 1 # to ensure loop exits
        inds[arm] += 1 # move to next contour next time

    # take arms home
    if finish_home:
        arms_home=[list(plan.getLastJoint(a)) == home for a in range(n_arms)]
        arm = 0
        retries = 0
        while not all(arms_home) and retries < n_arms:
            gohome = plan.plan_to(env.home, arm)
            if list(plan.getLastJoint(arm)) != home and \
                    gohome is not None and \
                    plan.appendTrajectory([gohome], arm)==0:
                print("Arm {} safely returned to home pose".format(arm))
                retries = 0
                arms_home[arm] = True
            else:
                print("Failed to take arm {} home, trying another arm".format(arm))
                retries +=1
            arm += 1
            if arm == n_arms:
                arm = 0
531

532
533
    e_layer = time.time()
    print("last layer completed in {}s".format(e_layer-s_layer))
534

535
    return combined_schedule, plan
536
537


538
539
540
541
542
def decmcts(
        sim_env,
        tracker,
        nth_layer=1,
        layers=None,
Jayant Khatkar's avatar
Jayant Khatkar committed
543
        nodes_per_layer=1000,
544
545
546
547
548
549
550
551
552
553
        interactive=False):
    """
    plan whole print with decmcts using decmcts_layer
    """
    if layers is None:
        n_layers = len(tracker.layers)
        layers = list(range(0, n_layers, nth_layer))

    plan = None
    fin_home = False
554
    print("Planning the following layers: {}".format(layers))
555
556

    for i in layers:
557
        print("PLANNING LAYER {}".format(i))
558
559
560
561
562
563
564
565
566
567
568

        # take arms home if final layer
        if i == layers[-1]:
            fin_home = True

        cs, plan = decmcts_layer(
                tracker.layers[i],
                tracker,
                sim_env,
                plan=plan,
                finish_home=fin_home,
Jayant Khatkar's avatar
Jayant Khatkar committed
569
                mcts_nodes=nodes_per_layer
570
571
572
573
574
575
576
577
                )

        if interactive:
            visualiseCombinedSchedule(cs)
            plan.visualise()
    return plan


578
579
580
if __name__ == "__main__":

    # Read gcode file
581
    model = "CuteOcto_x2"
582
    tracker = tu.get_tracker(model)
583

584
585
586
587
588
589
    # 3arms simulation
    env_desc3 = tu.PlanningEnv([[0,0,1,0],[0,0,1,0],[0,0,0,1]],
                       [[0.9, -0.35, 0],[-0.8, -0.35, 0.001],[0, 0.5, 0]],
                       [[0,0,0,1]]*3,
                       [[0.15,0,0]]*3) # WARNING - also edit urdf

590
591
592
    # 2arms real calibration
    env_desc2 = tu.read_env('calibrations/r1_tforms.yaml',
                            'calibrations/r2_tforms.yaml')
593
    sim_env = tu.SimEnv(env_desc=env_desc2, gui=True, home=home)
594

595
    # decmcts plan first n contours
596
    plan = decmcts(sim_env, tracker, layers=[1,5], interactive=True)