Merge branch 'full_ft'

2026-09-09 11:24:31 +08:00 · 2024-04-27 07:04:40 +08:00
parent 075de416c8 a82fe96158
commit 991f157fdd
32 changed files with 2537 additions and 1219 deletions
@@ -0,0 +1,42 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "test",
+            "type": "python",
+            "request": "launch",
+            "program": "${workspaceFolder}/src/main.py",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+            "autoReload": {"enable": true,},
+            "env": {"WANDB_MODE":"disabled"},
+            "args": [
+                // "'wandb.mode=disabled",
+                // "env.train.id=BreakoutNoFrameskip-v4",
+                "env.train.id=CrafterReward-v1",
+                // # make it start early
+                "training.tokenizer.start_after_epochs=1",
+                "training.world_model.start_after_epochs=2",
+                "training.actor_critic.start_after_epochs=3",
+                "training.tokenizer.steps_per_epoch=10",
+                "training.world_model.steps_per_epoch=10",
+                "training.actor_critic.steps_per_epoch=10",
+            ]
+        },
+        {
+            "name": "main",
+            "type": "python",
+            "request": "launch",
+            "program": "${workspaceFolder}/src/main.py",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+            "autoReload": {"enable": true,},
+            "args": [
+                "env.train.id=CrafterReward-v1",
+            ]
+        }
+    ]
+}
@@ -5,6 +5,19 @@ See also:
 - [AdaVAE](https://github.com/ImKeTT/AdaVAE)
 - [bigvae](https://github.com/JD-P/minihf/blob/adavae-moe/vae_infer.py)

+
+A fork of IRIS where I use a pretrained LLM as the tranformer (with LoRa). 
+
+My hypothesis: Pretrained LLM's make good world models by including a lot of world information!
+
+details:
+- for speed and cost I use a small 1.5B model. But it would be interesting to try a 7B one
+- for speed I use a smaller actor critic than in IRIS
+- max_blocks 20->10
+- batch smaller because of my small machine
+- actor_critic.steps_per_epoch 200->20
+- world_model.batch_num_sampler; 64->8 because the forzen transformer uses lots of gpu ram
+
 # Transformers are Sample-Efficient World Models (IRIS)

 [Transformers are Sample-Efficient World Models](https://openreview.net/forum?id=vhFu1Acb0xb) <br>
@@ -1 +1,2 @@
 use_original_obs: False
+lstm_dim: 512
@@ -1,8 +1,8 @@
 train:
-  _target_: dataset.EpisodesDatasetRamMonitoring
+  _target_: src.dataset.EpisodesDatasetRamMonitoring
  max_ram_usage: 30G
  name: train_dataset
 test:
-  _target_: dataset.EpisodesDataset
+  _target_: src.dataset.EpisodesDataset
  max_num_episodes: null
  name: test_dataset
@@ -1,5 +1,5 @@
 train:
-  _target_: envs.make_atari
+  _target_: src.envs.make_env
  id: null
  size: 64
  max_episode_steps: 20000
@@ -18,4 +18,4 @@ test:
  done_on_life_loss: False
  clip_reward: False

-keymap: atari/${.train.id}
+keymap: atari/${.train.id}
@@ -1,14 +1,14 @@
-_target_: models.tokenizer.Tokenizer
+_target_: src.models.tokenizer.Tokenizer

-vocab_size: 512
-embed_dim: 512
+vocab_size: ${..world_model.vocab_size}
+embed_dim: ${..world_model.embed_dim}
 encoder:
-  _target_: models.tokenizer.Encoder
+  _target_: src.models.tokenizer.Encoder
  config:
-    _target_: models.tokenizer.EncoderDecoderConfig
+    _target_: src.models.tokenizer.EncoderDecoderConfig
    resolution: 64
    in_channels: 3
-    z_channels: 512
+    z_channels: ${...vocab_size}
    ch: 64
    ch_mult: [1, 1, 1, 1, 1]
    num_res_blocks: 2
@@ -16,5 +16,5 @@ encoder:
    out_ch: 3
    dropout: 0.0
 decoder:
-  _target_: models.tokenizer.Decoder
-  config: ${..encoder.config}
+  _target_: src.models.tokenizer.Decoder
+  config: ${..encoder.config}
@@ -54,24 +54,24 @@ training:
  should: True
  learning_rate: 0.0001
  tokenizer:
-    batch_num_samples: 256
+    batch_num_samples: 128
    grad_acc_steps: 1
    max_grad_norm: 10.0
    start_after_epochs: 5
    steps_per_epoch: 200
  world_model:
-    batch_num_samples: 64
-    grad_acc_steps: 1
+    batch_num_samples: 8  # pretrained models use lots of ram
+    grad_acc_steps: 2
    max_grad_norm: 10.0
    weight_decay: 0.01
    start_after_epochs: 25
    steps_per_epoch: 200
  actor_critic:
-    batch_num_samples: 64
+    batch_num_samples: 16
    grad_acc_steps: 1
    max_grad_norm: 10.0
-    start_after_epochs: 50
-    steps_per_epoch: 200
+    start_after_epochs: 300
+    steps_per_epoch: 40
    imagine_horizon: ${common.sequence_length}
    burn_in: 20
    gamma: 0.995
@@ -92,3 +92,7 @@ evaluation:
    num_episodes_to_save: ${training.actor_critic.batch_num_samples}
    horizon: ${training.actor_critic.imagine_horizon}
    start_after_epochs: ${training.actor_critic.start_after_epochs}
+
+hydra:
+  job:
+    chdir: True
@@ -1,2 +1,8 @@
-_target_: models.BigVAEConfig
-tokens_per_block: 17
+_target_: src.models.TransformerConfig
+max_blocks: 10  # this is the rollout length when training policy
+tokens_per_block: 17 # how much info we can encode
+dropout: 0.1
+rank: 32 # lora rank
+model_name: "PY007/TinyLlama-1.1B-intermediate-step-715k-1.5T"
+vocab_size: 32000 # change to llm vocab dim
+embed_dim: 2048  # change this to whatever the embedding dimension is in your pretrained llm 2048 for llama. 2560 for stablelm
@@ -0,0 +1,29 @@
+set shell := ["zsh", "-cu"]
+
+breakout:
+    python src/main.py env.train.id=BreakoutNoFrameskip-v4 
+
+crafter:
+    python src/main.py env.train.id=CrafterReward-v1
+
+# minihack:
+#     python src/main.py env.train.id=MiniHack-River-v0
+
+# watch the latest runs
+watch_latest:
+    . ./.venv/bin/activate
+    cd ./outputs && \
+    cd *([-1]) && \
+    cd *([-1]) && \
+    scripts/play.sh -e -r -h
+
+
+resume_latest:
+    . ./.venv/bin/activate
+    cd ./outputs && \
+    cd *([-1]) && \
+    cd *([-1]) && \
+    scripts/resume.sh
+
+default: 
+    just --list
@@ -0,0 +1,747 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "# autoreload import your package\n",
+    "%load_ext autoreload\n",
+    "%autoreload 2\n",
+    "\n",
+    "import gym\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "%matplotlib inline\n",
+    "plt.style.use('ggplot')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Debug model components\n",
+    "\n",
+    "### Using trainer? :poop:\n",
+    "\n",
+    "Hyrda is really annoying\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n",
+      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'wandb': {'mode': 'disabled', 'project': 'iris', 'entity': None, 'name': None, 'group': None, 'tags': None, 'notes': None}, 'initialization': {'path_to_checkpoint': None, 'load_tokenizer': False, 'load_world_model': False, 'load_actor_critic': False}, 'common': {'epochs': 600, 'device': 'cuda:0', 'do_checkpoint': False, 'seed': 0, 'sequence_length': '${world_model.max_blocks}', 'resume': True}, 'collection': {'train': {'num_envs': 1, 'stop_after_epochs': 500, 'num_episodes_to_save': 10, 'config': {'epsilon': 0.01, 'should_sample': True, 'temperature': 1.0, 'num_steps': 200, 'burn_in': '${training.actor_critic.burn_in}'}}, 'test': {'num_envs': 8, 'num_episodes_to_save': '${collection.train.num_episodes_to_save}', 'config': {'epsilon': 0.0, 'should_sample': True, 'temperature': 0.5, 'num_episodes': 16, 'burn_in': '${training.actor_critic.burn_in}'}}}, 'training': {'should': True, 'learning_rate': 0.0001, 'tokenizer': {'batch_num_samples': 128, 'grad_acc_steps': 1, 'max_grad_norm': 10.0, 'start_after_epochs': 1, 'steps_per_epoch': 10}, 'world_model': {'batch_num_samples': 4, 'grad_acc_steps': 1, 'max_grad_norm': 10.0, 'weight_decay': 0.01, 'start_after_epochs': 1, 'steps_per_epoch': 10}, 'actor_critic': {'batch_num_samples': 4, 'grad_acc_steps': 1, 'max_grad_norm': 10.0, 'start_after_epochs': 1, 'steps_per_epoch': 10, 'imagine_horizon': '${common.sequence_length}', 'burn_in': 20, 'gamma': 0.995, 'lambda_': 0.95, 'entropy_weight': 0.001}}, 'evaluation': {'should': True, 'every': 5, 'tokenizer': {'batch_num_samples': '${training.tokenizer.batch_num_samples}', 'start_after_epochs': '${training.tokenizer.start_after_epochs}', 'save_reconstructions': True}, 'world_model': {'batch_num_samples': '${training.world_model.batch_num_samples}', 'start_after_epochs': '${training.world_model.start_after_epochs}'}, 'actor_critic': {'num_episodes_to_save': '${training.actor_critic.batch_num_samples}', 'horizon': '${training.actor_critic.imagine_horizon}', 'start_after_epochs': '${training.actor_critic.start_after_epochs}'}}, 'tokenizer': {'_target_': 'src.models.tokenizer.Tokenizer', 'vocab_size': 2048, 'embed_dim': 2048, 'encoder': {'_target_': 'src.models.tokenizer.Encoder', 'config': {'_target_': 'src.models.tokenizer.EncoderDecoderConfig', 'resolution': 64, 'in_channels': 3, 'z_channels': 2048, 'ch': 64, 'ch_mult': [1, 1, 1, 1, 1], 'num_res_blocks': 2, 'attn_resolutions': [8, 16], 'out_ch': 3, 'dropout': 0.0}}, 'decoder': {'_target_': 'src.models.tokenizer.Decoder', 'config': '${..encoder.config}'}}, 'world_model': {'_target_': 'src.models.TransformerConfig', 'max_blocks': 10, 'num_layers': 1, 'num_heads': 1, 'embed_dim': 2048, 'dropout': 0.1, 'model_name': 'PY007/TinyLlama-1.1B-intermediate-step-715k-1.5T', 'rank': 32, 'tokens_per_block': 17}, 'actor_critic': {'use_original_obs': False, 'lstm_dim': 512}, 'env': {'train': {'_target_': 'src.envs.make_env', 'id': 'CrafterReward-v1', 'size': 64, 'max_episode_steps': 20000, 'noop_max': 30, 'frame_skip': 4, 'done_on_life_loss': True, 'clip_reward': False}, 'test': {'_target_': '${..train._target_}', 'id': '${..train.id}', 'size': '${..train.size}', 'max_episode_steps': 108000, 'noop_max': 1, 'frame_skip': '${..train.frame_skip}', 'done_on_life_loss': False, 'clip_reward': False}, 'keymap': 'atari/${.train.id}'}, 'datasets': {'train': {'_target_': 'src.dataset.EpisodesDatasetRamMonitoring', 'max_ram_usage': '30G', 'name': 'train_dataset'}, 'test': {'_target_': 'src.dataset.EpisodesDataset', 'max_num_episodes': None, 'name': 'test_dataset'}}}\n",
+      "Tokenizer : shape of latent is (2048, 4, 4).\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
+      "  warnings.warn(\n",
+      "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=VGG16_Weights.IMAGENET1K_V1`. You can also use `weights=VGG16_Weights.DEFAULT` to get the most up-to-date weights.\n",
+      "  warnings.warn(msg)\n",
+      "Using pad_token, but it is not set yet.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "trainable params: 50,462,720 || all params: 1,150,511,104 || trainable%: 4.386113252149889\n",
+      "None\n",
+      "32314243 parameters in agent.tokenizer\n",
+      "752979973 parameters in agent.world_model\n",
+      "3224626 parameters in agent.actor_critic\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "import os\n",
+    "os.environ['WANDB_MODE'] = 'disabled'\n",
+    "\n",
+    "import hydra\n",
+    "from hydra import initialize, initialize_config_module, initialize_config_dir, compose\n",
+    "from omegaconf import OmegaConf\n",
+    "\n",
+    "from pathlib import Path\n",
+    "from datetime import datetime\n",
+    "\n",
+    "from src.trainer import Trainer\n",
+    "\n",
+    "\n",
+    "class Trainer2(Trainer):\n",
+    "    \n",
+    "    def load_checkpoint(self, *args, **kwargs):\n",
+    "        pass\n",
+    "\n",
+    "\n",
+    "\n",
+    "ts = datetime.now().strftime(\"%Y-%m-%d/%H-%M-%S\")\n",
+    "run_dir = Path(f\"..outputs/{ts}\").absolute()\n",
+    "run_dir.mkdir(parents=True, exist_ok=True)\n",
+    "abs_config_dir=os.path.abspath(\"../config\")\n",
+    "os.chdir(run_dir)\n",
+    "# with initialize_config_dir(version_base=None, config_dir=abs_config_dir):\n",
+    "with initialize(version_base=None, config_path=\"../config\"):\n",
+    "    cfg = compose(config_name='trainer', overrides=[\n",
+    "        f'hydra.run.dir={run_dir}',\n",
+    "        # f\"initialization.path_to_checkpoint={str(path_to_checkpoint.absolute())}\",\n",
+    "        'wandb.mode=disabled',\n",
+    "        \"env.train.id=CrafterReward-v1\",\n",
+    "        \"training.tokenizer.start_after_epochs=1\",\n",
+    "        \"training.world_model.start_after_epochs=1\",\n",
+    "        \"training.actor_critic.start_after_epochs=1\",\n",
+    "        \"training.tokenizer.steps_per_epoch=10\",\n",
+    "        \"training.world_model.steps_per_epoch=10\",\n",
+    "        \"training.actor_critic.steps_per_epoch=10\",\n",
+    "        \"common.do_checkpoint=False\",\n",
+    "        \"common.resume=True\",\n",
+    "        \"training.world_model.batch_num_samples=4\",\n",
+    "        \"training.actor_critic.batch_num_samples=4\",\n",
+    "        ])\n",
+    "    print(cfg)\n",
+    "\n",
+    "    with run_dir:\n",
+    "        Path('media/episodes/train').mkdir(parents=True, exist_ok=True)\n",
+    "        trainer = Trainer2(cfg)\n",
+    "    trainer\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "## Trainer train_agent\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Experience collection (train_dataset): 100%|██████████| 200/200 [00:03<00:00, 58.00it/s]\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "[{'train_dataset/episode_length': 183,\n",
+       "  'train_dataset/episode_return': tensor(0.1000),\n",
+       "  'train_dataset/episode_num': 0,\n",
+       "  'train_dataset/action_histogram': <wandb.sdk.data_types.histogram.Histogram at 0x7f97f7373eb0>},\n",
+       " {'train_dataset/#episodes': 2,\n",
+       "  'train_dataset/#steps': 200,\n",
+       "  'train_dataset/return': 0.100000024}]"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "self=trainer\n",
+    "epoch = 52\n",
+    "\n",
+    "# get out first exp\n",
+    "self.train_collector.collect(self.agent, epoch, **self.cfg.collection.train.config)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "self.agent.train()\n",
+    "self.agent.zero_grad()\n",
+    "\n",
+    "metrics_tokenizer, metrics_world_model, metrics_actor_critic = {}, {}, {}\n",
+    "\n",
+    "cfg_tokenizer = self.cfg.training.tokenizer\n",
+    "cfg_world_model = self.cfg.training.world_model\n",
+    "cfg_actor_critic = self.cfg.training.actor_critic\n",
+    "\n",
+    "# if epoch > cfg_tokenizer.start_after_epochs:\n",
+    "#     metrics_tokenizer = self.train_component(self.agent.tokenizer, self.optimizer_tokenizer, sequence_length=1, sample_from_start=True, **cfg_tokenizer)\n",
+    "# self.agent.tokenizer.eval()\n",
+    "\n",
+    "# if epoch > cfg_world_model.start_after_epochs:\n",
+    "#     metrics_world_model = self.train_component(self.agent.world_model, self.optimizer_world_model, sequence_length=self.cfg.common.sequence_length, sample_from_start=True, tokenizer=self.agent.tokenizer, **cfg_world_model)\n",
+    "# self.agent.world_model.eval()\n",
+    "\n",
+    "# if epoch > cfg_actor_critic.start_after_epochs:\n",
+    "#     metrics_actor_critic = self.train_component(self.agent.actor_critic, self.optimizer_actor_critic, sequence_length=1 + self.cfg.training.actor_critic.burn_in, sample_from_start=False, tokenizer=self.agent.tokenizer, world_model=self.agent.world_model, **cfg_actor_critic)\n",
+    "# self.agent.actor_critic.eval()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from torchinfo import summary\n",
+    "import torch\n",
+    "from einops import rearrange\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Directly benchmark models"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tokenizer = self.agent.tokenizer\n",
+    "world_model = self.agent.world_model\n",
+    "actor_critic = self.agent.actor_critic\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "4"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "batch_num_samples = cfg.training.world_model.batch_num_samples\n",
+    "sequence_length = cfg.common.sequence_length\n",
+    "sample_from_start = False\n",
+    "# train_dataset = instantiate(cfg.datasets.train)\n",
+    "batch_num_samples\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "batch = self.train_dataset.sample_batch(batch_num_samples, sequence_length, sample_from_start)\n",
+    "batch = {k: v.to(self.device) for k, v in batch.items()}\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CPU times: user 190 ms, sys: 4.87 ms, total: 194 ms\n",
+      "Wall time: 195 ms\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "<src.utils.LossWithIntermediateLosses at 0x7f97ed202a30>"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "%%time\n",
+    "self.agent.world_model.compute_loss(batch,  tokenizer=self.agent.tokenizer)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CPU times: user 9.35 s, sys: 17.3 ms, total: 9.37 s\n",
+      "Wall time: 9.37 s\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "<src.utils.LossWithIntermediateLosses at 0x7f97f274beb0>"
+      ]
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "%%time\n",
+    "# TODO: why is this so slow?\n",
+    "cfg_actor_critic = self.cfg.training.actor_critic\n",
+    "self.agent.actor_critic.compute_loss(batch,  tokenizer=self.agent.tokenizer, world_model=self.agent.world_model, **cfg_actor_critic)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CPU times: user 9.11 s, sys: 15.5 ms, total: 9.13 s\n",
+      "Wall time: 9.13 s\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%time\n",
+    "# is this the slow part... yes. damn\n",
+    "actor_critic.imagine(batch, tokenizer, world_model, horizon=10);\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# # takes 0.1 s, fast\n",
+    "# wm_env = WorldModelEnv(tokenizer, world_model, device)\n",
+    "# wm_env\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([4, 3, 64, 64])\n",
+      "CPU times: user 105 ms, sys: 207 µs, total: 105 ms\n",
+      "Wall time: 105 ms\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%time\n",
+    "# this takes 0.1 seconds and is run 10+ time. So 1 second. Hmm\n",
+    "from src.envs.world_model_env import WorldModelEnv, Categorical\n",
+    "initial_observations = batch['observations']\n",
+    "\n",
+    "# get the right obs\n",
+    "wm_env = WorldModelEnv(self.agent.tokenizer, self.agent.world_model, self.device)\n",
+    "obs = wm_env.reset_from_initial_observations(initial_observations[:, -1])\n",
+    "print(obs.shape)\n",
+    "\n",
+    "\n",
+    "# make sure hidden states are right\n",
+    "self.agent.actor_critic.reset(obs.shape[0])\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import gc\n",
+    "gc.collect()\n",
+    "torch.cuda.empty_cache()\n",
+    "# obs\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CPU times: user 1.6 ms, sys: 309 µs, total: 1.9 ms\n",
+      "Wall time: 1.72 ms\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%time\n",
+    "# 700us\n",
+    "# fast, executed 10+ times\n",
+    "outputs_ac = actor_critic(obs)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([4, 1, 17])"
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "outputs_ac.logits_actions.shape\n",
+    "# action_token.shape\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# %%timeit\n",
+    "# slow! takes 1s, executed 10+ times this is the culprit, not the lstm. hmm\n",
+    "k=3\n",
+    "horizon = 6\n",
+    "action_token = Categorical(logits=outputs_ac.logits_actions).sample()\n",
+    "obs, reward, done, _ = wm_env.step(action_token, should_predict_next_obs=(k < horizon - 1))\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "66.5 ms ± 1.53 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%timeit\n",
+    "# 62ms\n",
+    "# this is the slow part again. no grad and eval don't hepl\n",
+    "outputs_wm = world_model(action_token, past_keys_values=wm_env.keys_values_wm)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "num_steps=1\n",
+    "prev_steps=0\n",
+    "sequences = world_model.embedder(action_token, num_steps, prev_steps) + world_model.pos_emb(prev_steps + torch.arange(num_steps, device=action_token.device))\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "AssertionError",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
+      "\u001b[1;32m/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/notebooks/01_debug_models.ipynb Cell 24\u001b[0m line \u001b[0;36m1\n\u001b[0;32m----> <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/notebooks/01_debug_models.ipynb#X64sZmlsZQ%3D%3D?line=0'>1</a>\u001b[0m get_ipython()\u001b[39m.\u001b[39;49mrun_cell_magic(\u001b[39m'\u001b[39;49m\u001b[39mtimeit\u001b[39;49m\u001b[39m'\u001b[39;49m, \u001b[39m'\u001b[39;49m\u001b[39m'\u001b[39;49m, \u001b[39m\"\u001b[39;49m\u001b[39m# ofc it\u001b[39;49m\u001b[39m'\u001b[39;49m\u001b[39ms the transformer that\u001b[39;49m\u001b[39m'\u001b[39;49m\u001b[39ms slow. I guess we just call it was more than during training\u001b[39;49m\u001b[39m\\n\u001b[39;49;00m\u001b[39mpast_keys_values = wm_env.keys_values_wm\u001b[39;49m\u001b[39m\\n\u001b[39;49;00m\u001b[39mx = world_model.transformer(sequences, past_keys_values)\u001b[39;49m\u001b[39m\\n\u001b[39;49;00m\u001b[39m\"\u001b[39;49m)\n",
+      "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/IPython/core/interactiveshell.py:2515\u001b[0m, in \u001b[0;36mInteractiveShell.run_cell_magic\u001b[0;34m(self, magic_name, line, cell)\u001b[0m\n\u001b[1;32m   2513\u001b[0m \u001b[39mwith\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mbuiltin_trap:\n\u001b[1;32m   2514\u001b[0m     args \u001b[39m=\u001b[39m (magic_arg_s, cell)\n\u001b[0;32m-> 2515\u001b[0m     result \u001b[39m=\u001b[39m fn(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   2517\u001b[0m \u001b[39m# The code below prevents the output from being displayed\u001b[39;00m\n\u001b[1;32m   2518\u001b[0m \u001b[39m# when using magics with decorator @output_can_be_silenced\u001b[39;00m\n\u001b[1;32m   2519\u001b[0m \u001b[39m# when the last Python token in the expression is a ';'.\u001b[39;00m\n\u001b[1;32m   2520\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mgetattr\u001b[39m(fn, magic\u001b[39m.\u001b[39mMAGIC_OUTPUT_CAN_BE_SILENCED, \u001b[39mFalse\u001b[39;00m):\n",
+      "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/IPython/core/magics/execution.py:1189\u001b[0m, in \u001b[0;36mExecutionMagics.timeit\u001b[0;34m(self, line, cell, local_ns)\u001b[0m\n\u001b[1;32m   1186\u001b[0m         \u001b[39mif\u001b[39;00m time_number \u001b[39m>\u001b[39m\u001b[39m=\u001b[39m \u001b[39m0.2\u001b[39m:\n\u001b[1;32m   1187\u001b[0m             \u001b[39mbreak\u001b[39;00m\n\u001b[0;32m-> 1189\u001b[0m all_runs \u001b[39m=\u001b[39m timer\u001b[39m.\u001b[39;49mrepeat(repeat, number)\n\u001b[1;32m   1190\u001b[0m best \u001b[39m=\u001b[39m \u001b[39mmin\u001b[39m(all_runs) \u001b[39m/\u001b[39m number\n\u001b[1;32m   1191\u001b[0m worst \u001b[39m=\u001b[39m \u001b[39mmax\u001b[39m(all_runs) \u001b[39m/\u001b[39m number\n",
+      "File \u001b[0;32m~/miniforge3/lib/python3.9/timeit.py:205\u001b[0m, in \u001b[0;36mTimer.repeat\u001b[0;34m(self, repeat, number)\u001b[0m\n\u001b[1;32m    203\u001b[0m r \u001b[39m=\u001b[39m []\n\u001b[1;32m    204\u001b[0m \u001b[39mfor\u001b[39;00m i \u001b[39min\u001b[39;00m \u001b[39mrange\u001b[39m(repeat):\n\u001b[0;32m--> 205\u001b[0m     t \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mtimeit(number)\n\u001b[1;32m    206\u001b[0m     r\u001b[39m.\u001b[39mappend(t)\n\u001b[1;32m    207\u001b[0m \u001b[39mreturn\u001b[39;00m r\n",
+      "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/IPython/core/magics/execution.py:173\u001b[0m, in \u001b[0;36mTimer.timeit\u001b[0;34m(self, number)\u001b[0m\n\u001b[1;32m    171\u001b[0m gc\u001b[39m.\u001b[39mdisable()\n\u001b[1;32m    172\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> 173\u001b[0m     timing \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49minner(it, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mtimer)\n\u001b[1;32m    174\u001b[0m \u001b[39mfinally\u001b[39;00m:\n\u001b[1;32m    175\u001b[0m     \u001b[39mif\u001b[39;00m gcold:\n",
+      "File \u001b[0;32m<magic-timeit>:3\u001b[0m, in \u001b[0;36minner\u001b[0;34m(_it, _timer)\u001b[0m\n",
+      "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_compiled_call_impl(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)  \u001b[39m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1517\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_call_impl(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
+      "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1522\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1523\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1524\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1525\u001b[0m         \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1526\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1529\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m   1530\u001b[0m     result \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n",
+      "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/models/transformer.py:69\u001b[0m, in \u001b[0;36mTransformer.forward\u001b[0;34m(self, sequences, past_keys_values)\u001b[0m\n\u001b[1;32m     66\u001b[0m     \u001b[39m# k_size = (x.shape[0], x.shape[1], x.shape[1], 1)\u001b[39;00m\n\u001b[1;32m     67\u001b[0m     \u001b[39m# v_size = past_keys_values[0]._v_cache._cache.size()\u001b[39;00m\n\u001b[1;32m     68\u001b[0m     v_size \u001b[39m=\u001b[39m (k_size[\u001b[39m0\u001b[39m], k_size[\u001b[39m1\u001b[39m], x\u001b[39m.\u001b[39mshape[\u001b[39m1\u001b[39m], k_size[\u001b[39m3\u001b[39m])\n\u001b[0;32m---> 69\u001b[0m     past_keys_values[\u001b[39m0\u001b[39;49m]\u001b[39m.\u001b[39;49mupdate(torch\u001b[39m.\u001b[39;49mrand(v_size), torch\u001b[39m.\u001b[39;49mrand(v_size))\n\u001b[1;32m     70\u001b[0m \u001b[39mreturn\u001b[39;00m x\n",
+      "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/models/kv_caching.py:59\u001b[0m, in \u001b[0;36mKVCache.update\u001b[0;34m(self, k, v)\u001b[0m\n\u001b[1;32m     58\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mupdate\u001b[39m(\u001b[39mself\u001b[39m, k: torch\u001b[39m.\u001b[39mTensor, v: torch\u001b[39m.\u001b[39mTensor):\n\u001b[0;32m---> 59\u001b[0m     \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_k_cache\u001b[39m.\u001b[39;49mupdate(k)\n\u001b[1;32m     60\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_v_cache\u001b[39m.\u001b[39mupdate(v)\n",
+      "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/models/kv_caching.py:33\u001b[0m, in \u001b[0;36mCache.update\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     31\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mupdate\u001b[39m(\u001b[39mself\u001b[39m, x: torch\u001b[39m.\u001b[39mTensor) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m     32\u001b[0m     \u001b[39massert\u001b[39;00m (x\u001b[39m.\u001b[39mndim \u001b[39m==\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_cache\u001b[39m.\u001b[39mndim) \u001b[39mand\u001b[39;00m \u001b[39mall\u001b[39m([x\u001b[39m.\u001b[39msize(i) \u001b[39m==\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_cache\u001b[39m.\u001b[39msize(i) \u001b[39mfor\u001b[39;00m i \u001b[39min\u001b[39;00m (\u001b[39m0\u001b[39m, \u001b[39m1\u001b[39m, \u001b[39m3\u001b[39m)])\n\u001b[0;32m---> 33\u001b[0m     \u001b[39massert\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_size \u001b[39m+\u001b[39m x\u001b[39m.\u001b[39msize(\u001b[39m2\u001b[39m) \u001b[39m<\u001b[39m\u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_cache\u001b[39m.\u001b[39mshape[\u001b[39m2\u001b[39m]\n\u001b[1;32m     34\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_cache \u001b[39m=\u001b[39m AssignWithoutInplaceCheck\u001b[39m.\u001b[39mapply(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_cache, x, \u001b[39m2\u001b[39m, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_size, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_size \u001b[39m+\u001b[39m x\u001b[39m.\u001b[39msize(\u001b[39m2\u001b[39m))\n\u001b[1;32m     35\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_size \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m x\u001b[39m.\u001b[39msize(\u001b[39m2\u001b[39m)\n",
+      "\u001b[0;31mAssertionError\u001b[0m: "
+     ]
+    }
+   ],
+   "source": [
+    "%%timeit\n",
+    "# ofc it's the transformer that's slow. I guess we just call it was more than during training\n",
+    "past_keys_values = wm_env.keys_values_wm\n",
+    "x = world_model.transformer(sequences, past_keys_values)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# past_keys_values = wm_env.keys_values_wm\n",
+    "# x = world_model.transformer(sequences, past_keys_values)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%timeit\n",
+    "# ofc it's the transformer that's slow. I guess we just call it was more than during training\n",
+    "past_keys_values = wm_env.keys_values_wm\n",
+    "x = world_model.transformer(sequences, past_keys_values)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "logits_observations = world_model.head_observations(x, num_steps=num_steps, prev_steps=prev_steps)\n",
+    "logits_rewards = world_model.head_rewards(x, num_steps=num_steps, prev_steps=prev_steps)\n",
+    "logits_ends = world_model.head_ends(x, num_steps=num_steps, prev_steps=prev_steps)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Torchinfo model sizes\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "observations = self.agent.tokenizer.preprocess_input(rearrange(batch['observations'], 'b t c h w -> (b t) c h w'))\n",
+    "# z, z_quantized, reconstructions = self.agent.tokenizer(observations, should_preprocess=False, should_postprocess=False)\n",
+    "summary(self.agent.tokenizer, input_data=observations)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    obs_tokens = self.agent.tokenizer.encode(batch['observations'], should_preprocess=True).tokens  # (BL, K)\n",
+    "\n",
+    "act_tokens = rearrange(batch['actions'], 'b l -> b l 1')\n",
+    "tokens = rearrange(torch.cat((obs_tokens, act_tokens), dim=2), 'b l k1 -> b (l k1)')  # \n",
+    "\n",
+    "summary(self.agent.world_model, input_data=tokens)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from src.envs.world_model_env import WorldModelEnv\n",
+    "initial_observations = batch['observations']\n",
+    "\n",
+    "# get the right obs\n",
+    "wm_env = WorldModelEnv(self.agent.tokenizer, self.agent.world_model, self.device)\n",
+    "obs = wm_env.reset_from_initial_observations(initial_observations[:, -1])\n",
+    "obs.shape\n",
+    "\n",
+    "\n",
+    "# make sure hidden states are right\n",
+    "self.agent.actor_critic.reset(obs.shape[0])\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "\n",
+    "from torchinfo import summary\n",
+    "summary(self.agent.actor_critic, input_data=obs)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Debug env\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import minihack\n",
+    "env = gym.make(\"MiniHack-River-v0\", observation_keys=(\"pixel_crop\", \"pixel\", 'blstats', 'message'))\n",
+    "env.reset() # each reset generates a new environment instance\n",
+    "obs, reward, end, info = env.step(1)  # move agent '@' north\n",
+    "print(obs['pixel_crop'].shape)\n",
+    "plt.imshow(obs['pixel_crop'])\n",
+    "plt.show()\n",
+    "\n",
+    "print(obs['pixel'].shape)\n",
+    "plt.imshow(obs['pixel'])\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# # plt.imshow(obs['glyphs_crop'])\n",
+    "# obs['glyphs_crop'].shape\n",
+    "# obs['blstats']\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import minihack\n",
+    "env = gym.make(\"MiniHack-Room-5x5-v0\", observation_keys=(\"pixel_crop\", \"pixel\", 'blstats', 'message'))\n",
+    "env.reset() # each reset generates a new environment instance\n",
+    "obs, reward, end, info = env.step(1)  # move agent '@' north\n",
+    "print(obs['pixel_crop'].shape)\n",
+    "plt.imshow(obs['pixel_crop'])\n",
+    "plt.show()\n",
+    "\n",
+    "print(obs['pixel'].shape)\n",
+    "plt.imshow(obs['pixel'])\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import minihack\n",
+    "import crafter\n",
+    "env = gym.make(\"CrafterReward-v1\")\n",
+    "env.reset() # each reset generates a new environment instance\n",
+    "obs, reward, end, info = env.step(1)  # move agent '@' north\n",
+    "print(obs.shape)\n",
+    "plt.imshow(obs)\n",
+    "plt.show()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
@@ -1,30 +1,35 @@
 [tool.poetry]
 name = "src"
 version = "0.1.0"
-description = "Trying to use a AdaVAE (an LLM VAE) as a world model in an RL agent in a text game"
+description = ""
 authors = ["wassname <git@wassname.org>"]
-license = "MIT"
 readme = "README.md"

 [tool.poetry.dependencies]
-python = ">=3.11,<3.13"
+python = ">=3.9,<3.13"
 torch = {version = "^2.1.0+cu118", source = "pytorch"}
-simple-parsing = "^0.1.4"
-tqdm = "^4.66.1"
-numpy = "^1.26.1"
-pandas = "^2.1.1"
-lightning = "^2.1.0"
 matplotlib = "^3.8.0"
 loguru = "^0.7.2"
-einops = "^0.7.0"
-scikit-learn = "^1.3.1"
-pytorch-optimizer = "^2.12.0"
+einops = "^0.3.1"
 torchinfo = "^1.8.0"
 accelerate = "^0.24.1"
-datasets = "^2.14.6"
 peft = "^0.5.0"
 bitsandbytes = {url = "https://github.com/TimDettmers/bitsandbytes/releases/download/0.41.0/bitsandbytes-0.41.0-py3-none-any.whl"}
 transformers = "4.34.0"
+tqdm = "^4.66.1"
+wandb = "^0.12.6"
+ale-py = "^0.7.4"
+pygame = "^2.5.2"
+psutil = "^5.9.6"
+protobuf = "^3.10.0"
+opencv-python = "^4.8.1.78"
+hydra-core = "^1.3.2"
+torchvision = "^0.16.0"
+numpy = ">=1.18.0"
+gym = {version = "0.22.0", extras = ["accept-rom-license", "atari"]}
+scipy = "^1.11.3"
+crafter = "^1.8.2"
+minihack = "^0.1.5"

 [[tool.poetry.source]]
 name = "pytorch"
@@ -0,0 +1,407 @@
+# 2023-11-12 13:17:35
+
+Try IRIs but with pretrained transformer with LoRA adapter
+
+- [x] first can I run it yes with a 1/2 batch size
+- [ ] then can I add 3B with adapter...
+
+```sh
+poetry install
+. ./.venv/bin/activate
+python src/main.py env.train.id=BreakoutNoFrameskip-v4 common.device=cuda:0 wandb.mode=offline
+
+# or for quick debug
+WANDB_MODE=disabled python -m pdb src/main.py env.train.id=BreakoutNoFrameskip-v4
+```
+
+
+```sh
+# TODO use this code to load a transformer, and other code from my bigvae repo https://github.com/wassname/bigvae_wm
+def load_model(config, device='cuda'):
+    tokenizer = AutoTokenizer.from_pretrained(config.model_name, trust_remote_code=True)
+    tokenizer.padding_side = "left"
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+    )
+    base_model = AutoModelForCausalLM.from_pretrained(
+        config.model_name,
+        device_map={"": device},
+        quantization_config=bnb_config,
+        torch_dtype=torch.bfloat16, 
+        trust_remote_code=True
+    )
+    peft_config = peft.LoraConfig(
+        peft.TaskType.CAUSAL_LM,
+        inference_mode=False,
+        r=config.rank,
+        lora_alpha=8,
+        lora_dropout=config.dropout,
+        target_modules=[
+            "self_attn.q_proj",
+            "self_attn.k_proj",
+            "self_attn.v_proj",
+            "self_attn.o_proj",
+            "mlp.gate_proj",
+            "mlp.up_proj",
+            "mlp.down_proj",
+        ],
+    )
+    base_model_peft = peft.get_peft_model(base_model, peft_config)
+    vae_model = BigVAE(
+        base_model_peft, device, peft_config, z_dim=config.z_dim,
+    )
+    if config.start_from:
+        vae_model.load_pretrained(config.start_from)
+    base_model_peft.requires_grad_(False)
+    vae_model.vae_head.requires_grad_(False)
+    vae_model.vae_head.w_d.requires_grad_()
+    router = BigVAERouter(base_model_peft, vae_model, device, peft_config)
+    if config.start_from:
+        router.load_pretrained(config.start_from, is_trainable=True)
+    print(router.model.print_trainable_parameters())
+    router.model.set_adapter("router")
+```
+
+Debugging:
+    batch['observations'].shape
+    torch.Size([16, 20, 3, 64, 64])
+
+    obs_tokens.shape
+    torch.Size([16, 20, 16])
+
+    https://vscode.dev/github/wassname/iris_bigvae/blob/just_llms2/src/models/world_model.py#L105
+    tokens
+    tensor([[222, 222, 222,  ..., 409,  55,   2],
+            [222, 222, 222,  ..., 409, 139,   1],
+            [222, 222, 222,  ..., 168, 190,   3],
+            ...,
+            [222, 222, 222,  ..., 168,  55,   0],
+            [222, 222, 222,  ..., 237, 190,   3],
+            [222, 222, 222,  ..., 168,  55,   0]], device='cuda:0')
+    tokens.shape
+    torch.Size([16, 340])
+    where 16 is the batch size. 340 is the step size?. actions was 16,20 int
+
+    tokens.shape int
+    torch.Size([16, 340])
+
+    sequences.shape float32
+    torch.Size([16, 340, 256])
+
+    transfrmer
+    x.shape
+    torch.Size([16, 340, 256])
+
+# 2023-11-12 16:58:37
+
+So I got it training, but during imagination it passes in a single token with no past steps. But the slicer seems to need at least on block? And so I get none?
+
+hmm it's because num_kept_tokens is 16 not 1. So there should be a whole block passed in ?
+
+wait apparently it's also a problem in the normal repo.... I confuse! maybe it's my config! maybe I need >larger than block size. nope
+
+hmm it still happens in the original repo with my debug params. maybe it's my debug params
+
+
+... trying a full run without my debug params...
+
+note trains.world_model.batch_num_samples:4 fill 20GB gpu ram for the 3b stability ai llm
+
+ok even with a full run I get the error. I think it's a bug in the original repo. I'll try to debug it there.
+
+    Epoch 51 / 600
+
+    Experience collection (train_dataset): 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:03<00:00, 59.91it/s]
+    Training tokenizer: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:17<00:00, 11.53it/s]
+    Training world_model: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [02:11<00:00,  1.53it/s]
+    Training actor_critic:   0%|                                                                                                                                                                                     | 0/200 [00:00<?, ?it/s]
+    Error executing job with overrides: ['env.train.id=BreakoutNoFrameskip-v4', 'common.device=cuda:0', 'wandb.mode=offline']
+    Traceback (most recent call last):
+    File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/main.py", line 10, in main
+        trainer.run()
+    File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/trainer.py", line 111, in run
+        to_log += self.train_agent(epoch)
+    File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/trainer.py", line 146, in train_agent
+        metrics_actor_critic = self.train_component(self.agent.actor_critic, self.optimizer_actor_critic, sequence_length=1 + self.cfg.training.actor_critic.burn_in, sample_from_start=False, tokenizer=self.agent.tokenizer, world_model=self.agent.world_model, **cfg_actor_critic)
+    File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/trainer.py", line 161, in train_component
+        losses = component.compute_loss(batch, **kwargs_loss) / grad_acc_steps
+    File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/models/actor_critic.py", line 102, in compute_loss
+        outputs = self.imagine(batch, tokenizer, world_model, horizon=imagine_horizon)
+    File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/models/actor_critic.py", line 149, in imagine
+        obs, reward, done, _ = wm_env.step(action_token, should_predict_next_obs=(k < horizon - 1))
+    File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
+        return func(*args, **kwargs)
+    File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/envs/world_model_env.py", line 75, in step
+        reward = Categorical(logits=outputs_wm.logits_rewards).sample().float().cpu().numpy().reshape(-1) - 1   # (B,)
+    File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/torch/distributions/categorical.py", line 70, in __init__
+        super().__init__(batch_shape, validate_args=validate_args)
+    File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/torch/distributions/distribution.py", line 66, in __init__
+        valid = constraint.check(value)
+    File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/torch/distributions/constraints.py", line 226, in check
+        result = result.reshape(
+    RuntimeError: cannot reshape tensor of 0 elements into shape [8, 0, -1] because the unspecified dimension size -1 can be any value and is ambiguous
+
+Oh maybe it's because we don't keep track of KV cache, but it's actually used to track number of steps!!
+
+# 2023-11-13 20:11:51
+
+I go it working byt ut takes 30 seconds for one one actor critic batch, werird
+
+Experience collection (train_dataset): 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:03<00:00, 60.45it/s]
+Training tokenizer: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:17<00:00, 11.26it/s]
+Training world_model: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [02:12<00:00,  1.51it/s]
+Training actor_critic:  82%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▍                             | 165/200 [1:01:43<13:24, 22.99s/it]
+
+
+hm maybe it's just the face it has to backprop throguh the whole LLM :( damn... is there another way to train it? Daym. How many params did the original have?
+
+well running eval on the transformer brought it down from 100sec to 60, but it's still huge. 
+
+But then why is the model training fast? It makes not sense
+ 
+# 2023-11-16 12:54:48
+
+Why is agent so slow? Lets find out
+- look at diagram
+- look at train_agent
+  - to tokenizer.compute_loss is just tokenizer
+  - world_model.compute_loss user tokenizer with no grad
+  - actor_critic? takes an hour!!
+    - imagine (with grad?)
+      - x20 = horizon
+        - self(obs)
+        - WorldModelEnv.step this has no grad!
+          - transformer
+          - tokenizer with no grad
+    - compute_lambda_returns with no grad
+
+
+So changes:
+- the world model step always had no grad!
+- I just made the lstm smaller and the horizon smaller
+- from 1h to 3m. Reasonable.
+
+
+Experiment: 
+- try no grad on the model? ok it now takes 20 minutes to train... still slow
+with a smaller lstm and only 10 steos ut tajes 8 mins,
+
+![](img/2023-11-16-13-01-11.png)
+
+Experience collection (train_dataset): 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:05<00:00, 36.11it/s]
+Training tokenizer: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:55<00:00,  3.61it/s]
+Training world_model: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [01:12<00:00,  2.76it/s]
+Training tokenizer: 55sec
+Training world_model 72 sec
+train actor_critic 3min. It looks like it scales with lstm size!
+
+
+new changes 10mins
+lets try no lstm?
+Right now it will take 41 hours for on epoch lol
+
+
+
+ram during stages:
+- actor critic 20G/24
+
+
+how big does my actor critic need to be?
+- IRIS: large 512 lstm on 64,64,3 obs
+  - We ran our experiments with 8 Nvidia A100 40GB GPUs. With two Atari environments running on the same GPU, training takes around 7 days, resulting in an average of 3.5 days per environment.
+- twm: mlp 512
+
+
+How long to train?
+`600*10//6/24` = 41 days
+- 600 epochs * 10 minutes / 6 to get hours, 24 to get days
+
+# 2023-11-17 07:59:44
+
+so I've got it working with  these times. But maybe it's too small
+
+Epoch 148 / 600
+
+Experience collection (train_dataset): 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:10<00:00, 19.25it/s]
+Training tokenizer: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:59<00:00,  3.34it/s]
+Training world_model: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [01:15<00:00,  2.64it/s]
+Training actor_critic: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 20/20 [03:20<00:00, 10.05s/it]
+
+
+
+- what about resume? oh we seem to have that although the code doesn't make sense https://hydra.cc/docs/tutorials/basic/running_your_app/working_directory/ https://hydra.cc/docs/1.2/upgrades/1.1_to_1.2/changes_to_job_working_dir/ see eval.py
+- [ ] but it's still too damn slow. what about bfloat16? using auto case?
+- why does it take so long? it would be nice to have a reproduction notebook
+- also the model might be to small now....
+
+how to play
+```sh
+cd outputs/2023-11-17/07-59-44
+python scripts/play.sh
+```
+
+
+## Envs
+
+tl:dr just use pong or breakout or crafter (1m steps)
+
+for steps see [crafter paper](https://arxiv.org/pdf/2109.06780.pdf)
+
+Nethack learning env. What's the obs size? 21x79 of glyphs (5991 possibilities) and 21 dim of stats
+- they use an lstm of 128. 5 layer conv
+- requires 1B steps
+- 
+atari:
+- reqs 200M stpes
+
+progcen:
+- 200M steps
+
+minihack:
+- 2M steps for room 5xt
+- but needs editing to be atari compatible. e.g. 336 × 1264 × 3 pixels
+- pixel_crop 64,64,3 or 9x9 crop works!
+- lstm 256
+- The training on MiniHack’s Room-5x5 task for two million timesteps using our IMPALA baseline takes approximately 4:30 minutes (r
+
+crafter
+- reqs 1M steps
+- "All agents trained for 1M environment steps in under 24 hours on a single GPU and we repeated the training for 10 random seeds per method. The training reward curves are included in Appendix "
+
+
+
+
+from https://arxiv.org/pdf/2111.09794.pdf
+There are several PCG state-varying gridworld environments (
+- [MiniGrid](https://minigrid.farama.org/environments/minigrid/), ~~- BabyAI~~
+- Crafter,
+- 2019 [Rogue-gym,](https://github.com/kngwyu/rogue-gym) 
+- 2020 MarsExplorer, maxe exploration. 1M steps
+- NLE, 
+- MiniHack; 
+- [gym\_nethack](http://campbelljc.com/research/gym_nethack/)
+- 2018 [rogueinabox](https://github.com/rogueinabox/rogueinabox)
+- [rogue-gym](https://github.com/kngwyu/rogue-gym)
+- [MiniGrid](https://github.com/maximecb/gym-minigrid)
+- 2019 [CoinRun](https://github.com/openai/coinrun) no traction or maintanance
+- [MineRL](http://minerl.io/docs)
+- [Project Malmo](https://www.microsoft.com/en-us/research/project/project-malmo/) miencraft
+- [OpenAI Procgen Benchmark](https://openai.com/blog/procgen-benchmark/) 200M steps
+- 2020 [Obstacle Tower](https://github.com/Unity-Technologies/obstacle-tower-env) - 3d slow
+
+non-PCG observation-varying continuous control environments
+- (RoboSuite, DMC-Remastered, DMC-GB, DCS, KitchenShift, NaturalEnvs MuJoCo; Fan
+et al., 2021; Grigsby & Qi, 2020; Hansen & Wang, 2021; Stone et al., 2021; Xing et al., 2021a;
+Zhang et al., 2018a), and multi-task continuous control benchmarks which could be adapted
+to ZSG (CausalWorld, RLBench, Meta-world; Ahmed et al., 2020; James et al., 2019a; Yu
+et al., 2019).
+
+## investigating model slowness
+
+So it's all just that using the transformer to imagine takes almost 0.1s. but it's run so many more times than during training. All my ideas to speed it up don't work.
+
+- [x] eval. no grad
+- [x] remove the call for adapter, causal mask each time?
+- [ ] lower rank?
+
+
+Ok so it's all just the
+- rollout, controller by max block size. 10x
+- the fact that actor_critic can use a larger batch, therefore 4-8x more samples
+- for each one it imagines 2
+
+# 2023-11-18 06:17:55
+
+It trained overnight, now I would like to view a replay
+
+Hmm "delta-IRIS" ∆-IRIS 
+https://openreview.net/forum?id=o8IDoZggqO
+∆-IRIS encodes
+new frames by attending to the ongoing trajectory, effec-
+tively describing deltas between timesteps.
+This new ap-
+proach drastically reduces the number of tokens to encode
+frames, since they are not encoded independently as in IRIS.
+In the Crafter benchmark (Hafner, 2022), ∆-IRIS unlocks
+16 out of 22 objectives at the 10M frames mark
+
+
+# 2023-11-18 16:16:16
+
+Why is it not learning? It's because the dynamics model is total BS!!!
+
+- [ ] Well lets try training it for longer then. It's cheap to train so..
+- [ ] also maybe train tokenizer and model together? I have a lot of frozen layers, including the embeddings... so might be better
+  - [ ] oh no we do have an unforzen embedder before the transformer or more layers
+  - maybe I need a higher rank lora? after all I'm changing a lot from text tokens
+  - maybe no tokens, bypass to embedder?
+
+
+# 2023-11-19 06:45:34
+
+So I tried just trainign the world model for 200 epochs. And with a post_embedding layer. It helped the flickering. But not enougth to actually go for the obvious local minima of the next state equals the last
+  
+idea
+- bypass embedding?, but wait dreamerv3 needed quant z... 
+  - yes I am bypassing it by passing in the input_embeds... but maybe I shouldn't
+- [x] use same embedding everywhere. e.g. model embedding in encoder decoder?
+  - Our embedings is (embed_tokens): Embedding(32000, 2048). So we would need to encode to 32000!
+
+
+ok we need to freeze it, and change dtype
+OK it seems slightly better yay! Lets train it overnight and see
+
+next idea is to the delta-IRIS thing where the tokens only have to encode the diff(obs)
+
+# 2023-11-19 16:50:45 
+
+Seems to be working! Now let's plan delta-IRIS
+
+
+So IRIS has
+- Encoder $E(x_0, a_0) = t_0$
+  ```py
+  obs_tokens = self.tokenizer.encode(observations, should_preprocess=True).tokens    # (B, C, H, W) -> (B, K)
+  ```
+- Embed $Emb(t_0) = z_0$
+    ```py
+    embedded_tokens = self.tokenizer.embedding(self.obs_tokens)     # (B, K, E)
+    z = rearrange(embedded_tokens, 'b (h w) e -> b e h w', h=int(np.sqrt(self.num_observations_tokens)))
+    ```
+- Dynamics $D(z_0, a_0) = z_1$
+    ```py
+    outputs_wm = self.world_model(tokenRedmond AI, past_keys_values=self.keys_values_wm)
+    ```
+- Decoder $D(z_0, a_0) = x_1$
+    ```py
+    rec = self.tokenizer.decode(z, should_postprocess=True)         # (B, C, H, W)
+    ```
+
+
+but we have tokens vs z
+
+Questions:
+- wait why are we just passing in "action_token" to the transformer and not obs? that must have obs in it right... right??? confirm
+- in iris-delta how did they pass everything in? I guess obs_prev was tokenized too? I think the slices are annoying so maybe I should just pass things seperatly
+
+# 2023-11-24 10:56:40
+
+If I unfreeze the whole transformer, it seem to learn the most obvious dynamics (the next latent space is the same as the last). 
+
+To summarize
+- with Qlora it didn't learn that
+- with unfrozen head it didn't
+- when training transformer and obs embedding together it did not (frozen llm embeddings)
+
+
+no it didn't work with tokenizer sep hmm
+
+Oh it did with whole transfrmer and tokenizer at same time https://wandb.ai/wassname/iris/runs/w7lvs4gi?workspace=user-wassname
+wandb:           world_model/eval/loss_obs ▇█▃▄▄▄▄▃▃▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:       world_model/eval/loss_rewards ▁█▃█▇▅▇▁▅▆▂▄▃▇▆▃▄▇▃▄▂▄▃▂▃▄▃
+wandb:         world_model/eval/total_loss ▂█▃██▅▇▁▅▆▁▃▂█▇▃▃▆▃▄▂▃▃▂▂▄▂
@@ -14,22 +14,23 @@ while [ "$1" != "" ]; do
            ;;
        -h | --header )
            header=1
-            ;;
+            ;; # adds banner with env metadata like action
        -r | --reconstruction )
            reconstruction=1
-            ;;
+            ;; # 3 panes [original_obs, resized_obs, reconstructed], doesn't do anything if any of -w -a or -e are set. shows quality of encoder decoder
        -s | --save-mode )
            save_mode=1
-            ;;
+            ;; # lets you save the episode to mp4
        -a | --agent-world-model )
            mode="agent_in_world_model"
-            ;;
+            ;; # the agent plays in the world model env, shows the quality of the dynamics model
        -e | --episode )
            mode="episode_replay"
-            ;;
+            ;; # replay train, test, or imagined episodes. shows quality of dynamics model
+            # this is quick low resource way to check the dynamics model and agent while training
        -w | --world-model )
            mode="play_in_world_model"
-            ;;
+            ;; # human plays in world model
        * )
            echo Invalid usage : $1
            exit 1
@@ -4,10 +4,10 @@ import torch
 from torch.distributions.categorical import Categorical
 import torch.nn as nn

-from models.actor_critic import ActorCritic
-from models.tokenizer import Tokenizer
-from models.world_model import WorldModel
-from utils import extract_state_dict
+from src.models.actor_critic import ActorCritic
+from src.models.tokenizer import Tokenizer
+from src.models.world_model import WorldModel
+from src.utils import extract_state_dict


 class Agent(nn.Module):
@@ -34,4 +34,5 @@ class Agent(nn.Module):
        input_ac = obs if self.actor_critic.use_original_obs else torch.clamp(self.tokenizer.encode_decode(obs, should_preprocess=True, should_postprocess=True), 0, 1)
        logits_actions = self.actor_critic(input_ac).logits_actions[:, -1] / temperature
        act_token = Categorical(logits=logits_actions).sample() if should_sample else logits_actions.argmax(dim=-1)
+        # FIXME, is this really just an action and doesn't have an obs in?
        return act_token
@@ -8,11 +8,11 @@ import torch
 from tqdm import tqdm
 import wandb

-from agent import Agent
-from dataset import EpisodesDataset
-from envs import SingleProcessEnv, MultiProcessEnv
-from episode import Episode
-from utils import EpisodeDirManager, RandomHeuristic
+from src.agent import Agent
+from src.dataset import EpisodesDataset
+from src.envs import SingleProcessEnv, MultiProcessEnv
+from src.episode import Episode
+from src.utils import EpisodeDirManager, RandomHeuristic


 class Collector:
@@ -7,7 +7,7 @@ from typing import Dict, List, Optional, Tuple
 import psutil
 import torch

-from episode import Episode
+from src.episode import Episode

 Batch = Dict[str, torch.Tensor]

@@ -1,4 +1,4 @@
 from .multi_process_env import MultiProcessEnv
-from .wrappers import make_atari, ResizeObsWrapper
+from .wrappers import make_atari, make_crafter, make_env, ResizeObsWrapper
 from .single_process_env import SingleProcessEnv
 from .world_model_env import WorldModelEnv
@@ -65,12 +65,13 @@ class WorldModelEnv:

        token = action.clone().detach() if isinstance(action, torch.Tensor) else torch.tensor(action, dtype=torch.long)
        token = token.reshape(-1, 1).to(self.device)  # (B, 1)
+        

        for k in range(num_passes):  # assumption that there is only one action token.

+            # FIXME: hold on we are ONLY passing in the action token! should it not be obs too
            outputs_wm = self.world_model(token, past_keys_values=self.keys_values_wm)
            output_sequence.append(outputs_wm.output_sequence)
-
            if k == 0:
                reward = Categorical(logits=outputs_wm.logits_rewards).sample().float().cpu().numpy().reshape(-1) - 1   # (B,)
                done = Categorical(logits=outputs_wm.logits_ends).sample().cpu().numpy().astype(bool).reshape(-1)       # (B,)
@@ -7,11 +7,22 @@ from typing import Tuple
 import gym
 import numpy as np
 from PIL import Image
+import crafter
+
+
+def make_env(id, size=64, max_episode_steps=None, noop_max=30, frame_skip=4, done_on_life_loss=False, clip_reward=False):
+    if id.startswith('Crafter'):
+        return make_crafter(id, size=size, max_episode_steps=max_episode_steps, done_on_life_loss=done_on_life_loss)
+    if id.startswith('MiniHack'):
+        return make_minihack(size=size, max_episode_steps=max_episode_steps, done_on_life_loss=done_on_life_loss)
+    else:
+        return make_atari(id, size, max_episode_steps, noop_max, frame_skip, done_on_life_loss, clip_reward)


 def make_atari(id, size=64, max_episode_steps=None, noop_max=30, frame_skip=4, done_on_life_loss=False, clip_reward=False):
    env = gym.make(id)
-    assert 'NoFrameskip' in env.spec.id or 'Frameskip' not in env.spec
+    print(env.spec)
+    assert 'NoFrameskip' in env.spec.id or 'Frameskip' not in str(env.spec)
    env = ResizeObsWrapper(env, (size, size))
    if clip_reward:
        env = RewardClippingWrapper(env)
@@ -24,6 +35,26 @@ def make_atari(id, size=64, max_episode_steps=None, noop_max=30, frame_skip=4, d
        env = EpisodicLifeEnv(env)
    return env

+def make_crafter(id, size=64, max_episode_steps=None, done_on_life_loss=False):
+    # https://github.com/danijar/dreamerv2/blob/07d906e9c4322c6fc2cd6ed23e247ccd6b7c8c41/dreamerv2/common/envs.py#L242
+    # https://github.com/footoredo/torchbeast/blob/12939569cc46b6a8616e4c25b138d97248cc8581/torchbeast/atari_wrappers.py#L301
+    env = gym.make(id)
+    env = ResizeObsWrapper(env, (size, size))
+    return env
+
+
+def make_minihack(id, size=64, max_episode_steps=None, noop_max=30, frame_skip=4, done_on_life_loss=False, clip_reward=False):
+    # https://github.com/facebookresearch/minihack/blob/47065748f04714c49ba5b52fb74d166228c7acc1/minihack/agent/common/envs/wrapper.py#L117
+    # https://github.com/roger-creus/SOFE/blob/5551a115a9c7e1d632cf6996bf5dcabde59cdcc5/e3b/minihack/torchbeast/src/utils.py#L110
+    env = gym.make(id,
+                   # https://minihack.readthedocs.io/en/latest/getting-started/observation_spaces.html
+                   observation_keys=("pixel_crop"),
+                #    obs_crop_h=9,
+                #    obs_crop_w=9,
+                    )
+    env = ResizeObsWrapper(env, (size, size))
+    return env
+

 class ResizeObsWrapper(gym.ObservationWrapper):
    def __init__(self, env: gym.Env, size: Tuple[int, int]) -> None:
@@ -64,7 +95,7 @@ class NoopResetEnv(gym.Wrapper):
        if self.override_num_noops is not None:
            noops = self.override_num_noops
        else:
-            noops = self.unwrapped.np_random.randint(1, self.noop_max + 1)
+            noops = self.unwrapped.np_random.integers(1, self.noop_max + 1)
        assert noops > 0
        obs = None
        for _ in range(noops):
@@ -4,14 +4,14 @@ from PIL import Image
 import torch
 from torchvision.transforms.functional import InterpolationMode, resize

-from agent import Agent
-from envs import SingleProcessEnv, WorldModelEnv
-from game.keymap import get_keymap_and_action_names
+from src.agent import Agent
+from src.envs import SingleProcessEnv, WorldModelEnv
+from src.game.keymap import get_keymap_and_action_names


 class AgentEnv:
    def __init__(self, agent: Agent, env: SingleProcessEnv, keymap_name: str, do_reconstruction: bool) -> None:
-        assert isinstance(env, SingleProcessEnv) or isinstance(env, WorldModelEnv)
+        assert isinstance(env, SingleProcessEnv) or isinstance(env, WorldModelEnv), f"{env}"
        self.agent = agent
        self.env = env
        _, self.action_names = get_keymap_and_action_names(keymap_name)
@@ -12,6 +12,10 @@ def get_keymap_and_action_names(name):

    if name == 'atari':
        return ATARI_KEYMAP, ATARI_ACTION_NAMES
+    
+    if name == 'atari/CrafterReward-v1':
+        env_id = name.split('atari/')[1]
+        return CRAFTER_KEYMAP, gym.make(env_id).action_names 

    assert name.startswith('atari/')
    env_id = name.split('atari/')[1]
@@ -100,4 +104,25 @@ EMPTY_ACTION_NAMES = [
 ]

 EMPTY_KEYMAP = {
-}
+}
+
+CRAFTER_KEYMAP = {
+    pygame.K_a: 1,
+    pygame.K_d: 2,
+    pygame.K_w: 3,
+    pygame.K_s: 4,
+    pygame.K_SPACE: 5,
+    pygame.K_TAB: 6,
+
+    pygame.K_r: 7,
+    pygame.K_t: 8,
+    pygame.K_f: 9,
+    pygame.K_p: 10,
+
+    pygame.K_1: 11,
+    pygame.K_2: 12,
+    pygame.K_3: 13,
+    pygame.K_4: 14,
+    pygame.K_5: 15,
+    pygame.K_6: 16,
+}
@@ -2,7 +2,9 @@ import hydra
 from omegaconf import DictConfig

 from trainer import Trainer
-
+from loguru import logger
+import sys
+logger.add(sys.stderr, format="{time} {level} {message}", filter="my_module", level="INFO")

@hydra.main(config_path="../config", config_name="trainer")
 def main(cfg: DictConfig):
@@ -10,11 +10,11 @@ import torch.nn as nn
 import torch.nn.functional as F
 from tqdm import tqdm

-from dataset import Batch
-from envs.world_model_env import WorldModelEnv
-from models.tokenizer import Tokenizer
-from models.world_model import WorldModel
-from utils import compute_lambda_returns, LossWithIntermediateLosses
+from src.dataset import Batch
+from src.envs.world_model_env import WorldModelEnv
+from src.models.tokenizer import Tokenizer
+from src.models.world_model import WorldModel
+from src.utils import compute_lambda_returns, LossWithIntermediateLosses


@dataclass
@@ -34,24 +34,26 @@ class ImagineOutput:


 class ActorCritic(nn.Module):
-    def __init__(self, act_vocab_size, use_original_obs: bool = False) -> None:
+    def __init__(self, act_vocab_size, use_original_obs: bool = False, lstm_dim = 16) -> None:
        super().__init__()
+        shrink = 1
+        s = 1
        self.use_original_obs = use_original_obs
-        self.conv1 = nn.Conv2d(3, 32, 3, stride=1, padding=1)
+        self.conv1 = nn.Conv2d(3, 32//s, 3, stride=1, padding=1)
        self.maxp1 = nn.MaxPool2d(2, 2)
-        self.conv2 = nn.Conv2d(32, 32, 3, stride=1, padding=1)
+        self.conv2 = nn.Conv2d(32//s, 32//s, 3, stride=1, padding=1)
        self.maxp2 = nn.MaxPool2d(2, 2)
-        self.conv3 = nn.Conv2d(32, 64, 3, stride=1, padding=1)
+        self.conv3 = nn.Conv2d(32//s, 64//s, 3, stride=1, padding=1)
        self.maxp3 = nn.MaxPool2d(2, 2)
-        self.conv4 = nn.Conv2d(64, 64, 3, stride=1, padding=1)
+        self.conv4 = nn.Conv2d(64//s, 64//shrink, 3, stride=1, padding=1)
        self.maxp4 = nn.MaxPool2d(2, 2)

-        self.lstm_dim = 512
-        self.lstm = nn.LSTMCell(1024, self.lstm_dim)
+        self.lstm_dim = lstm_dim
+        self.lstm = nn.LSTMCell(1024//shrink, self.lstm_dim)
        self.hx, self.cx = None, None

-        self.critic_linear = nn.Linear(512, 1)
-        self.actor_linear = nn.Linear(512, act_vocab_size)
+        self.critic_linear = nn.Linear(self.lstm_dim, 1)
+        self.actor_linear = nn.Linear(self.lstm_dim, act_vocab_size)

    def __repr__(self) -> str:
        return "actor_critic"
@@ -85,7 +87,7 @@ class ActorCritic(nn.Module):
        x = F.relu(self.maxp2(self.conv2(x)))
        x = F.relu(self.maxp3(self.conv3(x)))
        x = F.relu(self.maxp4(self.conv4(x)))
-        x = torch.flatten(x, start_dim=1)
+        x = torch.flatten(x, start_dim=1) # [b=32, 64//shrink, 4, 4]

        if mask_padding is None:
            self.hx, self.cx = self.lstm(x, (self.hx, self.cx))
@@ -146,6 +148,8 @@ class ActorCritic(nn.Module):

            outputs_ac = self(obs)
            action_token = Categorical(logits=outputs_ac.logits_actions).sample()
+            
+            # FIXME shouldn't we pass in obs too?            
            obs, reward, done, _ = wm_env.step(action_token, should_predict_next_obs=(k < horizon - 1))

            all_actions.append(action_token)
@@ -31,7 +31,8 @@ class Head(Slicer):
        self.head_module = head_module

    def forward(self, x: torch.Tensor, num_steps: int, prev_steps: int) -> torch.Tensor:
-        x_sliced = x[:, self.compute_slice(num_steps, prev_steps)]  # x is (B, T, E)
+        s = self.compute_slice(num_steps, prev_steps)
+        x_sliced = x[:, s]  # x is (B, T, E)
        return self.head_module(x_sliced)


@@ -9,10 +9,10 @@ from einops import rearrange
 import torch
 import torch.nn as nn

-from dataset import Batch
+from src.dataset import Batch
 from .lpips import LPIPS
 from .nets import Encoder, Decoder
-from utils import LossWithIntermediateLosses
+from src.utils import LossWithIntermediateLosses


@dataclass
@@ -23,15 +23,15 @@ class TokenizerEncoderOutput:


 class Tokenizer(nn.Module):
-    def __init__(self, vocab_size: int, embed_dim: int, encoder: Encoder, decoder: Decoder, with_lpips: bool = True) -> None:
+    def __init__(self, transformer_embedding: nn.Embedding, vocab_size: int, embed_dim: int, encoder: Encoder, decoder: Decoder, with_lpips: bool = True) -> None:
        super().__init__()
        self.vocab_size = vocab_size
        self.encoder = encoder
        self.pre_quant_conv = torch.nn.Conv2d(encoder.config.z_channels, embed_dim, 1)
-        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        self.embedding = transformer_embedding  # pretrained transformer embedding
        self.post_quant_conv = torch.nn.Conv2d(embed_dim, decoder.config.z_channels, 1)
        self.decoder = decoder
-        self.embedding.weight.data.uniform_(-1.0 / vocab_size, 1.0 / vocab_size)
+        # self.embedding.weight.data.uniform_(-1.0 / vocab_size, 1.0 / vocab_size)
        self.lpips = LPIPS().eval() if with_lpips else None

    def __repr__(self) -> str:
@@ -46,6 +46,8 @@ class Tokenizer(nn.Module):
    def compute_loss(self, batch: Batch, **kwargs: Any) -> LossWithIntermediateLosses:
        assert self.lpips is not None
        observations = self.preprocess_input(rearrange(batch['observations'], 'b t c h w -> (b t) c h w'))
+        # TODO: in the delta-IRIS paper (https://openreview.net/forum?id=o8IDoZggqO) they encode(x0, a0, x1) -> z1 and decode(x0, a0, z1). In esense the tokens only need to encode the change
+        # note they also do dynamics(x0, a0, z1) -> z2. decode(x1, a1, z2) -> x2
        z, z_quantized, reconstructions = self(observations, should_preprocess=False, should_postprocess=False)

        # Codebook loss. Notes:
@@ -2,119 +2,163 @@
 # Credits to https://github.com/karpathy/minGPT
 # """

-# from dataclasses import dataclass
-# import math
-# from typing import Optional
-
-# from einops import rearrange
-# import torch
-# import torch.nn as nn
-# from torch.nn import functional as F
+from dataclasses import dataclass
+import math
+from typing import Optional
+from contextlib import contextmanager
+from einops import rearrange
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from loguru import logger

 # from .kv_caching import KeysValues, KVCache


-# @dataclass
-# class TransformerConfig:
-#     tokens_per_block: int
-#     max_blocks: int
-#     attention: str
+@dataclass
+class TransformerConfig:
+    # model_name: str = "stabilityai/stablelm-3b-4e1t"
+    # https://huggingface.co/PY007/TinyLlama-1.1B-intermediate-step-715k-1.5T
+    vocab_size: int = 32000
+    embed_dim: int = 2048
+    
+    max_blocks: int = 20
+    tokens_per_block: int = 17
+    
+    model_name: str = "PY007/TinyLlama-1.1B-intermediate-step-715k-1.5T"
+    dropout: float = 0.1
+    rank: int = 32

-#     num_layers: int
-#     num_heads: int
-#     embed_dim: int
-
-#     embed_pdrop: float
-#     resid_pdrop: float
-#     attn_pdrop: float
-
-#     @property
-#     def max_tokens(self):
-#         return self.tokens_per_block * self.max_blocks
+    @property
+    def max_tokens(self):
+        return self.tokens_per_block * self.max_blocks
+    
+    
+def freeze(n: nn.Module):
+    for p in n.parameters():
+        p.requires_grad = False
+    return n


-# class Transformer(nn.Module):
-#     def __init__(self, config: TransformerConfig) -> None:
-#         super().__init__()
-#         self.config = config
-#         self.drop = nn.Dropout(config.embed_pdrop)
-#         self.blocks = nn.ModuleList([Block(config) for _ in range(config.num_layers)])
-#         self.ln_f = nn.LayerNorm(config.embed_dim)
+class Transformer(nn.Module):
+    def __init__(self, config: TransformerConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.model = load_pretrained_model(config)
+        self.ln_f = nn.Linear(self.model.config.vocab_size, config.embed_dim)
+        self.embedding = freeze(self.model.base_model.embed_tokens.to(torch.float)) # HACK: custom path to embeddings layer for model

-#     def generate_empty_keys_values(self, n: int, max_tokens: int) -> KeysValues:
-#         device = self.ln_f.weight.device  # Assumption that all submodules are on the same device
-#         return KeysValues(n, self.config.num_heads, max_tokens, self.config.embed_dim, self.config.num_layers, device)
+    def generate_empty_keys_values(self, n: int, max_tokens: int) -> KeysValues:
+        device = self.ln_f.weight.device  # Assumption that all submodules are on the same device
+        return KeysValues(n, 1, max_tokens, self.config.embed_dim, 1, device)

-#     def forward(self, sequences: torch.Tensor, past_keys_values: Optional[KeysValues] = None) -> torch.Tensor:
-#         assert past_keys_values is None or len(past_keys_values) == len(self.blocks)
-#         x = self.drop(sequences)
-#         for i, block in enumerate(self.blocks):
-#             x = block(x, None if past_keys_values is None else past_keys_values[i])
+    def forward(self, sequences: torch.Tensor, past_keys_values: Optional[KeysValues] = None) -> torch.Tensor:
+        assert past_keys_values is None or len(past_keys_values) == 1
+        with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+            outputs = self.model(
+                inputs_embeds=sequences,
+                return_dict=True,
+                output_hidden_states=True,
+            )
+        x = outputs.logits
+        x = self.ln_f(x)
+        
+        # fake it, since it's used to keep track of steps
+        if past_keys_values is not None:
+            k_size = past_keys_values[0]._k_cache._cache.size()
+            v_size = (k_size[0], k_size[1], x.shape[1], k_size[3])
+            past_keys_values[0].update(torch.rand(v_size), torch.rand(v_size))
+        return x

-#         x = self.ln_f(x)
-#         return x
+        


-# class Block(nn.Module):
-#     def __init__(self, config: TransformerConfig) -> None:
-#         super().__init__()
-#         self.ln1 = nn.LayerNorm(config.embed_dim)
-#         self.ln2 = nn.LayerNorm(config.embed_dim)
-#         self.attn = SelfAttention(config)
-#         self.mlp = nn.Sequential(
-#             nn.Linear(config.embed_dim, 4 * config.embed_dim),
-#             nn.GELU(),
-#             nn.Linear(4 * config.embed_dim, config.embed_dim),
-#             nn.Dropout(config.resid_pdrop),
-#         )
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from peft import PeftModel, LoraConfig
+import peft

-#     def forward(self, x: torch.Tensor, past_keys_values: Optional[KeysValues] = None) -> torch.Tensor:
-#         x_attn = self.attn(self.ln1(x), past_keys_values)
-#         x = x + x_attn
-#         x = x + self.mlp(self.ln2(x))
-#         return x
+def load_pretrained_model(config, device="cuda:0"):
+    tokenizer = AutoTokenizer.from_pretrained(config.model_name, trust_remote_code=True)
+    tokenizer.padding_side = "left"
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+    )
+    base_model = AutoModelForCausalLM.from_pretrained(
+        config.model_name,
+        device_map={"": device},
+        quantization_config=bnb_config,
+        torch_dtype=torch.bfloat16, 
+        trust_remote_code=True
+    )
+    peft_config = peft.LoraConfig(
+        peft.TaskType.CAUSAL_LM,
+        inference_mode=False,
+        r=config.rank,
+        lora_alpha=config.rank*2, # Adjusting the LoRA rank is essential, and so is selecting an apt alpha value. A good heuristic is setting alpha at twice the rank's value. https://magazine.sebastianraschka.com/p/practical-tips-for-finetuning-llms
+        lora_dropout=config.dropout,
+        # TODO: If you're incorporating LoRA, ensure it's applied across all layers, not just to the Key and Value matrices, to maximize model performance.
+        target_modules=[
+            "self_attn.q_proj",
+            "self_attn.k_proj",
+            "self_attn.v_proj",
+            "self_attn.o_proj",
+            "mlp.gate_proj",
+            "mlp.up_proj",
+            "mlp.down_proj",
+            # "wte", "embed_tokens",
+            # "lm_head",
+        ],
+        # bias="lora_only",
+        # tune the embedding layer and prediction head
+        modules_to_save = ["lm_head",], # we want the classifier parameters to be trained too when fine-tuning the base model on our custom dataset. To ensure that the classifier parameters are also trained, we specify modules_to_save. 
+    )
+    base_model_peft = base_model
+    # base_model_peft = peft.get_peft_model(base_model, peft_config)
+    # base_model_peft.add_adapter(adapter_name="dynamics", peft_config=peft_config) # make and set an adapter
+    disable_causal_mask_always()
+    # print(base_model_peft.print_trainable_parameters())
+    logger.debug(f"loaded model {base_model_peft}")
+    return base_model_peft

+@contextmanager
+def set_adapter(model, adapter_name):
+    old_adapter_name = model.active_adapter
+    try:
+        if adapter_name is not None:
+            model.set_adapter(adapter_name)
+            yield model
+        else:
+            with model.disable_adapter():
+                yield model
+    finally:
+        model.set_adapter(old_adapter_name)

-# class SelfAttention(nn.Module):
-#     def __init__(self, config: TransformerConfig) -> None:
-#         super().__init__()
-#         assert config.embed_dim % config.num_heads == 0
-#         assert config.attention in ('causal', 'block_causal')
-#         self.num_heads = config.num_heads
-#         self.key = nn.Linear(config.embed_dim, config.embed_dim)
-#         self.query = nn.Linear(config.embed_dim, config.embed_dim)
-#         self.value = nn.Linear(config.embed_dim, config.embed_dim)
-#         self.attn_drop = nn.Dropout(config.attn_pdrop)
-#         self.resid_drop = nn.Dropout(config.resid_pdrop)
-#         self.proj = nn.Linear(config.embed_dim, config.embed_dim)
+def disable_causal_mask_always():
+    import transformers.models.llama.modeling_llama as modeling

-#         causal_mask = torch.tril(torch.ones(config.max_tokens, config.max_tokens))
-#         block_causal_mask = torch.max(causal_mask, torch.block_diag(*[torch.ones(config.tokens_per_block, config.tokens_per_block) for _ in range(config.max_blocks)]))
-#         self.register_buffer('mask', causal_mask if config.attention == 'causal' else block_causal_mask)
+    decoder_fn = modeling._make_causal_mask

-#     def forward(self, x: torch.Tensor, kv_cache: Optional[KVCache] = None) -> torch.Tensor:
-#         B, T, C = x.size()
-#         if kv_cache is not None:
-#             b, nh, L, c = kv_cache.shape
-#             assert nh == self.num_heads and b == B and c * nh == C
-#         else:
-#             L = 0
+    def encoder_fn(*args, **kwargs):
+        return torch.zeros_like(decoder_fn(*args, **kwargs))

-#         q = self.query(x).view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2)   # (B, nh, T, hs)
-#         k = self.key(x).view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2)     # (B, nh, T, hs)
-#         v = self.value(x).view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2)   # (B, nh, T, hs)
+    modeling._make_causal_mask = encoder_fn

-#         if kv_cache is not None:
-#             kv_cache.update(k, v)
-#             k, v = kv_cache.get()
+@contextmanager
+def disable_causal_mask():
+    import transformers.models.llama.modeling_llama as modeling

-#         att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
-#         att = att.masked_fill(self.mask[L:L + T, :L + T] == 0, float('-inf'))
-#         att = F.softmax(att, dim=-1)
-#         att = self.attn_drop(att)
-#         y = att @ v
-#         y = rearrange(y, 'b h t e -> b t (h e)')
+    decoder_fn = modeling._make_causal_mask

-#         y = self.resid_drop(self.proj(y))
+    def encoder_fn(*args, **kwargs):
+        return torch.zeros_like(decoder_fn(*args, **kwargs))

-#         return y
+    try:
+        modeling._make_causal_mask = encoder_fn
+        yield
+    finally:
+        modeling._make_causal_mask = decoder_fn
@@ -6,13 +6,12 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F

-from dataset import Batch
-from .kv_caching import KeysValues
-from .slicer import Embedder, Head
-from .tokenizer import Tokenizer
-# from .transformer import Transformer, TransformerConfig
-from .bigvae import BigVAE, BigVAEConfig
-from utils import init_weights, LossWithIntermediateLosses
+from src.dataset import Batch
+from src.models.kv_caching import KeysValues
+from src.models.slicer import Embedder, Head
+from src.models.tokenizer import Tokenizer
+from src.models.transformer import Transformer, TransformerConfig
+from src.utils import init_weights, LossWithIntermediateLosses


@dataclass
@@ -27,21 +26,35 @@ class WorldModel(nn.Module):
    def __init__(self, obs_vocab_size: int, act_vocab_size: int, config: BigVAEConfig) -> None:
        super().__init__()
        self.obs_vocab_size, self.act_vocab_size = obs_vocab_size, act_vocab_size
-        self.config = config
-        self.transformer = BigVAE(config)
+        self.config = config       

        all_but_last_obs_tokens_pattern = torch.ones(config.tokens_per_block)
        all_but_last_obs_tokens_pattern[-2] = 0
        act_tokens_pattern = torch.zeros(self.config.tokens_per_block)
        act_tokens_pattern[-1] = 1
        obs_tokens_pattern = 1 - act_tokens_pattern
+        
+        self.transformer = Transformer(config)
+        transformer_embedding = self.transformer.embedding

        self.pos_emb = nn.Embedding(config.max_tokens, config.embed_dim)
+        self.act_emb = nn.Embedding(act_vocab_size, config.embed_dim)

+        # FIXME: having slices is unclear. maybe it's better just to have obs and action embeddings?
        self.embedder = Embedder(
            max_blocks=config.max_blocks,
            block_masks=[act_tokens_pattern, obs_tokens_pattern],
-            embedding_tables=nn.ModuleList([nn.Embedding(act_vocab_size, config.embed_dim), nn.Embedding(obs_vocab_size, config.embed_dim)])
+            embedding_tables=nn.ModuleList([self.act_emb, transformer_embedding])
+        )
+        
+        # why have this? Well I worry that the transformer can't adapt, since so much is frozen
+        # TODO: If I get the dynamics model working, maybe try without it
+        self.post_embed = nn.Sequential(
+            nn.Linear(config.embed_dim, config.embed_dim),
+            nn.ReLU(),
+            nn.Linear(config.embed_dim, config.embed_dim),
+            # nn.ReLU(),
+            # nn.Linear(config.embed_dim, config.embed_dim)
        )

        self.head_observations = Head(
@@ -74,19 +87,28 @@ class WorldModel(nn.Module):
            )
        )

-        self.apply(init_weights)
+        # don't apply to transformer or transformer/obs embeddings
+        self.act_emb.apply(init_weights)
+        self.pos_emb.apply(init_weights)
+        self.post_embed.apply(init_weights)
+        self.head_observations.apply(init_weights)
+        self.head_rewards.apply(init_weights)
+        self.head_ends.apply(init_weights)
+        
+        

    def __repr__(self) -> str:
        return "world_model"

    def forward(self, tokens: torch.LongTensor, past_keys_values: Optional[KeysValues] = None) -> WorldModelOutput:

-        num_steps = tokens.size(1)  # (B, T)
+        num_steps = tokens.size(1)  # (B=8, T=170) where often the last 10 are actons
        assert num_steps <= self.config.max_tokens
        prev_steps = 0 if past_keys_values is None else past_keys_values.size

        sequences = self.embedder(tokens, num_steps, prev_steps) + self.pos_emb(prev_steps + torch.arange(num_steps, device=tokens.device))
-
+        # [batch=8, num_steps=170, embed_size=2048]
+        sequences = self.post_embed(sequences)
        x = self.transformer(sequences, past_keys_values)

        logits_observations = self.head_observations(x, num_steps=num_steps, prev_steps=prev_steps)
@@ -97,11 +119,13 @@ class WorldModel(nn.Module):

    def compute_loss(self, batch: Batch, tokenizer: Tokenizer, **kwargs: Any) -> LossWithIntermediateLosses:

-        with torch.no_grad():
-            obs_tokens = tokenizer.encode(batch['observations'], should_preprocess=True).tokens  # (BL, K)
+        # with torch.no_grad():
+        # [B=8, S=10, Colors=3, H=64, W=64] -> [B=8, S=10, 16]
+        obs_tokens = tokenizer.encode(batch['observations'], should_preprocess=True).tokens  # (BL, K)

        act_tokens = rearrange(batch['actions'], 'b l -> b l 1')
        tokens = rearrange(torch.cat((obs_tokens, act_tokens), dim=2), 'b l k1 -> b (l k1)')  # (B, L(K+1))
+        # So first 10 are observation, the last 10 tokens are actions

        outputs = self(tokens)

@@ -1,4 +1,4 @@
-from functools import partial 
+from functools import partial
 from pathlib import Path

 import hydra
@@ -6,55 +6,96 @@ from hydra.utils import instantiate
 from omegaconf import DictConfig
 import torch

-from agent import Agent
-from envs import SingleProcessEnv, WorldModelEnv
-from game import AgentEnv, EpisodeReplayEnv, Game
-from models.actor_critic import ActorCritic
-from models.world_model import WorldModel
+from src.agent import Agent
+from src.envs import SingleProcessEnv, WorldModelEnv
+from src.game import AgentEnv, EpisodeReplayEnv, Game
+from src.models.actor_critic import ActorCritic
+from src.models.world_model import WorldModel
+from src.models.tokenizer import Tokenizer


@hydra.main(config_path="../config", config_name="trainer")
 def main(cfg: DictConfig):
    device = torch.device(cfg.common.device)
-    assert cfg.mode in ('episode_replay', 'agent_in_env', 'agent_in_world_model', 'play_in_world_model')
+    assert cfg.mode in (
+        "episode_replay",
+        "agent_in_env",
+        "agent_in_world_model",
+        "play_in_world_model",
+    )

    env_fn = partial(instantiate, config=cfg.env.test)
    test_env = SingleProcessEnv(env_fn)

-    if cfg.mode.startswith('agent_in_'):
+    if cfg.mode.startswith("agent_in_"):
        h, w, _ = test_env.env.unwrapped.observation_space.shape
    else:
        h, w = 64, 64
    multiplier = 800 // h
    size = [h * multiplier, w * multiplier]
-    
-    if cfg.mode == 'episode_replay':
-        env = EpisodeReplayEnv(replay_keymap_name=cfg.env.keymap, episode_dir=Path('media/episodes'))
-        keymap = 'episode_replay'
+
+    if cfg.mode == "episode_replay":
+        env = EpisodeReplayEnv(
+            replay_keymap_name=cfg.env.keymap, episode_dir=Path("media/episodes")
+        )
+        keymap = "episode_replay"

    else:
-        tokenizer = instantiate(cfg.tokenizer)
-        world_model = WorldModel(obs_vocab_size=tokenizer.vocab_size, act_vocab_size=test_env.num_actions, config=instantiate(cfg.world_model))
-        actor_critic = ActorCritic(**cfg.actor_critic, act_vocab_size=test_env.num_actions)
+        # tokenizer = instantiate(cfg.tokenizer)
+        world_model = WorldModel(
+            obs_vocab_size=cfg.tokenizer.vocab_size,
+            act_vocab_size=test_env.num_actions,
+            config=instantiate(cfg.world_model),
+        )
+        transformer_embedding = world_model.transformer.embedding
+        tokenizer = Tokenizer(
+            transformer_embedding=transformer_embedding,
+            vocab_size=cfg.tokenizer.vocab_size,
+            embed_dim=cfg.tokenizer.embed_dim,
+            encoder=instantiate(cfg.tokenizer.encoder),
+            decoder=instantiate(cfg.tokenizer.decoder),
+        )
+        actor_critic = ActorCritic(
+            **cfg.actor_critic, act_vocab_size=test_env.num_actions
+        )
        agent = Agent(tokenizer, world_model, actor_critic).to(device)
-        agent.load(Path('checkpoints/last.pt'), device)        
+        agent.load(Path("checkpoints/last.pt"), device)

-        if cfg.mode == 'play_in_world_model':
-            env = WorldModelEnv(tokenizer=agent.tokenizer, world_model=agent.world_model, device=device, env=env_fn())
+        if cfg.mode == "play_in_world_model":
+            env = WorldModelEnv(
+                tokenizer=agent.tokenizer,
+                world_model=agent.world_model,
+                device=device,
+                env=env_fn(),
+            )
            keymap = cfg.env.keymap
-        
-        elif cfg.mode == 'agent_in_env':
-            env = AgentEnv(agent, test_env, cfg.env.keymap, do_reconstruction=cfg.reconstruction)
-            keymap = 'empty'
+
+        elif cfg.mode == "agent_in_env":
+            env = AgentEnv(
+                agent, test_env, cfg.env.keymap, do_reconstruction=cfg.reconstruction
+            )
+            keymap = "empty"
            if cfg.reconstruction:
                size[1] *= 3

-        elif cfg.mode == 'agent_in_world_model':
-            wm_env = WorldModelEnv(tokenizer=agent.tokenizer, world_model=agent.world_model, device=device, env=env_fn())
+        elif cfg.mode == "agent_in_world_model":
+            wm_env = WorldModelEnv(
+                tokenizer=agent.tokenizer,
+                world_model=agent.world_model,
+                device=device,
+                env=env_fn(),
+            )
            env = AgentEnv(agent, wm_env, cfg.env.keymap, do_reconstruction=False)
-            keymap = 'empty'
+            keymap = "empty"

-    game = Game(env, keymap_name=keymap, size=size, fps=cfg.fps, verbose=bool(cfg.header), record_mode=bool(cfg.save_mode))
+    game = Game(
+        env,
+        keymap_name=keymap,
+        size=size,
+        fps=cfg.fps,
+        verbose=bool(cfg.header),
+        record_mode=bool(cfg.save_mode),
+    )
    game.run()


@@ -14,14 +14,15 @@ import torch.nn as nn
 from tqdm import tqdm
 import wandb

-from agent import Agent
-from collector import Collector
-from envs import SingleProcessEnv, MultiProcessEnv
-from episode import Episode
-from make_reconstructions import make_reconstructions_from_batch
-from models.actor_critic import ActorCritic
-from models.world_model import WorldModel
-from utils import configure_optimizer, EpisodeDirManager, set_seed
+from src.agent import Agent
+from src.collector import Collector
+from src.envs import SingleProcessEnv, MultiProcessEnv
+from src.episode import Episode
+from src.make_reconstructions import make_reconstructions_from_batch
+from src.models.actor_critic import ActorCritic
+from src.models.world_model import WorldModel
+from src.utils import configure_optimizer, EpisodeDirManager, set_seed
+from src.models.tokenizer import Tokenizer


 class Trainer:
@@ -46,6 +47,7 @@ class Trainer:
        self.reconstructions_dir = self.media_dir / 'reconstructions'

        if not cfg.common.resume:
+            print('cwd', Path.cwd())
            config_dir = Path('config')
            config_path = config_dir / 'trainer.yaml'
            config_dir.mkdir(exist_ok=False, parents=False)
@@ -79,8 +81,15 @@ class Trainer:
        assert self.cfg.training.should or self.cfg.evaluation.should
        env = train_env if self.cfg.training.should else test_env

-        tokenizer = instantiate(cfg.tokenizer)
-        world_model = WorldModel(obs_vocab_size=tokenizer.vocab_size, act_vocab_size=env.num_actions, config=instantiate(cfg.world_model))
+        world_model = WorldModel(obs_vocab_size=cfg.tokenizer.vocab_size, act_vocab_size=env.num_actions, config=instantiate(cfg.world_model))
+        transformer_embedding = world_model.transformer.embedding
+        tokenizer = Tokenizer(
+            transformer_embedding=transformer_embedding,
+            vocab_size=cfg.tokenizer.vocab_size,
+            embed_dim=cfg.tokenizer.embed_dim,
+            encoder=instantiate(cfg.tokenizer.encoder),
+            decoder=instantiate(cfg.tokenizer.decoder),
+        )
        actor_critic = ActorCritic(**cfg.actor_critic, act_vocab_size=env.num_actions)
        self.agent = Agent(tokenizer, world_model, actor_critic).to(self.device)
        print(f'{sum(p.numel() for p in self.agent.tokenizer.parameters())} parameters in agent.tokenizer')
@@ -88,7 +97,11 @@ class Trainer:
        print(f'{sum(p.numel() for p in self.agent.actor_critic.parameters())} parameters in agent.actor_critic')

        self.optimizer_tokenizer = torch.optim.Adam(self.agent.tokenizer.parameters(), lr=cfg.training.learning_rate)
-        self.optimizer_world_model = configure_optimizer(self.agent.world_model, cfg.training.learning_rate, cfg.training.world_model.weight_decay)
+        # self.optimizer_world_model = configure_optimizer([self.agent.tokenizer, self.agent.world_model], cfg.training.learning_rate, cfg.training.world_model.weight_decay)
+        self.optimizer_world_model = torch.optim.Adam(
+            list(self.agent.tokenizer.parameters())+list(self.agent.world_model.parameters()),
+            lr=cfg.training.learning_rate
+        )
        self.optimizer_actor_critic = torch.optim.Adam(self.agent.actor_critic.parameters(), lr=cfg.training.learning_rate)

        if cfg.initialization.path_to_checkpoint is not None:
@@ -136,10 +149,10 @@ class Trainer:

        if epoch > cfg_tokenizer.start_after_epochs:
            metrics_tokenizer = self.train_component(self.agent.tokenizer, self.optimizer_tokenizer, sequence_length=1, sample_from_start=True, **cfg_tokenizer)
-        self.agent.tokenizer.eval()

        if epoch > cfg_world_model.start_after_epochs:
            metrics_world_model = self.train_component(self.agent.world_model, self.optimizer_world_model, sequence_length=self.cfg.common.sequence_length, sample_from_start=True, tokenizer=self.agent.tokenizer, **cfg_world_model)
+        self.agent.tokenizer.eval()
        self.agent.world_model.eval()

        if epoch > cfg_actor_critic.start_after_epochs:
@@ -169,6 +182,7 @@ class Trainer:
            if max_grad_norm is not None:
                torch.nn.utils.clip_grad_norm_(component.parameters(), max_grad_norm)

+ 
            optimizer.step()

        metrics = {f'{str(component)}/train/total_loss': loss_total_epoch, **intermediate_losses}
@@ -190,7 +204,7 @@ class Trainer:
        if epoch > cfg_world_model.start_after_epochs:
            metrics_world_model = self.eval_component(self.agent.world_model, cfg_world_model.batch_num_samples, sequence_length=self.cfg.common.sequence_length, tokenizer=self.agent.tokenizer)

-        if epoch > cfg_actor_critic.start_after_epochs:
+        if epoch > cfg_world_model.start_after_epochs:
            self.inspect_imagination(epoch)

        if cfg_tokenizer.save_reconstructions:
@@ -3,47 +3,64 @@ import cv2
 from pathlib import Path
 import random
 import shutil
-
+from loguru import logger
 import numpy as np
 import torch
 import torch.nn as nn

-from episode import Episode
+from src.episode import Episode
+
+from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS


-def configure_optimizer(model, learning_rate, weight_decay, *blacklist_module_names):
+def configure_optimizer(models, learning_rate, weight_decay, *blacklist_module_names):
    """Credits to https://github.com/karpathy/minGPT"""
+    # FIXME: check this is still good for LoRA
    # separate out all parameters to those that will and won't experience regularizing weight decay
    decay = set()
    no_decay = set()
+    decay_params = []
+    no_decay_params = []
+    param_dict = {}
    whitelist_weight_modules = (torch.nn.Linear, torch.nn.Conv1d)
-    blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding)
-    for mn, m in model.named_modules():
-        for pn, p in m.named_parameters():
-            fpn = '%s.%s' % (mn, pn) if mn else pn  # full param name
-            if any([fpn.startswith(module_name) for module_name in blacklist_module_names]):
-                no_decay.add(fpn)
-            elif 'bias' in pn:
-                # all biases will not be decayed
-                no_decay.add(fpn)
-            elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):
-                # weights of whitelist modules will be weight decayed
-                decay.add(fpn)
-            elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):
-                # weights of blacklist modules will NOT be weight decayed
-                no_decay.add(fpn)
+    blacklist_weight_modules = tuple(ALL_LAYERNORM_LAYERS+[torch.nn.Embedding])
+    for model in models:
+        for mn, m in model.named_modules():
+            for pn, p in m.named_parameters():
+                fpn = '%s.%s' % (mn, pn) if mn else pn  # full param name
+                if any([fpn.startswith(module_name) for module_name in blacklist_module_names]):
+                    no_decay.add(fpn)
+                    no_decay_params.append(p)
+                elif 'bias' in pn:
+                    # all biases will not be decayed
+                    no_decay.add(fpn)
+                    no_decay_params.append(p)
+                elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):
+                    # weights of whitelist modules will be weight decayed
+                    decay.add(fpn)
+                    decay_params.append(p)
+                elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):
+                    # weights of blacklist modules will NOT be weight decayed
+                    no_decay.add(fpn)
+                    no_decay_params.append(p)
+                else:
+                    logger.warning(f"Parameter {fpn} of module {mn} not handled!")
+                    # raise NotImplementedError(f"Parameter {fpn} of module {m} not handled!")
+                    decay.add(fpn)
+                    decay_params.append(p)

-    # validate that we considered every parameter
-    param_dict = {pn: p for pn, p in model.named_parameters()}
+        # validate that we considered every parameter
+        param_dict.update({pn: p for pn, p in model.named_parameters()})
    inter_params = decay & no_decay
    union_params = decay | no_decay
+    # logger.debug(f"decay {decay} no_decay {no_decay}")
    assert len(inter_params) == 0, f"parameters {str(inter_params)} made it into both decay/no_decay sets!"
    assert len(param_dict.keys() - union_params) == 0, f"parameters {str(param_dict.keys() - union_params)} were not separated into either decay/no_decay set!"

    # create the pytorch optimizer object
    optim_groups = [
-        {"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": weight_decay},
-        {"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0},
+        {"params": no_decay_params, "weight_decay": weight_decay},
+        {"params": decay_params, "weight_decay": 0.0},
    ]
    optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate)
    return optimizer