| |
| |
| |
| |
| @@ -21,6 +21,21 @@ conda activate jat |
| pip install -e .[dev] |
| ``` |
| |
| +## REGENT fork of sample-factory: Installation |
| +Following [this install ink](https://www.samplefactory.dev/01-get-started/installation/) but for the fork: |
| +```shell |
| +git clone https://github.com/kaustubhsridhar/sample-factory.git |
| +cd sample-factory |
| +pip install -e .[dev,mujoco,atari,envpool,vizdoom] |
| +``` |
| + |
| +# Regent fork of sample-factory: Train Unseen Env Policies and Generate Datasets |
| +Train policies using envpool's atari: |
| +```shell |
| +bash scripts_sample-factory/train_unseen_atari.sh |
| +``` |
| +Note that the training command inside the above script was obtained from the config files of Ed Beeching's Atari 57 models on Huggingface. An example is [here](https://huggingface.co/edbeeching/atari_2B_atari_mspacman_1111/blob/main/cfg.json#L124). See my discussion [here](https://huggingface.co/edbeeching/atari_2B_atari_mspacman_1111/discussions/2). |
| + |
| ## PREV Installation |
| |
| To get started with JAT, follow these steps: |
| @@ -155,12 +170,21 @@ python -u scripts_jat_regent/eval_RandP.py --task ${TASK} &> outputs/RandP/${TAS |
| ``` |
| |
| ### REGENT Analyze data |
| +Necessary: |
| ```shell |
| -python -u examples_regent/compare_datasets.py &> examples_regent/compare_datasets.txt & |
| - |
| python -u examples_regent/analyze_rows_tokenized.py &> examples_regent/analyze_rows_tokenized.txt & |
| +``` |
| |
| +Already ran and output dict in code: |
| +```shell |
| python -u examples_regent/get_dim_all_vector_tasks.py &> examples_regent/get_dim_all_vector_tasks.txt & |
| + |
| +python -u examples_regent/count_rows_to_consider.py &> examples_regent/count_rows_to_consider.txt & |
| +``` |
| + |
| +Optional: |
| +```shell |
| +python -u examples_regent/compare_datasets.py &> examples_regent/compare_datasets.txt & |
| ``` |
| |
| ## PREV Dataset |
| |
| deleted file mode 100644 |
| |
| |
| |
| @@ -1,38 +0,0 @@ |
| -import warnings |
| -from dataclasses import dataclass |
| -from typing import List, Optional, Tuple, Union |
| - |
| -import numpy as np |
| -import torch |
| -import torch.nn.functional as F |
| -from gymnasium import spaces |
| -from torch import BoolTensor, FloatTensor, LongTensor, Tensor, nn |
| -from transformers import GPTNeoModel, GPTNeoPreTrainedModel |
| -from transformers.modeling_outputs import ModelOutput |
| -from transformers.models.vit.modeling_vit import ViTPatchEmbeddings |
| - |
| -from jat.configuration_jat import JatConfig |
| -from jat.processing_jat import JatProcessor |
| - |
| - |
| -class RandP(): |
| - def __init__(self, dataset) -> None: |
| - self.steps = 0 |
| - # create an index for retrieval in vector obs envs (OR) collect all images in Atari |
| - |
| - def reset_rl(self): |
| - self.steps = 0 |
| - |
| - def get_next_action( |
| - self, |
| - processor: JatProcessor, |
| - continuous_observation: Optional[List[float]] = None, |
| - discrete_observation: Optional[List[int]] = None, |
| - text_observation: Optional[str] = None, |
| - image_observation: Optional[np.ndarray] = None, |
| - action_space: Union[spaces.Box, spaces.Discrete] = None, |
| - reward: Optional[float] = None, |
| - deterministic: bool = False, |
| - context_window: Optional[int] = None, |
| - ): |
| - pass |
| \ No newline at end of file |
| |
| deleted file mode 100644 |
| |
| |
| |
| |
| |
| @@ -8,23 +8,35 @@ from tqdm import tqdm |
| from autofaiss import build_index |
| |
| |
| +UNSEEN_TASK_NAMES = { # Total -- atari: 57, metaworld: 50, babyai: 39, mujoco: 11 |
| + |
| +} |
| + |
| def myprint(str): |
| - # check if first character of string is a newline character |
| - if str[0] == '\n': |
| - str_without_newline = str[1:] |
| + # check if first characters of string are newline character |
| + num_newlines = 0 |
| + while str[num_newlines] == '\n': |
| print() |
| - else: |
| - str_without_newline = str |
| + num_newlines += 1 |
| + str_without_newline = str[num_newlines:] |
| print(f'{datetime.now().strftime("%Y-%m-%d %H:%M:%S")}: {str_without_newline}') |
| |
| def is_png_img(item): |
| return isinstance(item, PngImagePlugin.PngImageFile) |
| |
| +def get_last_row_for_1M_states(task): |
| + last_row_idx = {'atari-alien': 14134, 'atari-amidar': 14319, 'atari-assault': 14427, 'atari-asterix': 14456, 'atari-asteroids': 14348, 'atari-atlantis': 14325, 'atari-bankheist': 14167, 'atari-battlezone': 13981, 'atari-beamrider': 13442, 'atari-berzerk': 13534, 'atari-bowling': 14110, 'atari-boxing': 14542, 'atari-breakout': 13474, 'atari-centipede': 14196, 'atari-choppercommand': 13397, 'atari-crazyclimber': 14026, 'atari-defender': 13504, 'atari-demonattack': 13499, 'atari-doubledunk': 14292, 'atari-enduro': 13260, 'atari-fishingderby': 14073, 'atari-freeway': 14016, 'atari-frostbite': 14075, 'atari-gopher': 13143, 'atari-gravitar': 14405, 'atari-hero': 14044, 'atari-icehockey': 14017, 'atari-jamesbond': 12678, 'atari-kangaroo': 14248, 'atari-krull': 14204, 'atari-kungfumaster': 14030, 'atari-montezumarevenge': 14219, 'atari-mspacman': 14120, 'atari-namethisgame': 13575, 'atari-phoenix': 13539, 'atari-pitfall': 14287, 'atari-pong': 14151, 'atari-privateeye': 14105, 'atari-qbert': 14026, 'atari-riverraid': 14275, 'atari-roadrunner': 14127, 'atari-robotank': 14079, 'atari-seaquest': 14097, 'atari-skiing': 14708, 'atari-solaris': 14199, 'atari-spaceinvaders': 12652, 'atari-stargunner': 13822, 'atari-surround': 13840, 'atari-tennis': 14062, 'atari-timepilot': 13896, 'atari-tutankham': 13121, 'atari-upndown': 13504, 'atari-venture': 14260, 'atari-videopinball': 14272, 'atari-wizardofwor': 13920, 'atari-yarsrevenge': 13981, 'atari-zaxxon': 13833, 'babyai-action-obj-door': 95000, 'babyai-blocked-unlock-pickup': 29279, 'babyai-boss-level-no-unlock': 12087, 'babyai-boss-level': 12101, 'babyai-find-obj-s5': 32974, 'babyai-go-to-door': 95000, 'babyai-go-to-imp-unlock': 9286, 'babyai-go-to-local': 95000, 'babyai-go-to-obj-door': 95000, 'babyai-go-to-obj': 95000, 'babyai-go-to-red-ball-grey': 95000, 'babyai-go-to-red-ball-no-dists': 95000, 'babyai-go-to-red-ball': 95000, 'babyai-go-to-red-blue-ball': 95000, 'babyai-go-to-seq': 13744, 'babyai-go-to': 18974, 'babyai-key-corridor': 9014, 'babyai-mini-boss-level': 38119, 'babyai-move-two-across-s8n9': 24505, 'babyai-one-room-s8': 95000, 'babyai-open-door': 95000, 'babyai-open-doors-order-n4': 95000, 'babyai-open-red-door': 95000, 'babyai-open-two-doors': 73291, 'babyai-open': 32559, 'babyai-pickup-above': 34084, 'babyai-pickup-dist': 89640, 'babyai-pickup-loc': 95000, 'babyai-pickup': 18670, 'babyai-put-next-local': 83187, 'babyai-put-next': 56986, 'babyai-synth-loc': 21605, 'babyai-synth-seq': 13049, 'babyai-synth': 19409, 'babyai-unblock-pickup': 17881, 'babyai-unlock-local': 71186, 'babyai-unlock-pickup': 50883, 'babyai-unlock-to-unlock': 23062, 'babyai-unlock': 11734, 'metaworld-assembly': 10000, 'metaworld-basketball': 10000, 'metaworld-bin-picking': 10000, 'metaworld-box-close': 10000, 'metaworld-button-press-topdown-wall': 10000, 'metaworld-button-press-topdown': 10000, 'metaworld-button-press-wall': 10000, 'metaworld-button-press': 10000, 'metaworld-coffee-button': 10000, 'metaworld-coffee-pull': 10000, 'metaworld-coffee-push': 10000, 'metaworld-dial-turn': 10000, 'metaworld-disassemble': 10000, 'metaworld-door-close': 10000, 'metaworld-door-lock': 10000, 'metaworld-door-open': 10000, 'metaworld-door-unlock': 10000, 'metaworld-drawer-close': 10000, 'metaworld-drawer-open': 10000, 'metaworld-faucet-close': 10000, 'metaworld-faucet-open': 10000, 'metaworld-hammer': 10000, 'metaworld-hand-insert': 10000, 'metaworld-handle-press-side': 10000, 'metaworld-handle-press': 10000, 'metaworld-handle-pull-side': 10000, 'metaworld-handle-pull': 10000, 'metaworld-lever-pull': 10000, 'metaworld-peg-insert-side': 10000, 'metaworld-peg-unplug-side': 10000, 'metaworld-pick-out-of-hole': 10000, 'metaworld-pick-place-wall': 10000, 'metaworld-pick-place': 10000, 'metaworld-plate-slide-back-side': 10000, 'metaworld-plate-slide-back': 10000, 'metaworld-plate-slide-side': 10000, 'metaworld-plate-slide': 10000, 'metaworld-push-back': 10000, 'metaworld-push-wall': 10000, 'metaworld-push': 10000, 'metaworld-reach-wall': 10000, 'metaworld-reach': 10000, 'metaworld-shelf-place': 10000, 'metaworld-soccer': 10000, 'metaworld-stick-pull': 10000, 'metaworld-stick-push': 10000, 'metaworld-sweep-into': 10000, 'metaworld-sweep': 10000, 'metaworld-window-close': 10000, 'metaworld-window-open': 10000, 'mujoco-ant': 4023, 'mujoco-doublependulum': 4002, 'mujoco-halfcheetah': 4000, 'mujoco-hopper': 4931, 'mujoco-humanoid': 4119, 'mujoco-pendulum': 4959, 'mujoco-pusher': 9000, 'mujoco-reacher': 9000, 'mujoco-standup': 4000, 'mujoco-swimmer': 4000, 'mujoco-walker': 4101} |
| + return last_row_idx[task] |
| + |
| +def get_last_row_for_100k_states(task): |
| + last_row_idx = {'atari-alien': 3135, 'atari-amidar': 3142, 'atari-assault': 3132, 'atari-asterix': 3181, 'atari-asteroids': 3127, 'atari-atlantis': 3128, 'atari-bankheist': 3156, 'atari-battlezone': 3136, 'atari-beamrider': 3131, 'atari-berzerk': 3127, 'atari-bowling': 3148, 'atari-boxing': 3227, 'atari-breakout': 3128, 'atari-centipede': 3176, 'atari-choppercommand': 3144, 'atari-crazyclimber': 3134, 'atari-defender': 3127, 'atari-demonattack': 3127, 'atari-doubledunk': 3175, 'atari-enduro': 3126, 'atari-fishingderby': 3155, 'atari-freeway': 3131, 'atari-frostbite': 3146, 'atari-gopher': 3128, 'atari-gravitar': 3202, 'atari-hero': 3144, 'atari-icehockey': 3138, 'atari-jamesbond': 3131, 'atari-kangaroo': 3160, 'atari-krull': 3162, 'atari-kungfumaster': 3143, 'atari-montezumarevenge': 3168, 'atari-mspacman': 3143, 'atari-namethisgame': 3131, 'atari-phoenix': 3127, 'atari-pitfall': 3131, 'atari-pong': 3160, 'atari-privateeye': 3158, 'atari-qbert': 3136, 'atari-riverraid': 3157, 'atari-roadrunner': 3150, 'atari-robotank': 3133, 'atari-seaquest': 3138, 'atari-skiing': 3271, 'atari-solaris': 3129, 'atari-spaceinvaders': 3128, 'atari-stargunner': 3129, 'atari-surround': 3143, 'atari-tennis': 3129, 'atari-timepilot': 3132, 'atari-tutankham': 3127, 'atari-upndown': 3127, 'atari-venture': 3148, 'atari-videopinball': 3130, 'atari-wizardofwor': 3138, 'atari-yarsrevenge': 3129, 'atari-zaxxon': 3133, 'babyai-action-obj-door': 15923, 'babyai-blocked-unlock-pickup': 2919, 'babyai-boss-level-no-unlock': 1217, 'babyai-boss-level': 1159, 'babyai-find-obj-s5': 3345, 'babyai-go-to-door': 18875, 'babyai-go-to-imp-unlock': 923, 'babyai-go-to-local': 18724, 'babyai-go-to-obj-door': 16472, 'babyai-go-to-obj': 20197, 'babyai-go-to-red-ball-grey': 16953, 'babyai-go-to-red-ball-no-dists': 20165, 'babyai-go-to-red-ball': 18730, 'babyai-go-to-red-blue-ball': 16934, 'babyai-go-to-seq': 1439, 'babyai-go-to': 1964, 'babyai-key-corridor': 900, 'babyai-mini-boss-level': 3789, 'babyai-move-two-across-s8n9': 2462, 'babyai-one-room-s8': 16994, 'babyai-open-door': 13565, 'babyai-open-doors-order-n4': 9706, 'babyai-open-red-door': 21185, 'babyai-open-two-doors': 7348, 'babyai-open': 3331, 'babyai-pickup-above': 3392, 'babyai-pickup-dist': 19693, 'babyai-pickup-loc': 16405, 'babyai-pickup': 1806, 'babyai-put-next-local': 8303, 'babyai-put-next': 5703, 'babyai-synth-loc': 2183, 'babyai-synth-seq': 1316, 'babyai-synth': 1964, 'babyai-unblock-pickup': 1886, 'babyai-unlock-local': 7118, 'babyai-unlock-pickup': 5107, 'babyai-unlock-to-unlock': 2309, 'babyai-unlock': 1177, 'metaworld-assembly': 1000, 'metaworld-basketball': 1000, 'metaworld-bin-picking': 1000, 'metaworld-box-close': 1000, 'metaworld-button-press-topdown-wall': 1000, 'metaworld-button-press-topdown': 1000, 'metaworld-button-press-wall': 1000, 'metaworld-button-press': 1000, 'metaworld-coffee-button': 1000, 'metaworld-coffee-pull': 1000, 'metaworld-coffee-push': 1000, 'metaworld-dial-turn': 1000, 'metaworld-disassemble': 1000, 'metaworld-door-close': 1000, 'metaworld-door-lock': 1000, 'metaworld-door-open': 1000, 'metaworld-door-unlock': 1000, 'metaworld-drawer-close': 1000, 'metaworld-drawer-open': 1000, 'metaworld-faucet-close': 1000, 'metaworld-faucet-open': 1000, 'metaworld-hammer': 1000, 'metaworld-hand-insert': 1000, 'metaworld-handle-press-side': 1000, 'metaworld-handle-press': 1000, 'metaworld-handle-pull-side': 1000, 'metaworld-handle-pull': 1000, 'metaworld-lever-pull': 1000, 'metaworld-peg-insert-side': 1000, 'metaworld-peg-unplug-side': 1000, 'metaworld-pick-out-of-hole': 1000, 'metaworld-pick-place-wall': 1000, 'metaworld-pick-place': 1000, 'metaworld-plate-slide-back-side': 1000, 'metaworld-plate-slide-back': 1000, 'metaworld-plate-slide-side': 1000, 'metaworld-plate-slide': 1000, 'metaworld-push-back': 1000, 'metaworld-push-wall': 1000, 'metaworld-push': 1000, 'metaworld-reach-wall': 1000, 'metaworld-reach': 1000, 'metaworld-shelf-place': 1000, 'metaworld-soccer': 1000, 'metaworld-stick-pull': 1000, 'metaworld-stick-push': 1000, 'metaworld-sweep-into': 1000, 'metaworld-sweep': 1000, 'metaworld-window-close': 1000, 'metaworld-window-open': 1000, 'mujoco-ant': 401, 'mujoco-doublependulum': 401, 'mujoco-halfcheetah': 400, 'mujoco-hopper': 491, 'mujoco-humanoid': 415, 'mujoco-pendulum': 495, 'mujoco-pusher': 1000, 'mujoco-reacher': 2000, 'mujoco-standup': 400, 'mujoco-swimmer': 400, 'mujoco-walker': 407} |
| + return last_row_idx[task] |
| + |
| def get_obs_dim(task): |
| assert task.startswith("babyai") or task.startswith("metaworld") or task.startswith("mujoco") |
| |
| all_obs_dims={'babyai-action-obj-door': 212, 'babyai-blocked-unlock-pickup': 212, 'babyai-boss-level-no-unlock': 212, 'babyai-boss-level': 212, 'babyai-find-obj-s5': 212, 'babyai-go-to-door': 212, 'babyai-go-to-imp-unlock': 212, 'babyai-go-to-local': 212, 'babyai-go-to-obj-door': 212, 'babyai-go-to-obj': 212, 'babyai-go-to-red-ball-grey': 212, 'babyai-go-to-red-ball-no-dists': 212, 'babyai-go-to-red-ball': 212, 'babyai-go-to-red-blue-ball': 212, 'babyai-go-to-seq': 212, 'babyai-go-to': 212, 'babyai-key-corridor': 212, 'babyai-mini-boss-level': 212, 'babyai-move-two-across-s8n9': 212, 'babyai-one-room-s8': 212, 'babyai-open-door': 212, 'babyai-open-doors-order-n4': 212, 'babyai-open-red-door': 212, 'babyai-open-two-doors': 212, 'babyai-open': 212, 'babyai-pickup-above': 212, 'babyai-pickup-dist': 212, 'babyai-pickup-loc': 212, 'babyai-pickup': 212, 'babyai-put-next-local': 212, 'babyai-put-next': 212, 'babyai-synth-loc': 212, 'babyai-synth-seq': 212, 'babyai-synth': 212, 'babyai-unblock-pickup': 212, 'babyai-unlock-local': 212, 'babyai-unlock-pickup': 212, 'babyai-unlock-to-unlock': 212, 'babyai-unlock': 212, 'metaworld-assembly': 39, 'metaworld-basketball': 39, 'metaworld-bin-picking': 39, 'metaworld-box-close': 39, 'metaworld-button-press-topdown-wall': 39, 'metaworld-button-press-topdown': 39, 'metaworld-button-press-wall': 39, 'metaworld-button-press': 39, 'metaworld-coffee-button': 39, 'metaworld-coffee-pull': 39, 'metaworld-coffee-push': 39, 'metaworld-dial-turn': 39, 'metaworld-disassemble': 39, 'metaworld-door-close': 39, 'metaworld-door-lock': 39, 'metaworld-door-open': 39, 'metaworld-door-unlock': 39, 'metaworld-drawer-close': 39, 'metaworld-drawer-open': 39, 'metaworld-faucet-close': 39, 'metaworld-faucet-open': 39, 'metaworld-hammer': 39, 'metaworld-hand-insert': 39, 'metaworld-handle-press-side': 39, 'metaworld-handle-press': 39, 'metaworld-handle-pull-side': 39, 'metaworld-handle-pull': 39, 'metaworld-lever-pull': 39, 'metaworld-peg-insert-side': 39, 'metaworld-peg-unplug-side': 39, 'metaworld-pick-out-of-hole': 39, 'metaworld-pick-place-wall': 39, 'metaworld-pick-place': 39, 'metaworld-plate-slide-back-side': 39, 'metaworld-plate-slide-back': 39, 'metaworld-plate-slide-side': 39, 'metaworld-plate-slide': 39, 'metaworld-push-back': 39, 'metaworld-push-wall': 39, 'metaworld-push': 39, 'metaworld-reach-wall': 39, 'metaworld-reach': 39, 'metaworld-shelf-place': 39, 'metaworld-soccer': 39, 'metaworld-stick-pull': 39, 'metaworld-stick-push': 39, 'metaworld-sweep-into': 39, 'metaworld-sweep': 39, 'metaworld-window-close': 39, 'metaworld-window-open': 39, 'mujoco-ant': 27, 'mujoco-doublependulum': 11, 'mujoco-halfcheetah': 17, 'mujoco-hopper': 11, 'mujoco-humanoid': 376, 'mujoco-pendulum': 4, 'mujoco-pusher': 23, 'mujoco-reacher': 11, 'mujoco-standup': 376, 'mujoco-swimmer': 8, 'mujoco-walker': 17} |
| - return all_obs_dims[task] |
| + return (all_obs_dims[task],) |
| |
| def get_act_dim(task): |
| assert task.startswith("babyai") or task.startswith("metaworld") or task.startswith("mujoco") |
| @@ -36,141 +48,188 @@ def get_act_dim(task): |
| elif task.startswith("mujoco"): |
| all_act_dims={'mujoco-ant': 8, 'mujoco-doublependulum': 1, 'mujoco-halfcheetah': 6, 'mujoco-hopper': 3, 'mujoco-humanoid': 17, 'mujoco-pendulum': 1, 'mujoco-pusher': 7, 'mujoco-reacher': 2, 'mujoco-standup': 17, 'mujoco-swimmer': 2, 'mujoco-walker': 6} |
| return all_act_dims[task] |
| - |
| -def process_row_atari(attn_mask, row_of_obs, task): |
| - """ |
| - Example for selection with bools: |
| - >>> a = np.array([0,1,2,3,4,5]) |
| - >>> b = np.array([1,0,0,0,0,1]).astype(bool) |
| - >>> a[b] |
| - array([0, 5]) |
| - """ |
| - attn_mask = np.array(attn_mask).astype(bool) |
| |
| - row_of_obs = torch.stack([to_tensor(np.array(img)) for img in row_of_obs]) |
| - row_of_obs = row_of_obs[attn_mask] |
| +def get_task_info(task): |
| + rew_key = 'rewards' |
| + attn_key = 'attention_mask' |
| + if task.startswith("atari"): |
| + obs_key = 'image_observations' |
| + act_key = 'discrete_actions' |
| + B = 32 # half of 54 |
| + obs_dim = (3, 4*84, 84) |
| + elif task.startswith("babyai"): |
| + obs_key = 'discrete_observations' # also has 'text_observations' only for raw dataset not for tokenized dataset (as it is combined into discrete_observation in tokenized dataset) |
| + act_key = 'discrete_actions' |
| + B = 256 # half of 512 |
| + obs_dim = get_obs_dim(task) |
| + elif task.startswith("metaworld") or task.startswith("mujoco"): |
| + obs_key = 'continuous_observations' |
| + act_key = 'continuous_actions' |
| + B = 256 |
| + obs_dim = get_obs_dim(task) |
| + |
| + return rew_key, attn_key, obs_key, act_key, B, obs_dim |
| + |
| +def process_row_of_obs_atari_full_without_mask(row_of_obs): |
| + |
| + if not isinstance(row_of_obs, torch.Tensor): |
| + row_of_obs = torch.stack([to_tensor(np.array(img)) for img in row_of_obs]) |
| row_of_obs = row_of_obs * 0.5 + 0.5 # denormalize from [-1, 1] to [0, 1] |
| - assert row_of_obs.shape == (sum(attn_mask), 84, 4, 84) |
| + assert row_of_obs.shape == (len(row_of_obs), 84, 4, 84) |
| row_of_obs = row_of_obs.permute(0, 2, 1, 3) # (*, 4, 84, 84) |
| - row_of_obs = row_of_obs.reshape(sum(attn_mask), 4*84, 84) # put side-by-side |
| + row_of_obs = row_of_obs.reshape(len(row_of_obs), 4*84, 84) # put side-by-side |
| row_of_obs = row_of_obs.unsqueeze(1).repeat(1, 3, 1, 1) # repeat for 3 channels |
| - assert row_of_obs.shape == (sum(attn_mask), 3, 4*84, 84) # sum(attn_mask) is the batch size dimension |
| - |
| - return attn_mask, row_of_obs |
| + assert row_of_obs.shape == (len(row_of_obs), 3, 4*84, 84) # sum(attn_mask) is the batch size dimension |
| + |
| + return row_of_obs |
| |
| -def process_row_vector(attn_mask, row_of_obs, task, return_numpy=False): |
| - attn_mask = np.array(attn_mask).astype(bool) |
| +def collect_all_atari_data(dataset, all_row_idxs=None): |
| + if all_row_idxs is None: |
| + all_row_idxs = list(range(len(dataset['train']))) |
| |
| - row_of_obs = np.array(row_of_obs) |
| - if not return_numpy: |
| - row_of_obs = torch.tensor(row_of_obs) |
| - row_of_obs = row_of_obs[attn_mask] |
| - assert row_of_obs.shape == (sum(attn_mask), get_obs_dim(task)) |
| - |
| - return attn_mask, row_of_obs |
| - |
| -def retrieve_atari(row_of_obs, # query: (row_B, 3, 4*84, 84) |
| - dataset, # to retrieve from |
| - all_rows_to_consider, # rows to consider |
| - num_to_retrieve, # top-k |
| + all_rows_of_obs = [] |
| + all_attn_masks = [] |
| + for row_idx in tqdm(all_row_idxs): |
| + datarow = dataset['train'][row_idx] |
| + row_of_obs = process_row_of_obs_atari_full_without_mask(datarow['image_observations']) |
| + attn_mask = np.array(datarow['attention_mask']).astype(bool) |
| + all_rows_of_obs.append(row_of_obs) # appending tensor |
| + all_attn_masks.append(attn_mask) # appending np array |
| + all_rows_of_obs = torch.stack(all_rows_of_obs, dim=0) # stacking tensors |
| + all_attn_masks = np.stack(all_attn_masks, axis=0) # concatenating np arrays |
| + assert (all_rows_of_obs.shape == (len(all_row_idxs), 32, 3, 4*84, 84) and |
| + all_attn_masks.shape == (len(all_row_idxs), 32)) |
| + return all_attn_masks, all_rows_of_obs |
| + |
| +def collect_all_data(dataset, task, obs_key): |
| + last_row_idx = get_last_row_for_100k_states(task) |
| + all_row_idxs = list(range(last_row_idx)) |
| + if task.startswith("atari"): |
| + myprint("Collecting all Atari images and Atari attention masks...") |
| + all_attn_masks_OG, all_rows_of_obs_OG = collect_all_atari_data(dataset, all_row_idxs) |
| + else: |
| + datarows = dataset['train'][all_row_idxs] |
| + all_rows_of_obs_OG = np.array(datarows[obs_key]) |
| + all_attn_masks_OG = np.array(datarows['attention_mask']).astype(bool) |
| + return all_rows_of_obs_OG, all_attn_masks_OG, all_row_idxs |
| + |
| +def collect_subset(all_rows_of_obs_OG, |
| + all_attn_masks_OG, |
| + all_rows_to_consider, |
| + kwargs |
| + ): |
| + """ |
| + Function to collect subset of data given all_rows_to_consider, reshape it, create all_indices and return. |
| + Used in both retrieve_atari() and retrieve_vector() --> build_index_vector(). |
| + """ |
| + myprint(f'\n\n\n' + ('-'*100) + f'Collecting subset...') |
| + # read kwargs |
| + B, task, obs_dim = kwargs['B'], kwargs['task'], kwargs['obs_dim'] |
| + |
| + # take subset based on all_rows_to_consider |
| + myprint(f'Taking subset of data based on all_rows_to_consider...') |
| + all_processed_rows_of_obs = all_rows_of_obs_OG[all_rows_to_consider] |
| + all_attn_masks = all_attn_masks_OG[all_rows_to_consider] |
| + assert (all_processed_rows_of_obs.shape == (len(all_rows_to_consider), B, *obs_dim) and |
| + all_attn_masks.shape == (len(all_rows_to_consider), B)) |
| + |
| + # reshape |
| + myprint(f'Reshaping data...') |
| + all_attn_masks = all_attn_masks.reshape(-1) |
| + all_processed_rows_of_obs = all_processed_rows_of_obs.reshape(-1, *obs_dim) |
| + all_processed_rows_of_obs = all_processed_rows_of_obs[all_attn_masks] |
| + assert (all_attn_masks.shape == (len(all_rows_to_consider) * B,) and |
| + all_processed_rows_of_obs.shape == (np.sum(all_attn_masks), *obs_dim)) |
| + |
| + # collect indices of data |
| + myprint(f'Collecting indices of data...') |
| + all_indices = np.array([[row_idx, i] for row_idx in all_rows_to_consider for i in range(B)]) |
| + all_indices = all_indices[all_attn_masks] # this is fine because all attn masks have 0s that only come after 1s |
| + assert all_indices.shape == (np.sum(all_attn_masks), 2) |
| + |
| + myprint(f'{all_indices.shape=}, {all_processed_rows_of_obs.shape=}') |
| + myprint(('-'*100) + '\n\n\n') |
| + return all_indices, all_processed_rows_of_obs |
| + |
| +def retrieve_atari(row_of_obs, # query: (xbdim, 3, 4*84, 84) / (xdim *obs_dim) |
| + all_processed_rows_of_obs, |
| + all_indices, |
| + num_to_retrieve, |
| kwargs |
| - ): |
| + ): |
| + """ |
| + Retrieval for Atari with images, ssim distance, and on GPU. |
| + """ |
| assert isinstance(row_of_obs, torch.Tensor) |
| |
| # read kwargs # Note: B = len of row |
| - B, attn_key, obs_key, device, task, batch_size_retrieval = kwargs['B'], kwargs['attn_key'], kwargs['obs_key'], kwargs['device'], kwargs['task'], kwargs['batch_size_retrieval'] |
| + B, device, batch_size_retrieval = kwargs['B'], kwargs['device'], kwargs['batch_size_retrieval'] |
| |
| # batch size of row_of_obs which can be <= B since we process before calling this function |
| - row_B = row_of_obs.shape[0] |
| - |
| + xbdim = row_of_obs.shape[0] |
| + |
| + # collect subset of data that we can retrieve from |
| + ydim = all_processed_rows_of_obs.shape[0] |
| + |
| # first argument for ssim |
| - repeated_row_og = row_of_obs.repeat_interleave(B, dim=0).to(device) |
| - assert repeated_row_og.shape == (row_B*B, 3, 4*84, 84) |
| + xbatch = row_of_obs.repeat_interleave(batch_size_retrieval, dim=0).to(device) |
| + assert xbatch.shape == (xbdim * batch_size_retrieval, 3, 4*84, 84) |
| |
| - # iterate over all other rows |
| + # iterate over data that we can retrieve from in batches |
| all_ssim = [] |
| - all_indices = [] |
| - total = 0 |
| - for other_row_idx in tqdm(all_rows_to_consider): |
| - other_attn_mask, other_row_of_obs = process_row_atari(dataset['train'][other_row_idx][attn_key], dataset['train'][other_row_idx][obs_key]) |
| - |
| - # batch size of other_row_of_obs |
| - other_row_B = other_row_of_obs.shape[0] |
| - total += other_row_B |
| - |
| - # first argument for ssim: RECHECK |
| - if other_row_B < B: # when other row has less observations than expected |
| - repeated_row = row_of_obs.repeat_interleave(other_row_B, dim=0).to(device) |
| - elif other_row_B == B: # otherwise just use the one created before the for loop |
| - repeated_row = repeated_row_og |
| - assert repeated_row.shape == (row_B*other_row_B, 3, 4*84, 84) |
| - |
| + for j in range(0, ydim, batch_size_retrieval): |
| # second argument for ssim |
| - repeated_other_row = other_row_of_obs.repeat(row_B, 1, 1, 1).to(device) |
| - assert repeated_other_row.shape == (row_B*other_row_B, 3, 4*84, 84) |
| + ybatch = all_processed_rows_of_obs[j:j+batch_size_retrieval] |
| + ybdim = ybatch.shape[0] |
| + ybatch = ybatch.repeat(xbdim, 1, 1, 1).to(device) |
| + assert ybatch.shape == (ybdim * xbdim, 3, 4*84, 84) |
| + |
| + if ybdim < batch_size_retrieval: # for last batch |
| + xbatch = row_of_obs.repeat_interleave(ybdim, dim=0).to(device) |
| + assert xbatch.shape == (xbdim * ybdim, 3, 4*84, 84) |
| |
| # compare via ssim and updated all_ssim |
| - ssim_score = ssim(repeated_row, repeated_other_row, data_range=1.0, size_average=False) |
| - ssim_score = ssim_score.reshape(row_B, other_row_B) |
| + ssim_score = ssim(xbatch, ybatch, data_range=1.0, size_average=False) |
| + ssim_score = ssim_score.reshape(xbdim, ybdim) |
| all_ssim.append(ssim_score) |
| |
| - # update all_indices |
| - all_indices.extend([[other_row_idx, i] for i in range(other_row_B)]) |
| - |
| # concat |
| all_ssim = torch.cat(all_ssim, dim=1) |
| - assert all_ssim.shape == (row_B, total) |
| + assert all_ssim.shape == (xbdim, ydim) |
| |
| - all_indices = np.array(all_indices) |
| - assert all_indices.shape == (total, 2) |
| + assert all_indices.shape == (ydim, 2) |
| |
| # get top-k indices |
| topk_values, topk_indices = torch.topk(all_ssim, num_to_retrieve, dim=1, largest=True) |
| topk_indices = topk_indices.cpu().numpy() |
| - assert topk_indices.shape == (row_B, num_to_retrieve) |
| + assert topk_indices.shape == (xbdim, num_to_retrieve) |
| |
| # convert topk indices to indices in the dataset |
| - retrieved_indices = np.array(all_indices[topk_indices]) |
| - assert retrieved_indices.shape == (row_B, num_to_retrieve, 2) |
| - |
| - # pad the above to expected B |
| - if row_B < B: |
| - retrieved_indices = np.concatenate([retrieved_indices, np.zeros((B-row_B, num_to_retrieve, 2), dtype=int)], axis=0) |
| - assert retrieved_indices.shape == (B, num_to_retrieve, 2) |
| + retrieved_indices = all_indices[topk_indices] |
| + assert retrieved_indices.shape == (xbdim, num_to_retrieve, 2) |
| |
| return retrieved_indices |
| |
| -def build_index_vector(all_rows_of_obs_og, |
| - all_attn_masks_og, |
| +def build_index_vector(all_rows_of_obs_OG, |
| + all_attn_masks_OG, |
| all_rows_to_consider, |
| kwargs |
| - ): |
| + ): |
| + """ |
| + Builds FAISS index for vector observation environments. |
| + """ |
| # read kwargs # Note: B = len of row |
| - B, attn_key, obs_key, device, task, batch_size_retrieval, nb_cores_autofaiss = kwargs['B'], kwargs['attn_key'], kwargs['obs_key'], kwargs['device'], kwargs['task'], kwargs['batch_size_retrieval'], kwargs['nb_cores_autofaiss'] |
| - obs_dim = get_obs_dim(task) |
| + nb_cores_autofaiss = kwargs['nb_cores_autofaiss'] |
| |
| - # take subset based on all_rows_to_consider |
| - myprint(f'Taking subset') |
| - all_rows_of_obs = all_rows_of_obs_og[all_rows_to_consider] |
| - all_attn_masks = all_attn_masks_og[all_rows_to_consider] |
| - assert (all_rows_of_obs.shape == (len(all_rows_to_consider), B, obs_dim) and |
| - all_attn_masks.shape == (len(all_rows_to_consider), B)) |
| - |
| - # reshape |
| - all_attn_masks = all_attn_masks.reshape(-1) |
| - all_rows_of_obs = all_rows_of_obs.reshape(-1, obs_dim) |
| - all_rows_of_obs = all_rows_of_obs[all_attn_masks] |
| - assert all_rows_of_obs.shape == (np.sum(all_attn_masks), obs_dim) |
| + # take subset based on all_rows_to_consider, reshape, and save indices of data |
| + all_indices, all_processed_rows_of_obs = collect_subset(all_rows_of_obs_OG, all_attn_masks_OG, all_rows_to_consider, kwargs) |
| |
| - # save indices of data to retrieve from |
| - myprint(f'Saving indices of data to retrieve from') |
| - all_indices = np.array([[row_idx, i] for row_idx in all_rows_to_consider for i in range(B)]) |
| - all_indices = all_indices[all_attn_masks] # this is fine because all attn masks have 0s that only come after 1s |
| - assert all_indices.shape == (np.sum(all_attn_masks), 2) |
| + # make sure input to build_index is float, otherwise you will get reading temp file error |
| + all_processed_rows_of_obs = all_processed_rows_of_obs.astype(float) |
| |
| # build index |
| - myprint(f'Building index...') |
| - knn_index, knn_index_infos = build_index(embeddings=all_rows_of_obs, # Note: embeddings have to be float to avoid errors in autofaiss / embedding_reader! |
| + myprint(('-'*100) + 'Building index...') |
| + knn_index, knn_index_infos = build_index(embeddings=all_processed_rows_of_obs, # Note: embeddings have to be float to avoid errors in autofaiss / embedding_reader! |
| save_on_disk=False, |
| min_nearest_neighbors_to_retrieve=20, # default: 20 |
| max_index_query_time_ms=10, # default: 10 |
| @@ -179,34 +238,32 @@ def build_index_vector(all_rows_of_obs_og, |
| metric_type='l2', |
| nb_cores=nb_cores_autofaiss, # default: None # "The number of cores to use, by default will use all cores" as seen in https://criteo.github.io/autofaiss/getting_started/quantization.html#the-build-index-command |
| ) |
| + myprint(('-'*100) + '\n\n\n') |
| |
| - return knn_index, all_indices |
| + return all_indices, knn_index |
| |
| -def retrieve_vector(row_of_obs, # query: (row_B, dim) |
| - dataset, # to retrieve from |
| - all_rows_to_consider, # rows to consider |
| - num_to_retrieve, # top-k |
| +def retrieve_vector(row_of_obs, # query: (xbdim, *obs_dim) |
| + knn_index, |
| + all_indices, |
| + num_to_retrieve, |
| kwargs |
| - ): |
| + ): |
| + """ |
| + Retrieval for vector observation environments. |
| + """ |
| assert isinstance(row_of_obs, np.ndarray) |
| |
| # read few kwargs |
| B = kwargs['B'] |
| |
| # batch size of row_of_obs which can be <= B since we process before calling this function |
| - row_B = row_of_obs.shape[0] |
| + xbdim = row_of_obs.shape[0] |
| |
| - # read dataset_tuple |
| - all_rows_of_obs, all_attn_masks = dataset |
| - |
| - # create index and all_indices |
| - knn_index, all_indices = build_index_vector(all_rows_of_obs, all_attn_masks, all_rows_to_consider, kwargs) |
| - |
| # retrieve |
| myprint(f'Retrieving...') |
| topk_indices, _ = knn_index.search(row_of_obs, 10 * num_to_retrieve) |
| topk_indices = topk_indices.astype(int) |
| - assert topk_indices.shape == (row_B, 10 * num_to_retrieve) |
| + assert topk_indices.shape == (xbdim, 10 * num_to_retrieve) |
| |
| # remove -1s and crop to num_to_retrieve |
| try: |
| @@ -219,16 +276,10 @@ def retrieve_vector(row_of_obs, # query: (row_B, dim) |
| print(f'-------------------------------------------------------------------------------------------------------------------------------------------') |
| print(f'Leaving some -1s in topk_indices and continuing') |
| topk_indices = np.array([indices[:num_to_retrieve] for indices in topk_indices]) |
| - assert topk_indices.shape == (row_B, num_to_retrieve) |
| + assert topk_indices.shape == (xbdim, num_to_retrieve) |
| |
| # convert topk indices to indices in the dataset |
| retrieved_indices = all_indices[topk_indices] |
| - assert retrieved_indices.shape == (row_B, num_to_retrieve, 2) |
| - |
| - # pad the above to expected B |
| - if row_B < B: |
| - retrieved_indices = np.concatenate([retrieved_indices, np.zeros((B-row_B, num_to_retrieve, 2), dtype=int)], axis=0) |
| - assert retrieved_indices.shape == (B, num_to_retrieve, 2) |
| + assert retrieved_indices.shape == (xbdim, num_to_retrieve, 2) |
| |
| - myprint(f'Returning') |
| return retrieved_indices |
| \ No newline at end of file |
| |
| |
| |
| |
| @@ -15,9 +15,10 @@ from transformers import AutoModelForCausalLM, AutoProcessor, HfArgumentParser |
| |
| from jat.eval.rl import TASK_NAME_TO_ENV_ID, make |
| from jat.utils import normalize, push_to_hub, save_video_grid |
| -from jat_regent.RandP import RandP |
| +from jat_regent.modeling_RandP import RandP |
| from datasets import load_from_disk |
| from datasets.config import HF_DATASETS_CACHE |
| +from jat_regent.utils import myprint |
| |
| |
| @dataclass |
| @@ -70,6 +71,7 @@ def eval_rl(model, processor, task, eval_args): |
| scores = [] |
| frames = [] |
| for episode in tqdm(range(eval_args.num_episodes), desc=task, unit="episode", leave=False): |
| + myprint(('-'*100) + f'{episode=}') |
| observation, _ = env.reset() |
| reward = None |
| rewards = [] |
| @@ -96,6 +98,7 @@ def eval_rl(model, processor, task, eval_args): |
| frames.append(np.array(env.render(), dtype=np.uint8)) |
| |
| scores.append(sum(rewards)) |
| + myprint(('-'*100) + '\n\n\n') |
| env.close() |
| |
| raw_mean, raw_std = np.mean(scores), np.std(scores) |
| @@ -145,7 +148,9 @@ def main(): |
| tasks.extend([env_id for env_id in TASK_NAME_TO_ENV_ID.keys() if env_id.startswith(domain)]) |
| |
| device = torch.device("cpu") if eval_args.use_cpu else get_default_device() |
| - processor = None |
| + processor = AutoProcessor.from_pretrained( |
| + 'jat-project/jat', cache_dir=None, trust_remote_code=True |
| + ) |
| |
| evaluations = {} |
| video_list = [] |
| @@ -153,14 +158,18 @@ def main(): |
| |
| for task in tqdm(tasks, desc="Evaluation", unit="task", leave=True): |
| if task in TASK_NAME_TO_ENV_ID.keys(): |
| + myprint(('-'*100) + f'{task=}') |
| dataset = load_from_disk(f'{HF_DATASETS_CACHE}/jat-project/jat-dataset-tokenized/{task}') |
| - model = RandP(dataset) |
| + model = RandP(task, |
| + dataset, |
| + device,) |
| scores, frames, fps = eval_rl(model, processor, task, eval_args) |
| evaluations[task] = scores |
| # Save the video |
| if eval_args.save_video: |
| video_list.append(frames) |
| input_fps.append(fps) |
| + myprint(('-'*100) + '\n\n\n') |
| else: |
| warnings.warn(f"Task {task} is not supported.") |
| |
| |
| |
| |
| |
| @@ -8,7 +8,7 @@ import time |
| from datetime import datetime |
| from datasets import load_from_disk |
| from datasets.config import HF_DATASETS_CACHE |
| -from jat_regent.utils import myprint, process_row_atari, process_row_vector, retrieve_atari, retrieve_vector |
| +from jat_regent.utils import myprint, get_task_info, collect_all_data, process_row_of_obs_atari_full_without_mask, retrieve_atari, retrieve_vector, collect_subset, build_index_vector |
| import logging |
| logging.basicConfig(level=logging.DEBUG) |
| |
| @@ -17,7 +17,8 @@ def main(): |
| parser = argparse.ArgumentParser(description='Build RAAGENT sequence indices') |
| parser.add_argument('--task', type=str, default='atari-alien', help='Task name') |
| parser.add_argument('--num_to_retrieve', type=int, default=100, help='Number of states/windows to retrieve') |
| - parser.add_argument('--nb_cores_autofaiss', type=int, default=16, help='Number of cores to use for faiss in vector observation environments') |
| + parser.add_argument('--nb_cores_autofaiss', type=int, default=16, help='Number of cores to use for faiss in vector obs envs') |
| + parser.add_argument('--batch_size_retrieval', type=int, default=1024, help='Batch size for retrieval in atari') |
| args = parser.parse_args() |
| |
| # load dataset, map, device, for task |
| @@ -25,77 +26,83 @@ def main(): |
| dataset_path = f"{HF_DATASETS_CACHE}/jat-project/jat-dataset-tokenized/{task}" |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| |
| - rew_key = 'rewards' |
| - attn_key = 'attention_mask' |
| - if task.startswith("atari"): |
| - obs_key = 'image_observations' |
| - act_key = 'discrete_actions' |
| - len_row_tokenized_known = 32 # half of 54 |
| - process_row_fn = process_row_atari |
| - retrieve_fn = retrieve_atari |
| - elif task.startswith("babyai"): |
| - obs_key = 'discrete_observations'# also has 'text_observations' only for raw dataset not for tokenized dataset (as it is combined into discrete_observation in tokenized dataset) |
| - act_key = 'discrete_actions' |
| - len_row_tokenized_known = 256 # half of 512 |
| - process_row_fn = lambda attn_mask, row_of_obs, task: process_row_vector(attn_mask, row_of_obs, task, return_numpy=True) |
| - retrieve_fn = retrieve_vector |
| - elif task.startswith("metaworld") or task.startswith("mujoco"): |
| - obs_key = 'continuous_observations' |
| - act_key = 'continuous_actions' |
| - len_row_tokenized_known = 256 |
| - process_row_fn = lambda attn_mask, row_of_obs, task: process_row_vector(attn_mask, row_of_obs, task, return_numpy=True) |
| - retrieve_fn = retrieve_vector |
| + rew_key, attn_key, obs_key, act_key, B, obs_dim = get_task_info(task) |
| |
| dataset = load_from_disk(dataset_path) |
| with open(f"{dataset_path}/map_from_rows_to_episodes_for_tokenized.json", 'r') as f: |
| map_from_rows_to_episodes_for_tokenized = json.load(f) |
| |
| # setup kwargs |
| - len_dataset = len(dataset['train']) |
| - B = len_row_tokenized_known |
| kwargs = {'B': B, |
| - 'attn_key':attn_key, |
| - 'obs_key':obs_key, |
| - 'device':device, |
| - 'task':task, |
| - 'batch_size_retrieval':None, |
| - 'nb_cores_autofaiss':None if task.startswith("atari") else args.nb_cores_autofaiss, |
| - } |
| + 'obs_dim': obs_dim, |
| + 'attn_key': attn_key, |
| + 'obs_key': obs_key, |
| + 'device': device, |
| + 'task': task, |
| + 'batch_size_retrieval': args.batch_size_retrieval, |
| + 'nb_cores_autofaiss': None if task.startswith("atari") else args.nb_cores_autofaiss, |
| + } |
| |
| # collect all observations in a single array (this takes some time) for vector observation environments |
| - if not task.startswith("atari"): |
| - myprint("Collecting all observations/attn_masks in a single array") |
| - all_rows_of_obs = np.array(dataset['train'][obs_key]) |
| - all_attn_masks = np.array(dataset['train'][attn_key]).astype(bool) |
| + myprint("Collecting all observations/attn_masks") |
| + all_rows_of_obs_OG, all_attn_masks_OG, all_row_idxs = collect_all_data(dataset, task, obs_key) |
| |
| # iterate over rows |
| all_retrieved_indices = [] |
| - for row_idx in range(len_dataset): |
| - myprint(f"\nProcessing row {row_idx}/{len_dataset}") |
| + for row_idx in all_row_idxs: |
| + myprint(f"\nProcessing row {row_idx}/{len(all_row_idxs)}") |
| current_ep = map_from_rows_to_episodes_for_tokenized[str(row_idx)] |
| |
| - attn_mask, row_of_obs = process_row_fn(dataset['train'][row_idx][attn_key], dataset['train'][row_idx][obs_key], task) |
| + # get row_of_obs and attn_mask |
| + datarow = dataset['train'][row_idx] |
| + attn_mask = np.array(datarow[attn_key]).astype(bool) |
| + if task.startswith("atari"): |
| + row_of_obs = process_row_of_obs_atari_full_without_mask(datarow[obs_key]) |
| + else: |
| + row_of_obs = np.array(datarow[obs_key]) |
| + row_of_obs = row_of_obs[attn_mask] |
| + assert row_of_obs.shape == (np.sum(attn_mask), *obs_dim) |
| |
| # compare with rows from all but the current episode |
| - all_other_rows = [idx for idx in range(len_dataset) if map_from_rows_to_episodes_for_tokenized[str(idx)] != current_ep] |
| + all_other_row_idxs = [idx for idx in all_row_idxs if map_from_rows_to_episodes_for_tokenized[str(idx)] != current_ep] |
| |
| # do the retrieval |
| - retrieved_indices = retrieve_fn(row_of_obs=row_of_obs, |
| - dataset=dataset if task.startswith("atari") else (all_rows_of_obs, all_attn_masks), |
| - all_rows_to_consider=all_other_rows, |
| - num_to_retrieve=args.num_to_retrieve, |
| - kwargs=kwargs, |
| - ) |
| + if task.startswith("atari"): |
| + all_indices, all_processed_rows_of_obs = collect_subset(all_rows_of_obs_OG=all_rows_of_obs_OG, |
| + all_attn_masks_OG=all_attn_masks_OG, |
| + all_rows_to_consider=all_row_idxs, |
| + kwargs=kwargs) |
| + retrieved_indices = retrieve_atari(row_of_obs=row_of_obs, |
| + all_processed_rows_of_obs=all_processed_rows_of_obs, |
| + all_indices=all_indices, |
| + num_to_retrieve=args.num_to_retrieve, |
| + kwargs=kwargs) |
| + else: |
| + all_indices, knn_index = build_index_vector(all_rows_of_obs_OG=all_rows_of_obs_OG, |
| + all_attn_masks_OG=all_attn_masks_OG, |
| + all_rows_to_consider=all_other_row_idxs, |
| + kwargs=kwargs) |
| + retrieved_indices = retrieve_vector(row_of_obs=row_of_obs, |
| + knn_index=knn_index, |
| + all_indices=all_indices, |
| + num_to_retrieve=args.num_to_retrieve, |
| + kwargs=kwargs) |
| + |
| + # pad the above to expected B |
| + xbdim = row_of_obs.shape[0] |
| + if xbdim < B: |
| + retrieved_indices = np.concatenate([retrieved_indices, np.zeros((B-xbdim, args.num_to_retrieve, 2), dtype=int)], axis=0) |
| + assert retrieved_indices.shape == (B, args.num_to_retrieve, 2) |
| |
| # collect retrieved indices |
| all_retrieved_indices.append(retrieved_indices) |
| |
| # concat |
| all_retrieved_indices = np.stack(all_retrieved_indices, axis=0) |
| - assert all_retrieved_indices.shape == (len_dataset, B, args.num_to_retrieve, 2) |
| + assert all_retrieved_indices.shape == (len(all_row_idxs), B, args.num_to_retrieve, 2) |
| |
| # save arrays as bin for easy memmap access and faster loading |
| - all_retrieved_indices.tofile(f"{dataset_path}/retrieved_indices_{len_dataset}_{B}_{args.num_to_retrieve}_2.bin") |
| + all_retrieved_indices.tofile(f"{dataset_path}/retrieved_indices_{len(all_row_idxs)}_{B}_{args.num_to_retrieve}_2.bin") |
| |
| if __name__ == "__main__": |
| main() |
| \ No newline at end of file |
|
|
