Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+replay.mp4 filter=lfs diff=lfs merge=lfs -text

.summary/0/events.out.tfevents.1718465455.koa03

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e94cf53b8f6360348849b631b4c38a3b6f7f36c0e9a04db91a35f5b7b9ccd542
+size 18285797

README.md

@@ -0,0 +1,56 @@
+---
+library_name: sample-factory
+tags:
+- deep-reinforcement-learning
+- reinforcement-learning
+- sample-factory
+model-index:
+- name: APPO
+  results:
+  - task:
+      type: reinforcement-learning
+      name: reinforcement-learning
+    dataset:
+      name: atari_airraid
+      type: atari_airraid
+    metrics:
+    - type: mean_reward
+      value: 465.00 +/- 182.76
+      name: mean_reward
+      verified: false
+---
+
+A(n) **APPO** model trained on the **atari_airraid** environment.
+
+This model was trained using Sample-Factory 2.0: https://github.com/alex-petrenko/sample-factory.
+Documentation for how to use Sample-Factory can be found at https://www.samplefactory.dev/
+
+
+## Downloading the model
+
+After installing Sample-Factory, download the model with:
+```
+python -m sample_factory.huggingface.load_from_hub -r ksridhar/atari_2B_atari_airraid_1111
+```
+
+    
+## Using the model
+
+To run the model after download, use the `enjoy` script corresponding to this environment:
+```
+python -m <path.to.enjoy.module> --algo=APPO --env=atari_airraid --train_dir=./train_dir --experiment=atari_2B_atari_airraid_1111
+```
+
+
+You can also upload models to the Hugging Face Hub using the same script with the `--push_to_hub` flag.
+See https://www.samplefactory.dev/10-huggingface/huggingface/ for more details
+    
+## Training with this model
+
+To continue training with this model, use the `train` script corresponding to this environment:
+```
+python -m <path.to.train.module> --algo=APPO --env=atari_airraid --train_dir=./train_dir --experiment=atari_2B_atari_airraid_1111 --restart_behavior=resume --train_for_env_steps=10000000000
+```
+
+Note, you may have to adjust `--train_for_env_steps` to a suitably high number as the experiment will resume at the number of steps it concluded at.
+    

checkpoint_p0/best_000820672_1680736256_reward_1255.320.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:04ea5891f52bccc613ae420f2383c694c8a4583d64cce037fe061a08240982b3
+size 20722280

checkpoint_p0/checkpoint_000976608_2000093184.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3f5b7f2c104719944467b95f32e263dff2734fcfce0a7940f053de310ca23fb9
+size 20722628

checkpoint_p0/checkpoint_000976624_2000125952.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e83c5aba7700b73f5f86d2763f21eadcaee8e1daabfd945d07def7a62c2fdfab
+size 20722628

checkpoint_p0/milestones/checkpoint_000025328_51871744.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c70ea6c746b1099395d74f65b51a553989d2373949cda55bdb936f6d339f163a
+size 20723568

checkpoint_p0/milestones/checkpoint_000052832_108199936.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:90153c065eaf8a16774217a7ce5a67685dd5fc0af1d6f8ce55b475a50deea4fe
+size 20723626

checkpoint_p0/milestones/checkpoint_000080576_165019648.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:48af7f9d07dc3e13cd2d01cff53a5d4c142afb6e5d5cb46c3f8fcbf78c529f52
+size 20723626

checkpoint_p0/milestones/checkpoint_000108128_221446144.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ed2dbe5193dcea57fbc0d9dce3e6d3bd959bb251e232176daa539bfc93a24530
+size 20723626

checkpoint_p0/milestones/checkpoint_000136192_278921216.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba07de6c1b32632354f3d77b63a47c28e412bfed43ca5ef82c3e33fa2616785e
+size 20723626

checkpoint_p0/milestones/checkpoint_000164160_336199680.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cf9e64c7241e64da17332f5779f522e2efd38d87cf5847eeb914f1b86bb90eeb
+size 20723626

checkpoint_p0/milestones/checkpoint_000191744_392691712.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fcce6b301593984a3c055b3fbe8f21f3823b1b3233c9a405a5c2c62168f8d45f
+size 20723626

checkpoint_p0/milestones/checkpoint_000220096_450756608.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:60e5d3544717f7983b4abe5a622035a4949d20a3a07b0a6d2732d54e15004567
+size 20723626

checkpoint_p0/milestones/checkpoint_000248000_507904000.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:04bef3f4845d5afb30cb49f9bc48dd421b1ae398e2e249b56dfeda5b1f089906
+size 20723626

checkpoint_p0/milestones/checkpoint_000276096_565444608.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d4eb1a15333413a30b47229e6e8e1236096c8fb60dc20560f0762f5d5a7ee491
+size 20723626

checkpoint_p0/milestones/checkpoint_000303856_622297088.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8ed4b042c87e40ee5edcc31499c946d83b3206d68b572009c7f0b9b1bb9e51e8
+size 20723626

checkpoint_p0/milestones/checkpoint_000331888_679706624.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:506f2656ecd9f241a5175470d7502dea2df70e4b5f9c0fb911219fe041c906e8
+size 20723626

checkpoint_p0/milestones/checkpoint_000359664_736591872.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e56a0ed00e07cb870607d3ac6727ae2d1c952cfff4156fbabf41954a816271fe
+size 20723626

checkpoint_p0/milestones/checkpoint_000387616_793837568.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c4206c3d4735f6aada46266d53b372afef0ecf8aee1516538fdde4128e73ca35
+size 20723626

checkpoint_p0/milestones/checkpoint_000415744_851443712.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cdbeb1c1d4facca0691907133f028f7a395d7a4631e024bb962fd5eae9e0473e
+size 20723626

checkpoint_p0/milestones/checkpoint_000444096_909508608.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:908e6f59f7f70604da8a49c0effa07c915a4c8181b2229bbf5060448eab34237
+size 20723626

checkpoint_p0/milestones/checkpoint_000472512_967704576.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3acfe6effa0eb87d9be6c64ded5e46f73548ab35f2b9324282a9fab72952cdd9
+size 20723626

checkpoint_p0/milestones/checkpoint_000500608_1025245184.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c0562aa2dd28df15165bbc10ac172196e667e69694225bb419790a0204872b89
+size 20723684

checkpoint_p0/milestones/checkpoint_000528672_1082720256.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a359cb85026316b46cec52c2b3072bb3f4c6d0c8f8c04b81652d53c320962ab
+size 20723684

checkpoint_p0/milestones/checkpoint_000556928_1140588544.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67e509d0cba38cab4cb931ce1b06b3c0dcd997802b1cd0e2cf62bf42bba8c964
+size 20723684

checkpoint_p0/milestones/checkpoint_000585408_1198915584.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:123f1825e68d6007a227a2809108a82997c6140e78cd4b441124d5de77909c72
+size 20723684

checkpoint_p0/milestones/checkpoint_000614144_1257766912.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:58fed7f6964ad4469525aa4840aab5182a39bd3fa981c5d0c85131980dff948b
+size 20723684

checkpoint_p0/milestones/checkpoint_000642240_1315307520.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2c2b24a591716ea07f740d151b7b912eb8da99d438b7421c1c90e458950ff079
+size 20723684

checkpoint_p0/milestones/checkpoint_000670720_1373634560.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:06d76f4a0427517572114f56ad297b0820be092421150f56f0b045289b88fba0
+size 20723684

checkpoint_p0/milestones/checkpoint_000698816_1431175168.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67eedd32ef45079fb6063fa17b3573d4cfa3756ff2b850d2ed9a7cfa3ff2e9f8
+size 20723684

checkpoint_p0/milestones/checkpoint_000726976_1488846848.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f799ebc4f68d484712629e50e5e50b61a5cf1c228487f35565bef8fad96877ef
+size 20723684

checkpoint_p0/milestones/checkpoint_000755328_1546911744.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:acf482fd0d35a9e9466747e0902a0113179aa3153d6457d9069b62d5006a9af1
+size 20723684

checkpoint_p0/milestones/checkpoint_000783136_1603862528.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a72d7d19a447c00346e260ca9efe1737505c1bdd124bc713d2fb5f600d758db2
+size 20723684

checkpoint_p0/milestones/checkpoint_000811584_1662124032.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2d129fb12044a642bcf9cce81c2a72773e839a07e7d83b400c28ffa6f6e77dda
+size 20723684

checkpoint_p0/milestones/checkpoint_000839936_1720188928.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:04aa797855e355aa77fc037b122b93c36ef03add6458a5964950c787ec16b547
+size 20723684

checkpoint_p0/milestones/checkpoint_000868144_1777958912.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6f2e77ad1f2d084ca4b8a0560788be4a2c00b2215895a53f0b1ce94ee419dc2f
+size 20723684

checkpoint_p0/milestones/checkpoint_000896240_1835499520.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8d45adf6d55ca09b5f89e4178e6eae4c7844bf2e42f0c19985040c7e348a4f6a
+size 20723684

checkpoint_p0/milestones/checkpoint_000924672_1893728256.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a8057f3255e306fda08291c8e367d5adc5bf1baea5f0da09fb14f76e3d035666
+size 20723684

checkpoint_p0/milestones/checkpoint_000953008_1951760384.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a07fc48118e729c1125f70c78b65c102fff6c9714e8cc37b8455fca3da50fa3
+size 20723684

config.json

@@ -0,0 +1,167 @@
+{
+  "help": false,
+  "algo": "APPO",
+  "env": "atari_airraid",
+  "experiment": "atari_2B_atari_airraid_1111",
+  "train_dir": "train_dir",
+  "restart_behavior": "resume",
+  "device": "gpu",
+  "seed": 1111,
+  "num_policies": 1,
+  "async_rl": true,
+  "serial_mode": false,
+  "batched_sampling": true,
+  "num_batches_to_accumulate": 2,
+  "worker_num_splits": 1,
+  "policy_workers_per_policy": 1,
+  "max_policy_lag": 1000,
+  "num_workers": 4,
+  "num_envs_per_worker": 1,
+  "batch_size": 1024,
+  "num_batches_per_epoch": 8,
+  "num_epochs": 2,
+  "rollout": 64,
+  "recurrence": 1,
+  "shuffle_minibatches": false,
+  "gamma": 0.99,
+  "reward_scale": 1.0,
+  "reward_clip": 1000.0,
+  "value_bootstrap": false,
+  "normalize_returns": true,
+  "exploration_loss_coeff": 0.0004677351413,
+  "value_loss_coeff": 0.5,
+  "kl_loss_coeff": 0.0,
+  "exploration_loss": "entropy",
+  "gae_lambda": 0.95,
+  "ppo_clip_ratio": 0.1,
+  "ppo_clip_value": 1.0,
+  "with_vtrace": false,
+  "vtrace_rho": 1.0,
+  "vtrace_c": 1.0,
+  "optimizer": "adam",
+  "adam_eps": 1e-05,
+  "adam_beta1": 0.9,
+  "adam_beta2": 0.999,
+  "max_grad_norm": 0.0,
+  "learning_rate": 0.0003033891184,
+  "lr_schedule": "linear_decay",
+  "lr_schedule_kl_threshold": 0.008,
+  "lr_adaptive_min": 1e-06,
+  "lr_adaptive_max": 0.01,
+  "obs_subtract_mean": 0.0,
+  "obs_scale": 255.0,
+  "normalize_input": true,
+  "normalize_input_keys": [
+    "obs"
+  ],
+  "decorrelate_experience_max_seconds": 1,
+  "decorrelate_envs_on_one_worker": true,
+  "actor_worker_gpus": [],
+  "set_workers_cpu_affinity": true,
+  "force_envs_single_thread": false,
+  "default_niceness": 0,
+  "log_to_file": true,
+  "experiment_summaries_interval": 3,
+  "flush_summaries_interval": 30,
+  "stats_avg": 100,
+  "summaries_use_frameskip": true,
+  "heartbeat_interval": 20,
+  "heartbeat_reporting_interval": 180,
+  "train_for_env_steps": 2000000000,
+  "train_for_seconds": 3600000,
+  "save_every_sec": 120,
+  "keep_checkpoints": 2,
+  "load_checkpoint_kind": "latest",
+  "save_milestones_sec": 1200,
+  "save_best_every_sec": 5,
+  "save_best_metric": "reward",
+  "save_best_after": 100000,
+  "benchmark": false,
+  "encoder_mlp_layers": [
+    512,
+    512
+  ],
+  "encoder_conv_architecture": "convnet_atari",
+  "encoder_conv_mlp_layers": [
+    512
+  ],
+  "use_rnn": false,
+  "rnn_size": 512,
+  "rnn_type": "gru",
+  "rnn_num_layers": 1,
+  "decoder_mlp_layers": [],
+  "nonlinearity": "relu",
+  "policy_initialization": "orthogonal",
+  "policy_init_gain": 1.0,
+  "actor_critic_share_weights": true,
+  "adaptive_stddev": false,
+  "continuous_tanh_scale": 0.0,
+  "initial_stddev": 1.0,
+  "use_env_info_cache": false,
+  "env_gpu_actions": false,
+  "env_gpu_observations": true,
+  "env_frameskip": 4,
+  "env_framestack": 4,
+  "pixel_format": "CHW",
+  "use_record_episode_statistics": true,
+  "episode_counter": false,
+  "with_wandb": false,
+  "wandb_user": null,
+  "wandb_project": "sample_factory",
+  "wandb_group": null,
+  "wandb_job_type": "SF",
+  "wandb_tags": [],
+  "with_pbt": false,
+  "pbt_mix_policies_in_one_env": true,
+  "pbt_period_env_steps": 5000000,
+  "pbt_start_mutation": 20000000,
+  "pbt_replace_fraction": 0.3,
+  "pbt_mutation_rate": 0.15,
+  "pbt_replace_reward_gap": 0.1,
+  "pbt_replace_reward_gap_absolute": 1e-06,
+  "pbt_optimize_gamma": false,
+  "pbt_target_objective": "true_objective",
+  "pbt_perturb_min": 1.1,
+  "pbt_perturb_max": 1.5,
+  "env_agents": 512,
+  "command_line": "--seed=1111 --experiment=atari_2B_atari_airraid_1111 --env=atari_airraid --train_for_seconds=3600000 --algo=APPO --gamma=0.99 --num_workers=4 --num_envs_per_worker=1 --worker_num_splits=1 --env_agents=512 --benchmark=False --max_grad_norm=0.0 --decorrelate_experience_max_seconds=1 --encoder_conv_architecture=convnet_atari --encoder_conv_mlp_layers 512 --nonlinearity=relu --num_policies=1 --normalize_input=True --normalize_input_keys obs --normalize_returns=True --async_rl=True --batched_sampling=True --train_for_env_steps=2000000000 --save_milestones_sec=1200 --train_dir train_dir --rollout 64 --exploration_loss_coeff 0.0004677351413 --num_epochs 2 --batch_size 1024 --num_batches_per_epoch 8 --learning_rate 0.0003033891184",
+  "cli_args": {
+    "algo": "APPO",
+    "env": "atari_airraid",
+    "experiment": "atari_2B_atari_airraid_1111",
+    "train_dir": "train_dir",
+    "seed": 1111,
+    "num_policies": 1,
+    "async_rl": true,
+    "batched_sampling": true,
+    "worker_num_splits": 1,
+    "num_workers": 4,
+    "num_envs_per_worker": 1,
+    "batch_size": 1024,
+    "num_batches_per_epoch": 8,
+    "num_epochs": 2,
+    "rollout": 64,
+    "gamma": 0.99,
+    "normalize_returns": true,
+    "exploration_loss_coeff": 0.0004677351413,
+    "max_grad_norm": 0.0,
+    "learning_rate": 0.0003033891184,
+    "normalize_input": true,
+    "normalize_input_keys": [
+      "obs"
+    ],
+    "decorrelate_experience_max_seconds": 1,
+    "train_for_env_steps": 2000000000,
+    "train_for_seconds": 3600000,
+    "save_milestones_sec": 1200,
+    "benchmark": false,
+    "encoder_conv_architecture": "convnet_atari",
+    "encoder_conv_mlp_layers": [
+      512
+    ],
+    "nonlinearity": "relu",
+    "env_agents": 512
+  },
+  "git_hash": "e259c57b8c7aa9c7f541e9efd1316f8e6f97a6db",
+  "git_repo_name": "https://github.com/kaustubhsridhar/jat_regent.git"
+}

git.diff

@@ -0,0 +1,712 @@
+diff --git a/README.md b/README.md
+index e51a12b..a6e1ca1 100644
+--- a/README.md
++++ b/README.md
+@@ -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
+diff --git a/jat_regent/RandP.py b/jat_regent/RandP.py
+deleted file mode 100644
+index b2bd8bf..0000000
+--- a/jat_regent/RandP.py
++++ /dev/null
+@@ -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
+
+diff --git a/jat_regent/modelling_jat_regent.py b/jat_regent/modelling_jat_regent.py
+deleted file mode 100644
+index e69de29..0000000
+diff --git a/jat_regent/utils.py b/jat_regent/utils.py
+index 56bfb44..36f6cca 100644
+--- a/jat_regent/utils.py
++++ b/jat_regent/utils.py
+@@ -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
+
+diff --git a/scripts_regent/eval_RandP.py b/scripts_regent/eval_RandP.py
+index 07e545c..146b347 100755
+--- a/scripts_regent/eval_RandP.py
++++ b/scripts_regent/eval_RandP.py
+@@ -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.")
+ 
+diff --git a/scripts_regent/offline_retrieval_jat_regent.py b/scripts_regent/offline_retrieval_jat_regent.py
+index c83d259..aad678a 100644
+--- a/scripts_regent/offline_retrieval_jat_regent.py
++++ b/scripts_regent/offline_retrieval_jat_regent.py
+@@ -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()
+

replay.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:07be059a9429a473b5b17baa3708722b914b49c1e81c2c57e350ea5acb4339b7
+size 1295354