diff --git a/configs.yaml b/configs.yaml index 03a46fa..2be2d69 100644 --- a/configs.yaml +++ b/configs.yaml @@ -62,7 +62,7 @@ defaults: initial: 'learned' # Training - batch_size: 256 + batch_size: 64 batch_length: 64 train_ratio: 512 pretrain: 100 @@ -136,18 +136,59 @@ craftax: action_repeat: 1 envs: 1 train_ratio: 512 - video_pred_log: false # FIXME + video_pred_log: false dyn_hidden: 1024 dyn_deter: 4096 units: 1024 - encoder: {cnn_keys: '$^', mlp_keys: "state", mlp_layers: 5, mlp_units: 1024, } - decoder: {cnn_keys: '$^', mlp_keys: "state", mlp_layers: 5, mlp_units: 1024} + encoder: {cnn_keys: '$^', mlp_keys: "state", mlp_layers: 4, mlp_units: 512, } + decoder: {cnn_keys: '$^', mlp_keys: "state", mlp_layers: 4, mlp_units: 512} actor: {layers: 5, dist: 'onehot', std: 'none'} value: {layers: 5} reward_head: {layers: 5} cont_head: {layers: 5} imag_gradient: 'reinforce' +craftax_small: + task: craftax_Craftax-Symbolic-AutoReset-v1 + step: 1e6 + action_repeat: 1 + envs: 1 + train_ratio: 512 + video_pred_log: false + dyn_hidden: 512 + dyn_deter: 512 + units: 512 + encoder: {cnn_keys: '$^', mlp_keys: "state", mlp_layers: 3, mlp_units: 256} + decoder: {cnn_keys: '$^', mlp_keys: "state", mlp_layers: 3, mlp_units: 256} + actor: {layers: 3, dist: 'onehot', std: 'none'} + value: {layers: 3} + reward_head: {layers: 3} + cont_head: {layers: 3} + imag_gradient: 'reinforce' + batch_size: 128 + batch_length: 16 + + +craftax_smaller: + task: craftax_Craftax-Symbolic-AutoReset-v1 + step: 1e6 + action_repeat: 1 + envs: 1 + train_ratio: 256 + video_pred_log: false + dyn_hidden: 256 + dyn_deter: 1024 + units: 256 + encoder: {cnn_keys: '$^', mlp_keys: "state", mlp_layers: 2, mlp_units: 256, } + decoder: {cnn_keys: '$^', mlp_keys: "state", mlp_layers: 2, mlp_units: 256} + actor: {layers: 2, dist: 'onehot', std: 'none'} + value: {layers: 2} + reward_head: {layers: 2} + cont_head: {layers: 2} + imag_gradient: 'reinforce' + batch_size: 256 + batch_length: 16 + atari100k: steps: 4e5 envs: 1 diff --git a/dreamer.py b/dreamer.py index c8fa493..ed940e0 100644 --- a/dreamer.py +++ b/dreamer.py @@ -308,6 +308,7 @@ def main(config): agent.load_state_dict(checkpoint["agent_state_dict"]) tools.recursively_load_optim_state_dict(agent, checkpoint["optims_state_dict"]) agent._should_pretrain._once = False + logger.warning(f"Loaded model from {logdir / 'latest.pt'}") # make sure eval will be executed once after config.steps with tqdm(total=config.steps + config.eval_every, unit='step') as pbar: @@ -356,13 +357,12 @@ def main(config): except Exception: pass - -if __name__ == "__main__": +def parse_args(argv=None): parser = argparse.ArgumentParser() parser.add_argument("--configs", nargs="+") - args, remaining = parser.parse_known_args() + args, remaining = parser.parse_known_args(argv[1:]) configs = yaml.safe_load( - (pathlib.Path(sys.argv[0]).parent / "configs.yaml").read_text() + (pathlib.Path(argv[0]).parent / "configs.yaml").read_text() ) def recursive_update(base, update): @@ -380,4 +380,8 @@ if __name__ == "__main__": for key, value in sorted(defaults.items(), key=lambda x: x[0]): arg_type = tools.args_type(value) parser.add_argument(f"--{key}", type=arg_type, default=arg_type(value)) - main(parser.parse_args(remaining)) + args = parser.parse_args(remaining) + return args + +if __name__ == "__main__": + main(parse_args()) diff --git a/envs/craftax_env.py b/envs/craftax_env.py index 32fc068..2ea73cf 100644 --- a/envs/craftax_env.py +++ b/envs/craftax_env.py @@ -224,6 +224,8 @@ class Craftax: def step(self, action): state, reward, done, info = self._env.step(action) + info2 = {k.replace('Ach','log_ach'):v for k,v in info.items()} + reward = np.float32(reward) obs = { "image": self.get_image(), @@ -231,7 +233,7 @@ class Craftax: "is_first": False, "is_last": done, "is_terminal": info["discount"] == 0, - **info, + **info2, } return obs, reward, done, info diff --git a/image.png b/image.png new file mode 100644 index 0000000..a57bec3 Binary files /dev/null and b/image.png differ diff --git a/justfile b/justfile index ffc7f2c..c6c8b0a 100644 --- a/justfile +++ b/justfile @@ -6,7 +6,7 @@ export TQDM_MININTERVAL := "30" main: . ./.venv/bin/activate - python dreamer.py --configs crafter --task crafter_reward --logdir ./logdir/crafter + python dreamer.py --configs craftax_small --logdir ./logdir/crafter logs: tensorboard --logdir logdir/craftax diff --git a/nbs/02_torchinfo.ipynb b/nbs/02_torchinfo.ipynb new file mode 100644 index 0000000..3e551c8 --- /dev/null +++ b/nbs/02_torchinfo.ipynb @@ -0,0 +1,553 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In this notebook we load a saved dreamer, and run it, to look at params, speed and improve hackability" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "%load_ext autoreload\n", + "%autoreload 2" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Loading textures from cache\n" + ] + } + ], + "source": [ + "# TODO make this a proper package\n", + "import os, sys\n", + "sys.path.append('..')\n", + "\n", + "\n", + "from dreamer import parse_args, main, make_env, make_dataset, count_steps,Dreamer" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "['../dreamer.py', '--configs', 'craftax_small', '--logdir', '../logdir/craftax_small']\n" + ] + }, + { + "data": { + "text/plain": [ + "Namespace(act='SiLU', action_repeat=1, actor={'layers': 3, 'dist': 'onehot', 'entropy': 0.0003, 'unimix_ratio': 0.01, 'std': 'none', 'min_std': 0.1, 'max_std': 1.0, 'temp': 0.1, 'lr': 3e-05, 'eps': 1e-05, 'grad_clip': 100.0, 'outscale': 1.0}, batch_length=16, batch_size=128, compile=True, cont_head={'layers': 3, 'loss_scale': 1.0, 'outscale': 1.0}, critic={'layers': 2, 'dist': 'symlog_disc', 'slow_target': True, 'slow_target_update': 1, 'slow_target_fraction': 0.02, 'lr': 3e-05, 'eps': 1e-05, 'grad_clip': 100.0, 'outscale': 0.0}, dataset_size=1000000, debug=False, decoder={'mlp_keys': 'state', 'cnn_keys': '$^', 'act': 'SiLU', 'norm': True, 'cnn_depth': 32, 'kernel_size': 4, 'minres': 4, 'mlp_layers': 3, 'mlp_units': 256, 'cnn_sigmoid': False, 'image_dist': 'mse', 'vector_dist': 'symlog_mse', 'outscale': 1.0}, deterministic_run=False, device='cuda:0', disag_action_cond=False, disag_layers=4, disag_log=True, disag_models=10, disag_offset=1, disag_target='stoch', disag_units=400, discount=0.997, discount_lambda=0.95, dyn_deter=512, dyn_discrete=32, dyn_hidden=256, dyn_mean_act='none', dyn_min_std=0.1, dyn_rec_depth=1, dyn_scale=0.5, dyn_std_act='sigmoid2', dyn_stoch=32, encoder={'mlp_keys': 'state', 'cnn_keys': '$^', 'act': 'SiLU', 'norm': True, 'cnn_depth': 32, 'kernel_size': 4, 'minres': 4, 'mlp_layers': 3, 'mlp_units': 256, 'symlog_inputs': True}, envs=1, eval_episode_num=10, eval_every=10000.0, eval_state_mean=False, evaldir=None, expl_behavior='greedy', expl_extr_scale=0.0, expl_intr_scale=1.0, expl_until=0, grad_clip=1000, grad_heads=('decoder', 'reward', 'cont'), grayscale=False, imag_gradient='reinforce', imag_gradient_mix=0.0, imag_horizon=15, initial='learned', kl_free=1.0, log_every=10000.0, logdir='../logdir/craftax_small', model_lr=0.0001, norm=True, offline_evaldir='', offline_traindir='', opt='adam', opt_eps=1e-08, parallel=False, precision=32, prefill=2500, pretrain=100, rep_scale=0.1, reset_every=0, reward_EMA=True, reward_head={'layers': 3, 'dist': 'symlog_disc', 'loss_scale': 1.0, 'outscale': 0.0}, seed=0, size=(64, 64), step=1000000.0, steps=1000000.0, task='craftax_Craftax-Symbolic-AutoReset-v1', time_limit=1000, train_ratio=512, traindir=None, unimix_ratio=0.01, units=512, value={'layers': 3}, video_pred_log=False, weight_decay=0.0)" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# emulate cli\n", + "argv = f\"../dreamer.py --configs craftax_small --logdir ../logdir/craftax_small\"\n", + "argv = argv.split()\n", + "print(argv)\n", + "config = parse_args(argv)\n", + "config" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\u001b[32m2024-06-06 13:35:50.147\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m\u001b[0m:\u001b[36m27\u001b[0m - \u001b[1mLogdir ../logdir/craftax_small\u001b[0m\n", + "\u001b[32m2024-06-06 13:35:50.153\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m\u001b[0m:\u001b[36m36\u001b[0m - \u001b[1mCreate envs.\u001b[0m\n", + "\u001b[32m2024-06-06 13:36:42.176\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m\u001b[0m:\u001b[36m57\u001b[0m - \u001b[1mAction Space Box(0.0, 1.0, (43,), float32)\u001b[0m\n", + "\u001b[32m2024-06-06 13:36:42.178\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m\u001b[0m:\u001b[36m63\u001b[0m - \u001b[1mPrefill dataset (0 steps).\u001b[0m\n", + "\u001b[32m2024-06-06 13:36:42.180\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m\u001b[0m:\u001b[36m92\u001b[0m - \u001b[1mLogger: (128521 steps).\u001b[0m\n", + "\u001b[32m2024-06-06 13:36:42.180\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m\u001b[0m:\u001b[36m94\u001b[0m - \u001b[1mSimulate agent.\u001b[0m\n" + ] + } + ], + "source": [ + "from loguru import logger\n", + "from tqdm.auto import tqdm\n", + "import pathlib\n", + "\n", + "import torch\n", + "from torch import nn\n", + "from torch import distributions as torchd\n", + "\n", + "import exploration as expl\n", + "import models\n", + "import tools\n", + "import envs.wrappers as wrappers\n", + "from parallel import Parallel, Damy\n", + "\n", + "# from main\n", + "tools.set_seed_everywhere(config.seed)\n", + "if config.deterministic_run:\n", + " tools.enable_deterministic_run()\n", + "logdir = pathlib.Path(config.logdir).expanduser()\n", + "config.traindir = config.traindir or logdir / \"train_eps\"\n", + "config.evaldir = config.evaldir or logdir / \"eval_eps\"\n", + "config.steps //= config.action_repeat\n", + "config.eval_every //= config.action_repeat\n", + "config.log_every //= config.action_repeat\n", + "config.time_limit //= config.action_repeat\n", + "\n", + "logger.info(f\"Logdir {logdir}\")\n", + "logdir.mkdir(parents=True, exist_ok=True)\n", + "config.traindir.mkdir(parents=True, exist_ok=True)\n", + "config.evaldir.mkdir(parents=True, exist_ok=True)\n", + "step = count_steps(config.traindir)\n", + "# step in logger is environmental step\n", + "tlogger = tools.Logger(logdir, config.action_repeat * step)\n", + "logger.add(logdir/\"logger.log\")\n", + "\n", + "logger.info(\"Create envs.\")\n", + "if config.offline_traindir:\n", + " directory = config.offline_traindir.format(**vars(config))\n", + "else:\n", + " directory = config.traindir\n", + "train_eps = tools.load_episodes(directory, limit=config.dataset_size)\n", + "if config.offline_evaldir:\n", + " directory = config.offline_evaldir.format(**vars(config))\n", + "else:\n", + " directory = config.evaldir\n", + "eval_eps = tools.load_episodes(directory, limit=1)\n", + "make = lambda mode, id: make_env(config, mode, id)\n", + "train_envs = [make(\"train\", i) for i in range(config.envs)]\n", + "eval_envs = [make(\"eval\", i) for i in range(config.envs)]\n", + "if config.parallel:\n", + " train_envs = [Parallel(env, \"process\") for env in train_envs]\n", + " eval_envs = [Parallel(env, \"process\") for env in eval_envs]\n", + "else:\n", + " train_envs = [Damy(env) for env in train_envs]\n", + " eval_envs = [Damy(env) for env in eval_envs]\n", + "acts = train_envs[0].action_space\n", + "logger.info(f\"Action Space {acts}\" )\n", + "config.num_actions = acts.n if hasattr(acts, \"n\") else acts.shape[0]\n", + "\n", + "state = None\n", + "if not config.offline_traindir:\n", + " prefill = max(0, config.prefill - count_steps(config.traindir))\n", + " logger.info(f\"Prefill dataset ({prefill} steps).\")\n", + " if hasattr(acts, \"discrete\"):\n", + " random_actor = tools.OneHotDist(\n", + " torch.zeros(config.num_actions).repeat(config.envs, 1)\n", + " )\n", + " else:\n", + " random_actor = torchd.independent.Independent(\n", + " torchd.uniform.Uniform(\n", + " torch.Tensor(acts.low).repeat(config.envs, 1),\n", + " torch.Tensor(acts.high).repeat(config.envs, 1),\n", + " ),\n", + " 1,\n", + " )\n", + "\n", + " def random_agent(o, d, s):\n", + " action = random_actor.sample()\n", + " logprob = random_actor.log_prob(action)\n", + " return {\"action\": action, \"logprob\": logprob}, None\n", + "\n", + " state = tools.simulate(\n", + " random_agent,\n", + " train_envs,\n", + " train_eps,\n", + " config.traindir,\n", + " tlogger,\n", + " limit=config.dataset_size,\n", + " steps=prefill,\n", + " )\n", + " tlogger.step += prefill * config.action_repeat\n", + " logger.info(f\"Logger: ({tlogger.step} steps).\")\n", + "\n", + "logger.info(\"Simulate agent.\")\n", + "train_dataset = make_dataset(train_eps, config)\n", + "eval_dataset = make_dataset(eval_eps, config)" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\u001b[32m2024-06-06 13:38:20.651\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mnetworks\u001b[0m:\u001b[36m__init__\u001b[0m:\u001b[36m323\u001b[0m - \u001b[1mEncoder CNN shapes: {}\u001b[0m\n", + "\u001b[32m2024-06-06 13:38:20.651\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mnetworks\u001b[0m:\u001b[36m__init__\u001b[0m:\u001b[36m324\u001b[0m - \u001b[1mEncoder MLP shapes: {'state': (16536,)}\u001b[0m\n", + "\u001b[32m2024-06-06 13:38:20.751\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mnetworks\u001b[0m:\u001b[36m__init__\u001b[0m:\u001b[36m390\u001b[0m - \u001b[1mDecoder CNN shapes: {}\u001b[0m\n", + "\u001b[32m2024-06-06 13:38:20.751\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mnetworks\u001b[0m:\u001b[36m__init__\u001b[0m:\u001b[36m391\u001b[0m - \u001b[1mDecoder MLP shapes: {'state': (16536,)}\u001b[0m\n", + "\u001b[32m2024-06-06 13:38:20.813\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mmodels\u001b[0m:\u001b[36m__init__\u001b[0m:\u001b[36m102\u001b[0m - \u001b[1mOptimizer model_opt has 15732120 variables.\u001b[0m\n", + "\u001b[32m2024-06-06 13:38:20.836\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mmodels\u001b[0m:\u001b[36m__init__\u001b[0m:\u001b[36m281\u001b[0m - \u001b[1mOptimizer actor_opt has 1335851 variables.\u001b[0m\n", + "\u001b[32m2024-06-06 13:38:20.837\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mmodels\u001b[0m:\u001b[36m__init__\u001b[0m:\u001b[36m292\u001b[0m - \u001b[1mOptimizer value_opt has 1181439 variables.\u001b[0m\n", + "\u001b[32m2024-06-06 13:38:21.032\u001b[0m | \u001b[33m\u001b[1mWARNING \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m\u001b[0m:\u001b[36m17\u001b[0m - \u001b[33m\u001b[1mLoaded model from ../logdir/craftax_small/latest.pt\u001b[0m\n" + ] + } + ], + "source": [ + "config = parse_args(argv)\n", + "config.num_actions = acts.n if hasattr(acts, \"n\") else acts.shape[0]\n", + "agent = Dreamer(\n", + " train_envs[0].observation_space,\n", + " train_envs[0].action_space,\n", + " config,\n", + " tlogger,\n", + " train_dataset,\n", + ").to(config.device)\n", + "# print(agent)\n", + "agent.requires_grad_(requires_grad=False)\n", + "if (logdir / \"latest.pt\").exists():\n", + " checkpoint = torch.load(logdir / \"latest.pt\")\n", + " agent.load_state_dict(checkpoint[\"agent_state_dict\"])\n", + " tools.recursively_load_optim_state_dict(agent, checkpoint[\"optims_state_dict\"])\n", + " agent._should_pretrain._once = False\n", + " logger.warning(f\"Loaded model from {logdir / 'latest.pt'}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Now lets play" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "(0, 0, array([ True]), array([0], dtype=int32), [None], None, [0])" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "assert state is not None\n", + "import numpy as np\n", + "\n", + "state" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "from tools import convert, add_to_cache\n", + "envs = train_envs\n", + "cache = train_eps\n", + "\n", + "step, episode = 0, 0\n", + "done = np.ones(len(envs), bool)\n", + "length = np.zeros(len(envs), np.int32)\n", + "obs = [None] * len(envs)\n", + "agent_state = None\n", + "reward = [0] * len(envs)\n", + "\n", + "indices = [index for index, d in enumerate(done) if d]\n", + "results = [envs[i].reset() for i in indices]\n", + "results = [r() for r in results]\n", + "for index, result in zip(indices, results):\n", + " t = result.copy()\n", + " t = {k: convert(v) for k, v in t.items()}\n", + " # action will be added to transition in add_to_cache\n", + " t[\"reward\"] = 0.0\n", + " t[\"discount\"] = 1.0\n", + " # initial state should be added to cache\n", + " add_to_cache(cache, envs[index].id, t)\n", + " # replace obs with done by initial state\n", + " obs[index] = result\n", + "# step agents" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\u001b[32m2024-06-06 13:38:34.000\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mtools\u001b[0m:\u001b[36mwrite\u001b[0m:\u001b[36m85\u001b[0m - \u001b[1m[128521] model_loss \u001b[31m22.2\u001b[0m\u001b[1m / model_grad_norm \u001b[31m14.4\u001b[0m\u001b[1m / state_loss \u001b[31m17.4\u001b[0m\u001b[1m / reward_loss \u001b[31m0.1\u001b[0m\u001b[1m / cont_loss \u001b[31m0.0\u001b[0m\u001b[1m / kl_free \u001b[31m1.0\u001b[0m\u001b[1m / dyn_scale \u001b[31m0.5\u001b[0m\u001b[1m / rep_scale \u001b[31m0.1\u001b[0m\u001b[1m / dyn_loss \u001b[31m7.8\u001b[0m\u001b[1m / rep_loss \u001b[31m7.8\u001b[0m\u001b[1m / kl \u001b[31m7.7\u001b[0m\u001b[1m / prior_ent \u001b[31m48.4\u001b[0m\u001b[1m / post_ent \u001b[31m40.7\u001b[0m\u001b[1m / normed_target_mean \u001b[31m0.4\u001b[0m\u001b[1m / normed_target_std \u001b[31m0.3\u001b[0m\u001b[1m / normed_target_min \u001b[31m-0.3\u001b[0m\u001b[1m / normed_target_max \u001b[31m1.8\u001b[0m\u001b[1m / EMA_005 \u001b[31m12.3\u001b[0m\u001b[1m / EMA_095 \u001b[31m26.4\u001b[0m\u001b[1m / value_mean \u001b[31m18.2\u001b[0m\u001b[1m / value_std \u001b[31m4.3\u001b[0m\u001b[1m / value_min \u001b[31m10.1\u001b[0m\u001b[1m / value_max \u001b[31m31.1\u001b[0m\u001b[1m / target_mean \u001b[31m18.4\u001b[0m\u001b[1m / target_std \u001b[31m4.7\u001b[0m\u001b[1m / target_min \u001b[31m8.4\u001b[0m\u001b[1m / target_max \u001b[31m37.8\u001b[0m\u001b[1m / imag_reward_mean \u001b[31m0.0\u001b[0m\u001b[1m / imag_reward_std \u001b[31m0.1\u001b[0m\u001b[1m / imag_reward_min \u001b[31m-0.2\u001b[0m\u001b[1m / imag_reward_max \u001b[31m1.0\u001b[0m\u001b[1m / imag_action_mean \u001b[31m10.0\u001b[0m\u001b[1m / imag_action_std \u001b[31m12.9\u001b[0m\u001b[1m / imag_action_min \u001b[31m0.0\u001b[0m\u001b[1m / imag_action_max \u001b[31m42.0\u001b[0m\u001b[1m / actor_entropy \u001b[31m0.9\u001b[0m\u001b[1m / actor_loss \u001b[31m0.1\u001b[0m\u001b[1m / actor_grad_norm \u001b[31m0.5\u001b[0m\u001b[1m / value_loss \u001b[31m1.3\u001b[0m\u001b[1m / value_grad_norm \u001b[31m0.9\u001b[0m\u001b[1m / update_count \u001b[31m1.0\u001b[0m\u001b[1m / fps \u001b[31m0.0\u001b[0m\u001b[1m\u001b[0m\n" + ] + } + ], + "source": [ + "# from tools.simulate\n", + "\n", + "# step\n", + "# step, episode, done, length, obs, agent_state, reward = state\n", + "obs = {k: np.stack([o[k] for o in obs]) for k in obs[0] if \"log_\" not in k}\n", + "action, agent_state = agent(obs, done, agent_state)" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "=====================================================================================\n", + "Layer (type:depth-idx) Param #\n", + "=====================================================================================\n", + "Dreamer --\n", + "├─OptimizedModule: 1-1 --\n", + "│ └─WorldModel: 2-1 --\n", + "│ │ └─MultiEncoder: 3-1 (4,365,824)\n", + "│ │ └─RSSM: 3-2 (3,831,808)\n", + "│ │ └─ModuleDict: 3-3 (7,534,488)\n", + "├─OptimizedModule: 1-2 --\n", + "│ └─ImagBehavior: 2-2 15,732,120\n", + "│ │ └─WorldModel: 3-4 (recursive)\n", + "│ │ └─MLP: 3-5 (1,335,851)\n", + "│ │ └─MLP: 3-6 (1,181,439)\n", + "│ │ └─MLP: 3-7 (1,181,439)\n", + "├─OptimizedModule: 1-3 (recursive)\n", + "│ └─ImagBehavior: 2-3 (recursive)\n", + "│ │ └─WorldModel: 3-8 (recursive)\n", + "│ │ └─MLP: 3-9 (recursive)\n", + "│ │ └─MLP: 3-10 (recursive)\n", + "│ │ └─MLP: 3-11 (recursive)\n", + "=====================================================================================\n", + "Total params: 35,162,969\n", + "Trainable params: 0\n", + "Non-trainable params: 35,162,969\n", + "=====================================================================================" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from torchinfo import summary\n", + "\n", + "summary(agent, input=(obs, done, agent_state), depth=3)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Fine grained torchinfo" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [], + "source": [ + "wm = agent._wm\n", + "data = next(agent._dataset) \n", + "# self._train()\n", + "# post, context, mets = wm._train(data)\n", + "data = wm.preprocess(data)\n", + "embed = wm.encoder(data)\n", + "post, prior = wm.dynamics.observe(\n", + " embed, data[\"action\"], data[\"is_first\"]\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "==========================================================================================\n", + "Layer (type:depth-idx) Output Shape Param #\n", + "==========================================================================================\n", + "MultiEncoder [128, 16, 256] --\n", + "├─MLP: 1-1 [128, 16, 256] --\n", + "│ └─Sequential: 2-1 [128, 16, 256] --\n", + "│ │ └─Linear: 3-1 [128, 16, 256] (4,233,216)\n", + "│ │ └─LayerNorm: 3-2 [128, 16, 256] (512)\n", + "│ │ └─SiLU: 3-3 [128, 16, 256] --\n", + "│ │ └─Linear: 3-4 [128, 16, 256] (65,536)\n", + "│ │ └─LayerNorm: 3-5 [128, 16, 256] (512)\n", + "│ │ └─SiLU: 3-6 [128, 16, 256] --\n", + "│ │ └─Linear: 3-7 [128, 16, 256] (65,536)\n", + "│ │ └─LayerNorm: 3-8 [128, 16, 256] (512)\n", + "│ │ └─SiLU: 3-9 [128, 16, 256] --\n", + "==========================================================================================\n", + "Total params: 4,365,824\n", + "Trainable params: 0\n", + "Non-trainable params: 4,365,824\n", + "Total mult-adds (M): 558.83\n", + "==========================================================================================\n", + "Input size (MB): 487.31\n", + "Forward/backward pass size (MB): 25.17\n", + "Params size (MB): 17.46\n", + "Estimated Total Size (MB): 529.94\n", + "==========================================================================================" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "summary(wm.encoder, input_data=(data,), depth=3, col_names=[\"output_size\", \"num_params\", ])" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "decoder\n", + "==========================================================================================\n", + "Layer (type:depth-idx) Output Shape Param #\n", + "==========================================================================================\n", + "MultiDecoder -- --\n", + "├─MLP: 1-1 -- --\n", + "│ └─Sequential: 2-1 [128, 16, 256] --\n", + "│ │ └─Linear: 3-1 [128, 16, 256] (393,216)\n", + "│ │ └─LayerNorm: 3-2 [128, 16, 256] (512)\n", + "│ │ └─SiLU: 3-3 [128, 16, 256] --\n", + "│ │ └─Linear: 3-4 [128, 16, 256] (65,536)\n", + "│ │ └─LayerNorm: 3-5 [128, 16, 256] (512)\n", + "│ │ └─SiLU: 3-6 [128, 16, 256] --\n", + "│ │ └─Linear: 3-7 [128, 16, 256] (65,536)\n", + "│ │ └─LayerNorm: 3-8 [128, 16, 256] (512)\n", + "│ │ └─SiLU: 3-9 [128, 16, 256] --\n", + "│ └─ModuleDict: 2-2 -- --\n", + "│ │ └─Linear: 3-10 [128, 16, 16536] (4,249,752)\n", + "==========================================================================================\n", + "Total params: 4,775,576\n", + "Trainable params: 0\n", + "Non-trainable params: 4,775,576\n", + "Total mult-adds (M): 611.27\n", + "==========================================================================================\n", + "Input size (MB): 12.58\n", + "Forward/backward pass size (MB): 296.09\n", + "Params size (MB): 19.10\n", + "Estimated Total Size (MB): 327.78\n", + "==========================================================================================\n", + "Summary Failed for reward Failed to run torchinfo. See above stack traces for more details. Executed layers up to: [Sequential: 1, Linear: 2, LayerNorm: 2, SiLU: 2, Linear: 2, LayerNorm: 2, SiLU: 2, Linear: 2, LayerNorm: 2, SiLU: 2, Linear: 1]\n", + "Summary Failed for cont Failed to run torchinfo. See above stack traces for more details. Executed layers up to: [Sequential: 1, Linear: 2, LayerNorm: 2, SiLU: 2, Linear: 2, LayerNorm: 2, SiLU: 2, Linear: 2, LayerNorm: 2, SiLU: 2, Linear: 1]\n" + ] + } + ], + "source": [ + "# heads\n", + "feat = wm.dynamics.get_feat(post)\n", + "for name, head in wm.heads.items():\n", + " try:\n", + " o = summary(head, input_data=(feat,), depth=3, col_names=[\"output_size\", \"num_params\", ])\n", + " print(name)\n", + " print(o)\n", + " except Exception as e:\n", + " print(f\"Summary Failed for {name} {e}\")\n", + " continue" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [], + "source": [ + "# fail as no call method\n", + "# summary(wm.dynamics, input_data=(embed, data[\"action\"], data[\"is_first\"]), depth=3, col_names=[\"output_size\", \"num_params\", ])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "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 +} diff --git a/networks.py b/networks.py index 6f22d0e..7661f7f 100644 --- a/networks.py +++ b/networks.py @@ -320,8 +320,8 @@ class MultiEncoder(nn.Module): for k, v in shapes.items() if len(v) in (1, 2) and re.match(mlp_keys, k) } - logger.info("Encoder CNN shapes:", self.cnn_shapes) - logger.info("Encoder MLP shapes:", self.mlp_shapes) + logger.info("Encoder CNN shapes: {}", self.cnn_shapes) + logger.info("Encoder MLP shapes: {}", self.mlp_shapes) self.outdim = 0 if self.cnn_shapes: @@ -387,8 +387,8 @@ class MultiDecoder(nn.Module): for k, v in shapes.items() if len(v) in (1, 2) and re.match(mlp_keys, k) } - logger.info("Decoder CNN shapes: %s", self.cnn_shapes) - logger.info("Decoder MLP shapes: %s", self.mlp_shapes) + logger.info("Decoder CNN shapes: {}", self.cnn_shapes) + logger.info("Decoder MLP shapes: {}", self.mlp_shapes) if self.cnn_shapes: some_shape = list(self.cnn_shapes.values())[0] diff --git a/poetry.lock b/poetry.lock index d9ceae6..7cee8d2 100644 --- a/poetry.lock +++ b/poetry.lock @@ -3578,6 +3578,17 @@ type = "legacy" url = "https://download.pytorch.org/whl/cu121" reference = "pytorch" +[[package]] +name = "torchinfo" +version = "1.8.0" +description = "Model summary in PyTorch, based off of the original torchsummary." +optional = false +python-versions = ">=3.7" +files = [ + {file = "torchinfo-1.8.0-py3-none-any.whl", hash = "sha256:2e911c2918603f945c26ff21a3a838d12709223dc4ccf243407bce8b6e897b46"}, + {file = "torchinfo-1.8.0.tar.gz", hash = "sha256:72e94b0e9a3e64dc583a8e5b7940b8938a1ac0f033f795457f27e6f4e7afa2e9"}, +] + [[package]] name = "tornado" version = "6.4" @@ -3836,4 +3847,4 @@ test = ["big-O", "jaraco.functools", "jaraco.itertools", "jaraco.test", "more-it [metadata] lock-version = "2.0" python-versions = "^3.9" -content-hash = "8d04aef5b114f7ae76dc03bc61d308f1b239d390b0c71fab7d0c8f467cc95dd4" +content-hash = "0275da73363d94f6a5cdadc9662c1b254ef50310aabce3aa663552aa4802b001" diff --git a/pyproject.toml b/pyproject.toml index f378589..5af5591 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,7 @@ imageio = "^2.34.1" craftax = {path = "/media/wassname/SGIronWolf/projects5/2024/Craftax", develop = true } # craftax = {git = "https://github.com/wassname/Craftax" , develop = true } chex = "^0.1.86" +torchinfo = "^1.8.0" [tool.poetry.group.dev.dependencies] ipywidgets = "^8.1.3"