{ "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': },\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": [ "" ] }, "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": [ "" ] }, "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----> 1\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: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 }