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:bc78c1066c3576737f3dfa984fb5bb70008d2b895ed34a14de7b0952798a7a3c
+size 19343400

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_journeyescape
+      type: atari_journeyescape
+    metrics:
+    - type: mean_reward
+      value: -21220.00 +/- 7108.14
+      name: mean_reward
+      verified: false
+---
+
+A(n) **APPO** model trained on the **atari_journeyescape** 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_journeyescape_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_journeyescape --train_dir=./train_dir --experiment=atari_2B_atari_journeyescape_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_journeyescape --train_dir=./train_dir --experiment=atari_2B_atari_journeyescape_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_000436848_894664704_reward_-0.120.pth

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

checkpoint_p0/checkpoint_000976608_2000093184.pth

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

checkpoint_p0/checkpoint_000976624_2000125952.pth

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

checkpoint_p0/milestones/checkpoint_000024512_50200576.pth

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

checkpoint_p0/milestones/checkpoint_000049728_101842944.pth

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

checkpoint_p0/milestones/checkpoint_000075328_154271744.pth

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

checkpoint_p0/milestones/checkpoint_000100816_206471168.pth

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

checkpoint_p0/milestones/checkpoint_000126208_258473984.pth

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

checkpoint_p0/milestones/checkpoint_000151936_311164928.pth

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

checkpoint_p0/milestones/checkpoint_000177600_363724800.pth

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

checkpoint_p0/milestones/checkpoint_000202912_415563776.pth

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

checkpoint_p0/milestones/checkpoint_000228480_467927040.pth

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

checkpoint_p0/milestones/checkpoint_000253888_519962624.pth

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

checkpoint_p0/milestones/checkpoint_000279104_571604992.pth

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

checkpoint_p0/milestones/checkpoint_000304512_623640576.pth

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

checkpoint_p0/milestones/checkpoint_000329728_675282944.pth

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

checkpoint_p0/milestones/checkpoint_000354848_726728704.pth

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

checkpoint_p0/milestones/checkpoint_000380032_778305536.pth

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

checkpoint_p0/milestones/checkpoint_000405344_830144512.pth

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

checkpoint_p0/milestones/checkpoint_000430848_882376704.pth

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

checkpoint_p0/milestones/checkpoint_000456192_934281216.pth

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

checkpoint_p0/milestones/checkpoint_000481664_986447872.pth

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

checkpoint_p0/milestones/checkpoint_000507136_1038614528.pth

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

checkpoint_p0/milestones/checkpoint_000532416_1090387968.pth

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

checkpoint_p0/milestones/checkpoint_000557760_1142292480.pth

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

checkpoint_p0/milestones/checkpoint_000583200_1194393600.pth

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

checkpoint_p0/milestones/checkpoint_000608640_1246494720.pth

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

checkpoint_p0/milestones/checkpoint_000633792_1298006016.pth

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

checkpoint_p0/milestones/checkpoint_000659008_1349648384.pth

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

checkpoint_p0/milestones/checkpoint_000684480_1401815040.pth

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

checkpoint_p0/milestones/checkpoint_000709440_1452933120.pth

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

checkpoint_p0/milestones/checkpoint_000734656_1504575488.pth

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

checkpoint_p0/milestones/checkpoint_000759744_1555955712.pth

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

checkpoint_p0/milestones/checkpoint_000784448_1606549504.pth

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

checkpoint_p0/milestones/checkpoint_000809584_1658028032.pth

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

checkpoint_p0/milestones/checkpoint_000834608_1709277184.pth

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

checkpoint_p0/milestones/checkpoint_000859648_1760559104.pth

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

checkpoint_p0/milestones/checkpoint_000887168_1816920064.pth

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

checkpoint_p0/milestones/checkpoint_000931376_1907458048.pth

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

config.json

@@ -0,0 +1,167 @@
+{
+  "help": false,
+  "algo": "APPO",
+  "env": "atari_journeyescape",
+  "experiment": "atari_2B_atari_journeyescape_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_journeyescape_1111 --env=atari_journeyescape --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_journeyescape",
+    "experiment": "atari_2B_atari_journeyescape_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:28db7b90b8fb514eb5627dee99b37d5ba301f4cb67f2bcb9f0cb2dd93d7cafd9
+size 4872793