<jupyter_start><jupyter_text>for each concept: sample random demos from batch, create smaller flat batches, save as one pkl<jupyter_code>cuda6 = torch.device('cuda:6')<jupyter_output><empty_output><jupyter_text>concept 1: search for key for locked door<jupyter_code>"""
for demos that have unlock action:
    _color_ = get door color
    label timesteps [] to [pick up _color_ key]
"""

# save inds and concept inds, generate all memories, then slice to concept inds
for i in range(128):

    batch = []
    concept_inds = []

    # filter for 1280 demos that have unlock action
    count = 0
    while count < 1280:
        unlock_idx = None
        pickup_idx = None

        while unlock_idx is None or pickup_idx is None:
            rand_idx = 0
            dem = some_demos[rand_idx]
            dem_transformed = demo_list[rand_idx]

            # if demo has Toggle (unlock) action
            if env.actions.toggle in dem[3] and env.actions.pickup in dem[3]:
                # unlock_idx = dem[3].index(env.actions.toggle)
                # check if object is door
                # print(dem_transformed[unlock_idx][0]['image'][5][3][0])
                # while dem_transformed[unlock_idx][0]['image'][5][3][0] != 4:
                #     temp = unlock_idx
                #     print(temp)
                #     unlock_idx = dem[3].index(env.actions.toggle, temp)
                color = None
                for t in range(len(dem_transformed)):
                    isdoor = (dem_transformed[t][0]['image'][5][3][0] == 4)
                    istoggle = (dem[3][t] == env.actions.toggle)
                    if isdoor and istoggle:
                        unlock_idx = t
                        color = dem_transformed[unlock_idx][0]['image'][5][3][1]
                    
                    if color:
                        iskey = (dem_transformed[t][0]['image'][5][3][0] == 5)
                        iscolor = (dem_transformed[t][0]['image'][5][3][1] == c)
                        if iskey and istoggle:
                            pickup_idx = t
                            


        print('exit')
        # get door color
        c = dem_transformed[unlock_idx][0]['image'][5][3][1]
        # get pick up _c_ key idx
        pickup_idx = dem[3].index(env.actions.pickup)
        # check if object is key
        assert dem_transformed[pickup_idx][0]['image'][5][3][0] == 5
        assert dem_transformed[pickup_idx][0]['image'][5][3][1] == c

        # (start, end) idx
        batch.append(dem_transformed)
        concept_inds.append((pickup_idx, unlock_idx))

        count += 1

    batch.sort(key=len, reverse=True)
    # Constructing flat batch and indices pointing to start of each demonstration
    flat_batch = []
    inds = [0]
    flat_concept_inds = [(0,0)]

    for demo in batch:
        flat_batch += demo
        inds.append(inds[-1] + len(demo))
        flat_concept_inds.append((inds[-1] + pickup_idx, inds[-1] + unlock_idx))

    flat_batch = np.array(flat_batch)
    inds = inds[:-1]
    flat_concept_inds = flat_concept_inds[:-1]
    num_frames = len(flat_batch)

    assert len(flat_concept_inds) == len(inds) 

    mask = np.ones([len(flat_batch)], dtype=np.float64)
    mask[inds] = 0
    mask = torch.tensor(mask, device=cuda6, dtype=torch.float).unsqueeze(1)

    # get batch with only concepts
    flat_concept_batch = []
    concept_inds = [0]
    for pair in flat_concept_inds:
        concept_data = flat_batch[pair[0]:pair[1]]
        flat_concept_batch += concept_data
        concept_inds.append(concept_inds[-1] + (pair[1] - pair[0]))

    concept_mask = np.ones([len(flat_concept_batch)], dtype=np.float64)
    concept_mask[concept_inds] = 0
    concept_mask = torch.tensor(concept_mask, device=cuda6, dtype=torch.float).unsqueeze(1)

    # Observations, true action, values and done for each of the stored demostration
    obss, action_true, done = flat_batch[:, 0], flat_batch[:, 1], flat_batch[:, 2]

    episode_ids = np.zeros(len(flat_batch))
    inds_copy = inds.copy()

    # Loop terminates when every observation in the flat_batch has been handled
    while True:
        # taking observations and done located at inds
        done_step = done[inds]
        episode_ids[inds] = range(len(inds))

        # Updating inds, by removing those indices corresponding to which the demonstrations have finished
        inds = inds[:len(inds) - sum(done_step)]
        if len(inds) == 0:
            break

        # Incrementing the remaining indices
        inds = [index + 1 for index in inds]

    # make concept_episode_ids
    obss, action_true, done = flat_concept_batch[:, 0], flat_concept_batch[:, 1], flat_concept_batch[:, 2]

    concept_episode_ids = np.zeros(len(flat_concept_batch))
    concept_inds_copy = concept_inds.copy()

    # Loop terminates when every observation in the flat_batch has been handled
    while True:
        # taking observations and done located at inds
        done_step = done[concept_inds]
        concept_episode_ids[concept_inds] = range(len(concept_inds))

        # Updating inds, by removing those indices corresponding to which the demonstrations have finished
        concept_inds = concept_inds[:len(concept_inds) - sum(done_step)]
        if len(concept_inds) == 0:
            break

        # Incrementing the remaining indices
        concept_inds = [index + 1 for index in concept_inds]

    # {flat_batch, mask, episode_ids, inds_copy, 
    #   flat_batch_concept_inds, 
    #   concept_mask, concept_episode_ids, concept_inds} <-- inds after batch only has concepts
    batch_dict_list = {'flat_batch': flat_batch,
                        'mask': mask,
                        'episode_ids': episode_ids,
                        'inds': inds_copy,
                        'flat_batch_concept_inds': flat_concept_inds,
                        'concept_mask': concept_mask,
                        'concept_episode_ids': concept_episode_ids,
                        'concept_inds': concept_inds}

    path = '/data/graceduansu/concepts/1_search_for_key/batch' + str(i)
    with open(path, 'wb') as handle:
        pickle.dump(batch_dict_list, handle, protocol=pickle.HIGHEST_PROTOCOL)

### concept 2: take key to locked door
### concept 3: search for target object (default subgoal)<jupyter_output><empty_output><jupyter_text>testing stuff<jupyter_code>flat_batch = np.arange(200)
# batch_size = 50, 
demo_length = 40
memory_size = 25
num_frames = len(flat_batch)
recurrence = 5

inds = np.arange(0, len(flat_batch), demo_length)


memories = torch.zeros([len(flat_batch), memory_size])
episode_ids = np.zeros(len(flat_batch))
memory = torch.zeros([50, memory_size])

mask = np.ones([len(flat_batch)], dtype=np.float64)
mask[inds] = 0
mask = torch.tensor(mask, dtype=torch.float).unsqueeze(1)
# print(mask)
while True:
    # taking observations and done located at inds
    done_step = np.arange(0, len(flat_batch), demo_length)[1:]
    new_memory = np.arange(1000, 1025)
    new_memory = torch.tensor(new_memory)
    memories[inds, :] = memory[:len(inds), :]
    memory[:len(inds), :] = new_memory
    episode_ids[inds] = range(len(inds))

    # Updating inds, by removing those indices corresponding to which the demonstrations have finished
    inds = inds[:len(inds) - sum(done_step)]
    if len(inds) == 0:
        break

    # Incrementing the remaining indices
    inds = [index + 1 for index in inds]
print(episode_ids)
indexes = np.arange(0, num_frames, recurrence)
mask_step = mask[indexes]
print(len(mask_step))
print(len(mask))
len(indexes)
mask.shape
mask_step.shape
print(memory)
memory.shape
print(memories.shape)
memory = memories[indexes]
print(memory.shape)
memory = torch.rand((40,25))
# print(memory)
res = memory * mask_step
print(res.size())
# print(res)
indexes += 1
print(indexes)
mask_step = mask[indexes]<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1119440914866582, 'fraction_numerical': 0.017153748411689963, 'mean_word_length': 2.7445290199809707, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '236599', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 2650, 'n_tokens_neox': 2490, 'n_words': 593, 'path': 'generate_concepts_dataset.ipynb', 'repo_name': 'graceduansu/babyai', 'size': '7870'}