示例#1
0
def caction_generator2(uw, known_path, is_radius, amag):
    assert isinstance(
        known_path, list
    ), 'known_path should be a list rather than something else like np.array'
    print("### caction_generator2: FINDING INITIAL STATE")
    DEBUG = False
    if DEBUG:
        ni = 50
        iv = known_path[ni - 1]
    else:
        while True:
            # iv = uw_random.gen_unit_init_state(uw, is_radius)
            iv = _gen_unit_init_state(uw, known_path[0], is_radius)
            distances = pyosr.multi_distance(iv, known_path)
            ni = np.argmin(distances)
            # ni = 0
            '''
            if uw.is_valid_transition(iv, known_path[ni], amag / 8.0):
                break
            '''
            fstate, done, ratio = uw.transit_state_to(known_path[ni], iv,
                                                      amag / 8.0)
            if ratio > 1e-3:
                iv = fstate
                break
    vs = [iv] + known_path[ni:]
    # print('VS\n{}'.format(vs))
    # vs = known_path
    for items in trajectory_to_caction(vs, uw, amag):
        yield items
 def _get_gt_action_from_sample(self, q):
     distances = pyosr.multi_distance(q, self.unit_tunnel_v)
     ni = np.argmin(distances)
     close = self.unit_tunnel_v[ni]
     tr, aa, dq = pyosr.differential(q, close)
     assert pyosr.distance(close, pyosr.apply(
         q, tr, aa)) < 1e-6, "pyosr.differential is not pyosr.apply ^ -1"
     return np.concatenate([tr, aa], axis=-1), close
def pickup_key_configuration_old(args, ws):
    import pyosr
    scratch_dir = ws.local_ws(_KEYCAN_SCRATCH)
    os.makedirs(scratch_dir, exist_ok=True)
    top_k = ws.config.getint('TrainingKeyConf', 'KeyConf')
    fn_list = []
    # Do not process only_wait, we need to restart if ssh is broken
    for fn in pathlib.Path(scratch_dir).glob(
            'unitary_clearance_from_keycan-*.npz'):
        fn_list.append(fn)
    fn_list = sorted(fn_list)
    median_list = []
    max_list = []
    min_list = []
    mean_list = []
    stddev_list = []
    for fn in progressbar(fn_list[:]):
        d = matio.load(fn)
        distances = pyosr.multi_distance(d['FROM_V'], d['FREE_V'])
        median_list.append(np.median(distances))
        max_list.append(np.max(distances))
        min_list.append(np.min(distances))
        mean_list.append(np.mean(distances))
        stddev_list.append(np.std(distances))
    # mean works better than median, for some reason
    top_k_indices = np.array(mean_list).argsort()[:top_k]
    util.log('[pickup_key_configuration] top k {}'.format(top_k_indices))
    kq_ompl = []
    kq = []
    for k in top_k_indices:
        fn = fn_list[k]
        d = matio.load(fn)
        kq_ompl.append(d['FROM_V_OMPL'])
        kq.append(d['FROM_V'])
    key_out = ws.local_ws(util.KEY_FILE)
    stat_out = np.array(
        [median_list, max_list, min_list, mean_list, stddev_list])
    util.log(
        '[pickup_key_configuration] writting results to {}'.format(key_out))
    np.savez(key_out,
             KEYQ_OMPL=kq_ompl,
             KEYQ=kq,
             _STAT=stat_out,
             _TOP_K=top_k_indices)
    util.xz(key_out)
def pickup_key_configuration(args, ws):
    import pyosr
    uw = util.create_unit_world(
        ws.local_ws(util.TRAINING_DIR, util.PUZZLE_CFG_FILE))

    scratch_dir = ws.local_ws(_KEYCAN_SCRATCH)
    os.makedirs(scratch_dir, exist_ok=True)
    '''
    Pick up top_k from each trajectory
    '''
    top_k = ws.config.getint('TrainingKeyConf', 'KeyConf')
    candidate_file = _get_candidate_file(ws)
    cf = matio.load(candidate_file)
    trajs = sorted(list(cf.keys()))
    ntraj = len(trajs)
    util.log('[pickup_key_configuration] # of trajs {}'.format(ntraj))
    util.log('[pickup_key_configuration] actual trajs {}'.format(trajs))
    nq = cf[trajs[0]].shape[0]
    fn_list = sorted(
        pathlib.Path(scratch_dir).glob(
            'unitary_clearance_from_keycan-batch_*.hdf5.xz'))
    # fn_list = fn_list[:3] # Debug
    '''
    Load distances to traj_mean_list (dict of dict of index to list)
    '''
    # CAVEAT: stat_out may not be continuous
    # stat_out = np.full((ntraj, nq, 5), np.finfo(np.float64).max, dtype=np.float64)
    stat_dict = {
    }  # np.full((ntraj, nq, 5), np.finfo(np.float64).max, dtype=np.float64)
    for fn in progressbar(fn_list):
        # util.log("loading {}".format(fn)) # Debug
        d = matio.load(fn)
        # trajectory level
        for traj_name in d.keys():
            traj_id = int(traj_name[len('traj_'):])
            traj_grp = d[traj_name]
            for q_name in traj_grp.keys():
                q = traj_grp[q_name]
                if False:
                    # FIMXE: this distance is buggy
                    distances = uw.multi_kinetic_energy_distance(
                        q['FROM_V'], q['FREE_V'])
                else:
                    distances = pyosr.multi_distance(q['FROM_V'], q['FREE_V'])
                idx = int(str(q_name))
                # util.log("checking traj_name {} traj_id {} idx {}".format(traj_name, traj_id, idx)) # Debug
                m = np.mean(distances)
                if traj_id not in stat_dict:
                    stat_dict[traj_id] = np.full((nq, 5),
                                                 np.finfo(np.float64).max,
                                                 dtype=np.float64)
                stat_dict[traj_id][idx] = np.array([
                    np.median(distances),
                    np.max(distances),
                    np.min(distances), m,
                    np.std(distances)
                ])
                # util.log("checking traj_name {} traj_id {} idx {} mean {}".format(traj_name, traj_id, idx, m)) # Debug
    kq_ompl = []
    top_k_indices = []
    for traj_name in progressbar(trajs):
        # util.log("checking traj_name {}".format(traj_name)) # Debug
        traj_id = int(traj_name[len('traj_'):])
        stats = stat_dict[traj_id]
        current_top_k_indices = np.array(stats[:, 3]).argsort()[:top_k]
        # util.log("stat of traj_name {} traj_id {} top_k {}".format(traj_name, traj_id, current_top_k_indices)) # Debug
        for k in current_top_k_indices:
            kq_ompl.append(cf[traj_name][k])
            top_k_indices.append([traj_id, k])
    kq = uw.translate_ompl_to_unit(kq_ompl)
    key_out = ws.local_ws(util.KEY_FILE)
    # np.savez(key_out, KEYQ_OMPL=kq_ompl, KEYQ=kq, _STAT=stat_out, _TOP_K_PER_TRAJ=top_k_indices)
    stat_keys = []
    stat_values = []
    for key, value in stat_dict.items():
        stat_keys.append(key)
        stat_values.append(value)
    np.savez(key_out,
             KEYQ_OMPL=kq_ompl,
             KEYQ=kq,
             _STAT_KEYS=stat_keys,
             _STAT_VALUES=stat_values,
             _TOP_K_PER_TRAJ=top_k_indices)
    util.ack('[pickup_key_configuration] key configurations are written to {}'.
             format(key_out))
    def __iter__(self):
        envir = self.envir
        sess = self.sess
        advcore = self.advcore
        for _ in self.sample_generator():
            q = uw_random.random_state(0.75)
            envir.qstate = q
            if False:
                for i in range(len(self.unit_tunnel_v)):
                    envir.qstate = self.unit_tunnel_v[i]
                    yield envir.vstate[0][self.gview]
            DEBUG2 = False
            if DEBUG2:
                assert self.args.samplein, "DEBUG2=True requires --samplein"
                fn = '{}/{}.npz'.format(self.args.samplein, self.sample_index)
                print(fn)
                self.sample_index += 1
                d = np.load(fn)
                qs = d['QS']
                dqs = d['DQS']
                closes = d['CLOSES']
                for q, dq, close in zip(qs, dqs, closes):
                    #tr,aa,_ = pyosr.differential(q, close)
                    #dq = np.concatenate([tr, aa], axis=-1)
                    app = pyosr.apply(q, dq[:3], dq[3:])
                    print("q {}\ndq {}\nclose {}\napplied {}".format(
                        q, dq, close, app))
                    print("diff {}\n".format(dq))
                    assert pyosr.distance(close, pyosr.apply(
                        q, dq[:3], dq[3:])) < 1e-3
                    # continue
                    envir.qstate = q
                    yield envir.vstate[0][self.gview]
                    input("########## PRESS ENTER FOR APPLIED DQ ##########")
                    envir.qstate = app
                    yield envir.vstate[0][self.gview]
                    input("########## PRESS ENTER FOR CLOSEST Q ##########")
                    envir.qstate = close
                    yield envir.vstate[0][self.gview]
                    input("######## PRESS ENTER FOR THE NEXT TUPLE ########")
                continue  # Skip the remaining
            DEBUG3 = False
            if DEBUG3:
                for q in self.unit_tunnel_v:
                    envir.qstate = q
                    yield envir.vstate[0][self.gview]
                continue
            DEBUG = False
            if DEBUG:
                yield envir.vstate[0][self.gview]
                distances = pyosr.multi_distance(q, self.unit_tunnel_v)
                ni = np.argmin(distances)
                close = self.unit_tunnel_v[ni]
                envir.qstate = close
                yield envir.vstate[0][self.gview]
                continue
            assert self.args.sampleout, '--sampleout is required for checking CR results'
            for i in range(3):
                vstate = envir.vstate
                yield vstate[0][self.gview]
                if envir.r.is_valid_state(envir.qstate):
                    print("===== Current State is Valid ====")
                    if i == 1:  # Only capture the state after first movement
                        self.cr_noneed += 1
                else:
                    print("!!!!! Current State is not Valid ====")
                    print("!DEBUG UNIT: {}".format(envir.qstate))
                    nustate = envir.r.translate_from_unit_state(envir.qstate)
                    print("!DEBUG ORIGINAL: {}".format(nustate))
                    SEARCH_NUMERICAL = False
                    SEARCH_SA_FORCE = True
                    if SEARCH_NUMERICAL and i == 1:  # Only capture the state after first movement
                        early_exit = False
                        current = envir.qstate
                        self.to_cr_states.append(current)
                        for k in range(32):
                            nexts = []
                            for tr, aa in zip(self.delta_tr, self.delta_aa):
                                nexts.append(pyosr.apply(current, tr, aa))
                            nexts = np.array(nexts)
                            SAs = envir.r.intersection_region_surface_areas(
                                nexts, True)
                            print(SAs)
                            min_index = np.argmin(SAs)
                            current = nexts[min_index]
                            envir.qstate = current
                            vstate = envir.vstate
                            yield vstate[0][self.gview]
                            if envir.r.is_valid_state(envir.qstate):
                                early_exit = True
                                break
                        if early_exit:
                            self.cr_success += 1
                            self.cr_state_indicators.append(1)
                        else:
                            self.cr_fail += 1
                            self.cr_state_indicators.append(-1)
                        self.cr_states.append(envir.qstate)
                        self.sample_index += 1
                    if SEARCH_SA_FORCE and i == 1:  # Force from optimizing surface area
                        # TODO: Unfinished, just print the force mag and direction
                        tup = envir.r.force_from_intersecting_surface_area(
                            envir.qstate)
                        print('Force Pos:\n{}'.format(tup[0]))
                        print('Force Direction:\n{}'.format(tup[1]))
                        print('Force Mag:\n{}'.format(tup[2]))
                dic = {
                    advcore.rgb_tensor: [vstate[0]],
                    advcore.dep_tensor: [vstate[1]],
                }
                disp_pred = sess_no_hook(sess,
                                         advcore.finder_pred,
                                         feed_dict=dic)
                disp_pred = disp_pred[0][
                    0]  # first batch first view for [B,V,...]

                nstate = pyosr.apply(envir.qstate, disp_pred[:3],
                                     disp_pred[3:])
                distances = pyosr.multi_distance(q, self.unit_tunnel_v)
                ni = np.argmin(distances)
                close = self.unit_tunnel_v[ni]
                # nstate = envir.r.translate_to_unit_state(nstate_raw)
                print("Prediction {}".format(disp_pred))
                print("Next (Unit) {}".format(nstate))
                print("Error {}".format(pyosr.distance(nstate, close)))
                envir.qstate = nstate
                # input("########## PRESS ENTER TO CONTINUE ##########")
        print("CR No Need {}\nCR SUCCESS {}\nCR FAIL{}".format(
            self.cr_noneed, self.cr_success, self.cr_fail))
        np.savez(self.args.sampleout,
                 CRQS=self.cr_states,
                 TOCRQS=self.to_cr_states,
                 CRI=self.cr_state_indicators)