model-experiments.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: CUBLAS_WORKSPACE_CONFIG=:4096:8\n"
     ]
    }
   ],
   "source": [
    "%set_env CUBLAS_WORKSPACE_CONFIG=:4096:8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/theomichel/work/pyDreamer/env/lib/python3.8/site-packages/tqdm/auto.py:22: 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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Namespace(action_dim=18, actor_dist='onehot', actor_grad='reinforce', adam_eps=1e-05, adam_lr=0.0003, adam_lr_actor=0.0001, adam_lr_critic=0.0001, allow_mid_reset=True, amp=False, aux_critic=False, aux_critic_weight=1.0, batch_length=48, batch_size=32, buffer_size=10000000, buffer_size_offline=0, clip_rewards='tanh', cnn_depth=48, data_workers=4, deter_dim=1024, device='cuda:0', enable_profiler=False, entropy=0.001, env_action_repeat=4, env_id='Atari-Alien-V5', env_id_eval=None, env_no_terminal=False, env_time_limit=27000, eval_batch_size=32, eval_batches=61, eval_interval=2000, eval_samples=1, eval_save_size=1, gamma=0.99, gamma_aux=0.99, generator_prefill_policy='random', generator_prefill_steps=50000, generator_workers=1, generator_workers_eval=0, generator_workers_train=0, goals_size=0, grad_clip=200, grad_clip_ac=200, gru_layers=1, gru_type='gru', hidden_dim=1000, imag_horizon=15, image_categorical=False, image_channels=3, image_decoder='cnn', image_decoder_layers=0, image_decoder_min_prob=0, image_encoder='cnn', image_encoder_layers=0, image_key='image', image_size=64, image_weight=1.0, iwae_samples=1, keep_state=True, kl_balance=0.8, kl_weight=0.1, lambda_gae=0.95, lambda_gae_aux=0.95, layer_norm=True, limit_step_ratio=0, log_interval=100, logbatch_interval=1000, map_categorical=True, map_channels=0, map_decoder='dense', map_hidden_dim=1024, map_hidden_layers=4, map_key=None, map_size=0, model='dreamer', n_env_steps=200000000, n_steps=100000000, offline_data_dir=None, offline_eval_dir=None, offline_prefill_dir=None, offline_test_dir=None, probe_gradients=False, probe_model='none', reset_interval=200, reward_decoder_categorical=None, reward_decoder_layers=4, reward_input=False, reward_weight=1.0, save_interval=500, stoch_dim=32, stoch_discrete=32, target_interval=100, target_interval_aux=1000, terminal_decoder_layers=4, terminal_weight=1.0, test_batch_size=10, test_batches=61, test_save_size=1, vecobs_size=0, vecobs_weight=1.0, verbose=False)\n",
      "[launcher]  Started mlflow run 11600f3c3ab042fe8d69e1cf2fe7481d (None) in local/0\u001b[0m\n",
      "torch.Size([48, 32, 18])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import argparse\n",
    "from pydreamer.models.dreamer import Dreamer\n",
    "from pydreamer.tools import mlflow_load_checkpoint\n",
    "from pydreamer.tools import (configure_logging, mlflow_log_params,\n",
    "                             mlflow_init, print_once, read_yamls)\n",
    "from distutils.util import strtobool\n",
    "import os\n",
    "import json\n",
    "import mlflow\n",
    "from torchsummary import summary\n",
    "import pickle\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def get_worker_info():\n",
    "    worker_type = None\n",
    "    worker_index = None\n",
    "\n",
    "    if 'TF_CONFIG' in os.environ:\n",
    "        # TF_CONFIG indicates Google Vertex AI run\n",
    "        tf_config = json.loads(os.environ['TF_CONFIG'])\n",
    "        print_once('TF_CONFIG is set:', tf_config)\n",
    "        if tf_config['cluster'].get('worker'):\n",
    "            # If there are workers in the cluster, then it's a distributed run\n",
    "            worker_type = {\n",
    "                'chief': 'learner',\n",
    "                'worker': 'generator',\n",
    "            }[str(tf_config['task']['type'])]\n",
    "            worker_index = int(tf_config['task']['index'])\n",
    "            print_once('Distributed run detected, current worker is:', f'{worker_type} ({worker_index})')\n",
    "\n",
    "    return worker_type, worker_index\n",
    "\n",
    "configure_logging('[launcher]')\n",
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument('--configs', nargs='+', required=True)\n",
    "# args, remaining = parser.parse_known_args()\n",
    "#--configs defaults atari --env_id Atari-Alien-V5\n",
    "# Config from YAML\n",
    "args_list = ['defaults', 'atari']\n",
    "remaining = ['--env_id', 'Atari-Alien-V5']\n",
    "conf = {}\n",
    "configs = read_yamls('./config')\n",
    "for name in args_list:\n",
    "    if ',' in name:\n",
    "        for n in name.split(','):\n",
    "            conf.update(configs[n])\n",
    "    else:\n",
    "        conf.update(configs[name])\n",
    "\n",
    "# Override config from command-line\n",
    "\n",
    "parser = argparse.ArgumentParser()\n",
    "for key, value in conf.items():\n",
    "    type_ = type(value) if value is not None else str\n",
    "    if type_ == bool:\n",
    "        type_ = lambda x: bool(strtobool(x))\n",
    "    parser.add_argument(f'--{key}', type=type_, default=value)\n",
    "conf = parser.parse_args(remaining)\n",
    "\n",
    "print(conf)\n",
    "\n",
    "worker_type, worker_index = get_worker_info()\n",
    "is_main_worker = worker_type is None or worker_type == 'learner'\n",
    "mlrun = mlflow_init(wait_for_resume=not is_main_worker)\n",
    "artifact_uri = mlrun.info.artifact_uri\n",
    "mlflow_log_params(vars(conf))\n",
    "model = Dreamer(conf)\n",
    "optimizers=tuple()\n",
    "map_location=None#Same place that it is loaded from\n",
    "path = \"/home/theomichel/work/pyDreamer/pydreamer-minigrid/mlruns/0/c1a7c69b35fa4570915c6be36f57f2c9/artifacts/checkpoints/latest.pt\"\n",
    "checkpoint = torch.load(path, map_location=map_location)\n",
    "model.load_state_dict(checkpoint['model_state_dict'])\n",
    "for i, opt in enumerate(optimizers):\n",
    "    opt.load_state_dict(checkpoint[f'optimizer_{i}_state_dict'])\n",
    "    checkpoint['epoch']\n",
    "\n",
    "\n",
    "## Evaluate the model\n",
    "from train import evaluate\n",
    "from pydreamer.data import DataSequential, MlflowEpisodeRepository\n",
    "from torch.utils.data import DataLoader\n",
    "from pydreamer.preprocessing import Preprocessor, WorkerInfoPreprocess\n",
    "\n",
    "preprocess = Preprocessor(image_categorical=conf.image_channels if conf.image_categorical else None,\n",
    "                            image_key=conf.image_key,\n",
    "                            map_categorical=conf.map_channels if conf.map_categorical else None,\n",
    "                            map_key=conf.map_key,\n",
    "                            action_dim=conf.action_dim,\n",
    "                            clip_rewards=conf.clip_rewards,\n",
    "                            amp=conf.amp and device.type == 'cuda')\n",
    "\n",
    "\n",
    "device = torch.device(conf.device)\n",
    "steps = 10\n",
    "model.to(device)#Important\n",
    "\n",
    "# To be replaced with your state\n",
    "with open('states/in_state_alien_new10.pkl','rb') as f:\n",
    "    in_state = pickle.load(f)\n",
    "with open('states/obs_alien_new10.pkl','rb') as f:\n",
    "    obs = pickle.load(f)\n",
    "print(obs['action'].shape)\n",
    "for key in obs:\n",
    "    obs[key] = obs[key].to(device)\n",
    "#TODO in state are being resaved in the loop, get one clean\n",
    "\n",
    "## IMPORTANT : The following commented cells have not been retester and might need some adjustments to work\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "in_state_new = (in_state[0].to(device),in_state[1].to(device))\n",
    "\n",
    "#Forward is just used when the model is translating the world into features for the policy to take a decision\n",
    "# features, out_state = model.wm.forward(obs,in_state_new)\n",
    "#Training step does all evem the image prediction\n",
    "\n",
    "# loss, features, states, out_state, metrics, tensors = model.wm.training_step(obs,in_state_new,do_image_pred=True)\n",
    "\n",
    "def tensor_to_image(tensors,image_num=0,dream_num=0):\n",
    "    image_cpu = tensors['image_pred'].cpu().numpy()\n",
    "    image = image_cpu[image_num,dream_num,:,:,:].transpose(1,2,0)\n",
    "    image_final = ((image + 0.5) * 255.0).clip(0, 255).astype('uint8')\n",
    "    img_rgb = image_final#[...,::-1]\n",
    "    del image_cpu\n",
    "    return img_rgb\n",
    "\n",
    "def save_avg(obs,in_state_new,img_shape,save_path=\"images/avg\",image_num=2,dream_num=2,l=1000):\n",
    "    # fig, axs = plt.subplots(1, l,figsize=(20,80))\n",
    "    avg = np.zeros(img_shape)\n",
    "    for i in tqdm(range(l)):\n",
    "        with torch.no_grad():#imag_horizon=15,\n",
    "            _, _, _, tensors, dream_tensors = model.training_step(obs,in_state_new,imag_horizon=15,do_image_pred=True,do_dream_tensors=True)#Changed to model.training_step\n",
    "        img_rgb = tensor_to_image(dream_tensors,image_num=image_num,dream_num=dream_num)#dream tensors or tensors ?\n",
    "        avg += img_rgb\n",
    "        del tensors\n",
    "        del dream_tensors\n",
    "    avg = avg/l\n",
    "    im = Image.fromarray((avg).astype(np.uint8))\n",
    "    im.save(f\"{save_path}-{l}.png\")\n",
    "    return avg\n",
    "\n",
    "# im = Image.fromarray(img_rgb)\n",
    "# im.save(\"predicted_dream_alien.jpeg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([48, 32, 2048])\n",
      "torch.Size([48, 32, 3, 64, 64])\n"
     ]
    }
   ],
   "source": [
    "# with open('features_dream2/feature_alien_new8.pkl','rb') as f:\n",
    "with open('features_dream2/feature_alien_new50.pkl','rb') as f:\n",
    "    features_dream = pickle.load(f)\n",
    "#test directed decoding with the world model\n",
    "print(features_dream.shape)\n",
    "with torch.no_grad():\n",
    "    out = model.wm.decoder.image.forward(features_dream)\n",
    "print(out.shape)\n",
    "# show image\n",
    "img = out[2,2,:,:,:].cpu().numpy().transpose(1,2,0)\n",
    "img = ((img + 0.5) * 255.0).clip(0, 255).astype('uint8')\n",
    "im = Image.fromarray(img)\n",
    "im.save(\"dream_from_features_new_2.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# #drop each latent variable to zero and see the image prediction\n",
    "# with torch.no_grad():#imag_horizon=15,\n",
    "#     losses, new_state, loss_metrics, tensors, dream_tensors = model.training_step(obs,in_state_new,imag_horizon=15,do_image_pred=True,do_dream_tensors=True)#Changed to model.training_step\n",
    "\n",
    "# #pick image and dream\n",
    "\n",
    "# origin_prediction = tensor_to_image(tensors,image_num=2,dream_num=2)\n",
    "# print(tensors.keys())\n",
    "# del tensors\n",
    "\n",
    "# print(origin_prediction.shape)\n",
    "\n",
    "\n",
    "# l = 10\n",
    "# fig, axs = plt.subplots(1, l,figsize=(20,80))\n",
    "# for i in range(l):\n",
    "#     # instate_copy = in_state_new[1].clone()\n",
    "#     # instate_copy[:,i] = 0\n",
    "#     # in_state_new =(in_state_new[0],instate_copy)\n",
    "#     with torch.no_grad():#imag_horizon=15,\n",
    "#         losses, new_state, loss_metrics, tensors, dream_tensors = model.training_step(obs,in_state_new,imag_horizon=15,do_image_pred=True,do_dream_tensors=True)#Changed to model.training_step\n",
    "#     img_rgb = tensor_to_image(tensors,image_num=2,dream_num=2)\n",
    "#     del tensors\n",
    "#     axs[i].imshow(img_rgb)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3/3 [00:01<00:00,  2.39it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([32, 1024])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3/3 [00:01<00:00,  2.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(64, 64, 3)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "#Compare 2 images\n",
    "in_state_new = (in_state[0].to(device),in_state[1].to(device))\n",
    "\n",
    "with torch.no_grad():#imag_horizon=15,\n",
    "    losses, new_state, loss_metrics, tensors, dream_tensors = model.training_step(obs,in_state_new,imag_horizon=15,do_image_pred=True,do_dream_tensors=True)#Changed to model.training_step\n",
    "\n",
    "#pick image and dream\n",
    "origin_prediction = tensor_to_image(tensors,image_num=2,dream_num=2)\n",
    "del tensors\n",
    "\n",
    "#now we compare the difference when modyfing h\n",
    "avg1 = save_avg(obs,in_state_new,origin_prediction.shape,save_path=\"images/avg_new_unmodif\",image_num=2,dream_num=2,l=3)\n",
    "#modify the in state and see the effect\n",
    "a = in_state_new[1].clone()#set everything to zero in h\n",
    "print(a.shape)\n",
    "a[:,0:8] = 1\n",
    "\n",
    "b = in_state_new[0].clone()\n",
    "in_state_new = (b,a)\n",
    "\n",
    "avg2 = save_avg(obs,in_state_new,origin_prediction.shape,save_path=\"images/avg_new_modif_one_8\",image_num=2,dream_num=2,l=3)\n",
    "\n",
    "diff = np.sum(np.abs(avg1 - avg2)*3,axis=2)\n",
    "diff = diff.astype(np.uint8)\n",
    "im = Image.fromarray(diff)\n",
    "im.save(\"images/diff-same_state_ones_8.png\")\n",
    "\n",
    "\n",
    "\n",
    "#overlap the two images avg1 and diff, appering in red\n",
    "\n",
    "overlap = avg1.copy()\n",
    "\n",
    "print(overlap.shape)\n",
    "overlap[:,:,0] = overlap[:,:,0] + diff\n",
    "im = Image.fromarray(overlap.astype(np.uint8))\n",
    "im.save(\"images/overlap_8.png\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Compare the difference when modifying the latent variables\n",
    "# in_state_new = (in_state[0].to(device),in_state[1].to(device))\n",
    "\n",
    "# with torch.no_grad():#imag_horizon=15,\n",
    "#     losses, new_state, loss_metrics, tensors, dream_tensors = model.training_step(obs,in_state_new,imag_horizon=15,do_image_pred=True,do_dream_tensors=True)#Changed to model.training_step\n",
    "\n",
    "# #pick image and dream\n",
    "# origin_prediction = tensor_to_image(tensors,image_num=2,dream_num=2)\n",
    "# del tensors\n",
    "\n",
    "# #now we compare the difference when modyfing h\n",
    "# avg1 = save_avg(obs,in_state_new,origin_prediction.shape,save_path=\"images/avg_new_unmodif\",image_num=2,dream_num=2,l=100)\n",
    "# copy = (in_state_new[0].clone(),in_state_new[1].clone())\n",
    "# #modify the in state and see the effect\n",
    "# a = torch.zeros_like(in_state_new[1])#set everything to zero in h\n",
    "# b = in_state_new[0].clone()\n",
    "# in_state_new = (b,a)\n",
    "\n",
    "# avg2 = save_avg(obs,in_state_new,origin_prediction.shape,save_path=\"images/avg_new_modif\",image_num=2,dream_num=2,l=100)\n",
    "\n",
    "# diff = np.sum(np.abs(avg1 - avg2),axis=2)\n",
    "# diff = diff.astype(np.uint8)\n",
    "# im = Image.fromarray(diff)\n",
    "# im.save(\"images/diff-dif_state.jpeg\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We could imagine trying to do state interpolation or do advanced neural network interpretation Like LRP."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 ('env')",
   "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.8.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "efb447acfb63a1a0f180b1e27388da43dbd6168661e9674043ca2e4a62bff135"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}