exp_dreamer_v3_dmc.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check obsevation space:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "domain name:  acrobot \n",
      "\n",
      "task name:  swingup \n",
      "\n",
      "domain name:  acrobot \n",
      "\n",
      "task name:  swingup_sparse \n",
      "\n",
      "domain name:  ball_in_cup \n",
      "\n",
      "task name:  catch \n",
      "\n",
      "domain name:  cartpole \n",
      "\n",
      "task name:  balance \n",
      "\n",
      "domain name:  cartpole \n",
      "\n",
      "task name:  balance_sparse \n",
      "\n",
      "domain name:  cartpole \n",
      "\n",
      "task name:  swingup \n",
      "\n",
      "domain name:  cartpole \n",
      "\n",
      "task name:  swingup_sparse \n",
      "\n",
      "domain name:  cheetah \n",
      "\n",
      "task name:  run \n",
      "\n",
      "domain name:  finger \n",
      "\n",
      "task name:  spin \n",
      "\n",
      "domain name:  finger \n",
      "\n",
      "task name:  turn_easy \n",
      "\n",
      "domain name:  finger \n",
      "\n",
      "task name:  turn_hard \n",
      "\n",
      "domain name:  fish \n",
      "\n",
      "task name:  upright \n",
      "\n",
      "domain name:  fish \n",
      "\n",
      "task name:  swim \n",
      "\n",
      "domain name:  hopper \n",
      "\n",
      "task name:  stand \n",
      "\n",
      "domain name:  hopper \n",
      "\n",
      "task name:  hop \n",
      "\n",
      "domain name:  humanoid \n",
      "\n",
      "task name:  stand \n",
      "\n",
      "domain name:  humanoid \n",
      "\n",
      "task name:  walk \n",
      "\n",
      "domain name:  humanoid \n",
      "\n",
      "task name:  run \n",
      "\n",
      "domain name:  manipulator \n",
      "\n",
      "task name:  bring_ball \n",
      "\n",
      "domain name:  pendulum \n",
      "\n",
      "task name:  swingup \n",
      "\n",
      "domain name:  point_mass \n",
      "\n",
      "task name:  easy \n",
      "\n",
      "domain name:  reacher \n",
      "\n",
      "task name:  easy \n",
      "\n",
      "domain name:  reacher \n",
      "\n",
      "task name:  hard \n",
      "\n",
      "domain name:  swimmer \n",
      "\n",
      "task name:  swimmer6 \n",
      "\n",
      "domain name:  swimmer \n",
      "\n",
      "task name:  swimmer15 \n",
      "\n",
      "domain name:  walker \n",
      "\n",
      "task name:  stand \n",
      "\n",
      "domain name:  walker \n",
      "\n",
      "task name:  walk \n",
      "\n",
      "domain name:  walker \n",
      "\n",
      "task name:  run \n",
      "\n"
     ]
    }
   ],
   "source": [
    "from dm_control import suite\n",
    "# Iterate over a task set:\n",
    "for domain_name , task_name in suite.BENCHMARKING: \n",
    "    print('domain name: ', domain_name, '\\n')\n",
    "    print('task name: ', task_name, '\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-07-03 11:01:07,815][absl][INFO] - MUJOCO_GL is not set, so an OpenGL backend will be chosen automatically.\n",
      "[2024-07-03 11:01:07,870][absl][INFO] - Successfully imported OpenGL backend: glfw\n",
      "[2024-07-03 11:01:07,957][absl][INFO] - MuJoCo library version is: 3.1.6\n",
      "Error executing job with overrides: ['agent=dreamer_v3', 'env=dmc', 'env.wrapper.domain_name=pendulum', 'env.wrapper.task_name=swingup', 'env.grayscale=True', 'algo.cnn_keys.encoder=[rgb]', 'algo.mlp_keys.encoder=[state]']\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/jianheng/.local/share/ov/pkg/isaac-sim-4.0.0/sheeprl/examples/observation_space.py\", line 42, in <module>\n",
      "    main()\n",
      "  File \"/home/jianheng/miniconda3/envs/isaac-sim/lib/python3.10/site-packages/hydra/main.py\", line 90, in decorated_main\n",
      "    _run_hydra(\n",
      "  File \"/home/jianheng/miniconda3/envs/isaac-sim/lib/python3.10/site-packages/hydra/_internal/utils.py\", line 394, in _run_hydra\n",
      "    _run_app(\n",
      "  File \"/home/jianheng/miniconda3/envs/isaac-sim/lib/python3.10/site-packages/hydra/_internal/utils.py\", line 457, in _run_app\n",
      "    run_and_report(\n",
      "  File \"/home/jianheng/miniconda3/envs/isaac-sim/lib/python3.10/site-packages/hydra/_internal/utils.py\", line 222, in run_and_report\n",
      "    raise ex\n",
      "  File \"/home/jianheng/miniconda3/envs/isaac-sim/lib/python3.10/site-packages/hydra/_internal/utils.py\", line 219, in run_and_report\n",
      "    return func()\n",
      "  File \"/home/jianheng/miniconda3/envs/isaac-sim/lib/python3.10/site-packages/hydra/_internal/utils.py\", line 458, in <lambda>\n",
      "    lambda: hydra.run(\n",
      "  File \"/home/jianheng/miniconda3/envs/isaac-sim/lib/python3.10/site-packages/hydra/_internal/hydra.py\", line 132, in run\n",
      "    _ = ret.return_value\n",
      "  File \"/home/jianheng/miniconda3/envs/isaac-sim/lib/python3.10/site-packages/hydra/core/utils.py\", line 260, in return_value\n",
      "    raise self._return_value\n",
      "  File \"/home/jianheng/miniconda3/envs/isaac-sim/lib/python3.10/site-packages/hydra/core/utils.py\", line 186, in run_job\n",
      "    ret.return_value = task_function(task_cfg)\n",
      "  File \"/home/jianheng/.local/share/ov/pkg/isaac-sim-4.0.0/sheeprl/examples/observation_space.py\", line 28, in main\n",
      "    env: gym.Env = make_env(cfg, cfg.seed, 0)()\n",
      "  File \"/home/jianheng/miniconda3/envs/isaac-sim/lib/python3.10/site-packages/sheeprl/utils/env.py\", line 211, in thunk\n",
      "    if cfg.env.actions_as_observation.num_stack > 0 and \"diambra\" not in cfg.env.wrapper._target_:\n",
      "AttributeError: 'NoneType' object has no attribute 'num_stack'\n"
     ]
    }
   ],
   "source": [
    "!HYDRA_FULL_ERROR=1 python examples/observation_space.py agent=dreamer_v3 env=dmc env.wrapper.domain_name=pendulum env.wrapper.task_name=swingup env.grayscale=True algo.cnn_keys.encoder=[rgb] algo.mlp_keys.encoder=[state]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Experiments"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. Using grascale image as observation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/lightning/fabric/__init__.py:41: Deprecated call to `pkg_resources.declare_namespace('lightning.fabric')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/pkg_resources/__init__.py:2350: DeprecationWarning: Deprecated call to `pkg_resources.declare_namespace('lightning')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "  declare_namespace(parent)\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/lightning/pytorch/__init__.py:37: Deprecated call to `pkg_resources.declare_namespace('lightning.pytorch')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/pkg_resources/__init__.py:2350: DeprecationWarning: Deprecated call to `pkg_resources.declare_namespace('lightning')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "  declare_namespace(parent)\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/lightning_fabric/__init__.py:41: Deprecated call to `pkg_resources.declare_namespace('lightning_fabric')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/pytorch_lightning/__init__.py:37: Deprecated call to `pkg_resources.declare_namespace('pytorch_lightning')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "\u001b[2mCONFIG\u001b[0m\n",
      "\u001b[2m├── \u001b[0m\u001b[2malgo\u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m└── \u001b[0m\u001b[2;91;40mname\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mdreamer_v3                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mtotal_steps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m10000                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mper_rank_batch_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m32                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mrun_test\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcnn_keys\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mencoder\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mrgb                                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdecoder\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mrgb                                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mmlp_keys\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mencoder\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m[]                                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdecoder\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m[]                                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mworld_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40moptimizer\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.optim.Adam                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlr\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.0001                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0e-08                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mweight_decay\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mbetas\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.9                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.999                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdiscrete_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m32                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mstochastic_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m32                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkl_dynamic\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.5                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkl_representation\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.1                                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkl_free_nats\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkl_regularizer\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mcontinue_scale_factor\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mclip_gradients\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1000.0                                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdecoupled_rssm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mlearnable_initial_recurrent_state\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mencoder\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcnn_channels_multiplier\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m32                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcnn_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcnn_layer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNormChannelLast                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmlp_layer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mrecurrent_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mrecurrent_state_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mtransition_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mhidden_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mrepresentation_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mhidden_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mobservation_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcnn_channels_multiplier\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m32                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcnn_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcnn_layer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNormChannelLast                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmlp_layer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mreward_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mbins\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m255                                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdiscount_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlearnable\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mactor\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                  \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40moptimizer\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.optim.Adam                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlr\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m8.0e-05                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0e-05                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mweight_decay\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mbetas\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.9                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.999                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.algos.dreamer_v3.agent.Actor                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40ment_coef\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.0003                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mmin_std\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.1                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mmax_std\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40minit_std\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2.0                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mclip_gradients\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m100.0                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40munimix\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.01                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40maction_clip\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mmoments\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdecay\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.99                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmax\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mpercentile\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mlow\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.05                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mhigh\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.95                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcritic\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40moptimizer\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.optim.Adam                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlr\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m8.0e-05                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0e-05                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mweight_decay\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mbetas\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.9                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.999                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mper_rank_target_network_update_freq\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mtau\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.02                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mbins\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m255                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mclip_gradients\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m100.0                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mgamma\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.996996996996997                                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mlmbda\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.95                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mhorizon\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m15                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mreplay_ratio\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.5                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mlearning_starts\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1000                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mper_rank_pretrain_steps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mper_rank_sequence_length\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m64                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcnn_layer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNormChannelLast                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mmlp_layer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                                  \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcnn_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                                  \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40munimix\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.01                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mhafner_initialization\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mplayer\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdiscrete_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m32                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m                                                                        \u001b[0m\n",
      "\u001b[2m├── \u001b[0m\u001b[2mbuffer\u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m└── \u001b[0m\u001b[2;91;40msize\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m10000                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mmemmap\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mvalidate_args\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mfrom_numpy\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcheckpoint\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m                                                                        \u001b[0m\n",
      "\u001b[2m├── \u001b[0m\u001b[2mcheckpoint\u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m└── \u001b[0m\u001b[2;91;40mevery\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1000                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mresume_from\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mnull                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40msave_last\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mkeep_last\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m5                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m                                                                        \u001b[0m\n",
      "\u001b[2m├── \u001b[0m\u001b[2menv\u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m└── \u001b[0m\u001b[2;91;40mid\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mwalker_walk                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mnum_envs\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m4                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mframe_stack\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40msync_env\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mscreen_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m64                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40maction_repeat\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mgrayscale\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mclip_rewards\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcapture_video\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mframe_stack_dilation\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mmax_episode_steps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m-1                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mreward_as_observation\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mwrapper\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.envs.dmc.DMCWrapper                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdomain_name\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mwalker                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mtask_name\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mwalk                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mwidth\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m64                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mheight\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m64                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mseed\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mnull                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mfrom_pixels\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mfrom_vectors\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m                                                                        \u001b[0m\n",
      "\u001b[2m├── \u001b[0m\u001b[2mfabric\u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m└── \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mlightning.fabric.Fabric                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mdevices\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mnum_nodes\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mstrategy\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mauto                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40maccelerator\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mcuda                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mprecision\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mbf16-mixed                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcallbacks\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.utils.callback.CheckpointCallback                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkeep_last\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m5                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m                                                                        \u001b[0m\n",
      "\u001b[2m└── \u001b[0m\u001b[2mmetric\u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m└── \u001b[0m\u001b[2;91;40mlog_every\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m100                                                          \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mdisable_timer\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                    \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mlog_level\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1                                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                  \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40maggregator\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                             \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.utils.metric.MetricAggregator                       \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mraise_on_missing\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                               \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mmetrics\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mRewards/rew_avg\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                    \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mGame/ep_len_avg\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                    \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/world_model_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                              \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/value_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                    \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/policy_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                   \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/observation_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                              \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/reward_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                   \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/state_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                    \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/continue_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mState/kl\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                           \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mState/post_entropy\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mState/prior_entropy\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mGrads/world_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                  \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mGrads/actor\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                        \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mGrads/critic\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                       \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mlogger\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mlightning.fabric.loggers.TensorBoardLogger                  \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mname\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2024-05-05_14-12-21_dreamer_v3_walker_walk_5                    \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mroot_dir\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mlogs/runs/dreamer_v3/walker_walk                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mversion\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mnull                                                         \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdefault_hp_metric\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                               \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mprefix\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;93;40m'\u001b[0m\u001b[2;93;40m'\u001b[0m\u001b[2;40m                                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40msub_dir\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mnull                                                         \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m                                                                        \u001b[0m\n",
      "[2024-05-05 14:12:21,293][pytorch_lightning.utilities.rank_zero][INFO] - Using bfloat16 Automatic Mixed Precision (AMP)\n",
      "INFO: Seed set to 5\n",
      "[2024-05-05 14:12:21,313][lightning.fabric.utilities.seed][INFO] - Seed set to 5\n",
      "Log dir: logs/runs/dreamer_v3/walker_walk/2024-05-05_14-12-21_dreamer_v3_walker_walk_5/version_0\n",
      "Encoder CNN keys: ['rgb']\n",
      "Encoder MLP keys: []\n",
      "Decoder CNN keys: ['rgb']\n",
      "Decoder MLP keys: []\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/torchmetrics/utilities/prints.py:43: UserWarning: The ``compute`` method of metric MeanMetric was called before the ``update`` method which may lead to errors, as metric states have not yet been updated.\n",
      "  warnings.warn(*args, **kwargs)  # noqa: B028\n",
      "Rank-0: policy_step=2000, reward_env_0=34.312042236328125\n",
      "Rank-0: policy_step=2000, reward_env_1=41.49196243286133\n",
      "Rank-0: policy_step=2000, reward_env_2=33.82015609741211\n",
      "Rank-0: policy_step=2000, reward_env_3=42.34160614013672\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/torchmetrics/utilities/prints.py:43: UserWarning: The ``compute`` method of metric MeanMetric was called before the ``update`` method which may lead to errors, as metric states have not yet been updated.\n",
      "  warnings.warn(*args, **kwargs)  # noqa: B028\n",
      "Rank-0: policy_step=4000, reward_env_0=25.6603946685791\n",
      "Rank-0: policy_step=4000, reward_env_1=69.88949584960938\n",
      "Rank-0: policy_step=4000, reward_env_2=52.5946159362793\n",
      "Rank-0: policy_step=4000, reward_env_3=107.42786407470703\n",
      "Rank-0: policy_step=6000, reward_env_0=98.66031646728516\n",
      "Rank-0: policy_step=6000, reward_env_1=95.7126235961914\n",
      "Rank-0: policy_step=6000, reward_env_2=79.24365234375\n",
      "Rank-0: policy_step=6000, reward_env_3=38.695152282714844\n",
      "Rank-0: policy_step=8000, reward_env_0=110.80186462402344\n",
      "Rank-0: policy_step=8000, reward_env_1=123.6385269165039\n",
      "Rank-0: policy_step=8000, reward_env_2=91.18580627441406\n",
      "Rank-0: policy_step=8000, reward_env_3=108.86883544921875\n",
      "Rank-0: policy_step=10000, reward_env_0=126.08735656738281\n",
      "Rank-0: policy_step=10000, reward_env_1=35.294979095458984\n",
      "Rank-0: policy_step=10000, reward_env_2=110.68358612060547\n",
      "Rank-0: policy_step=10000, reward_env_3=51.870697021484375\n",
      "Test - Reward: 99.27343657301383\n"
     ]
    }
   ],
   "source": [
    "!python sheeprl.py exp=dreamer_v3_dmc_walk_walk_grayscale"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. Using grascale image and vector values as observation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/lightning/fabric/__init__.py:41: Deprecated call to `pkg_resources.declare_namespace('lightning.fabric')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/pkg_resources/__init__.py:2350: DeprecationWarning: Deprecated call to `pkg_resources.declare_namespace('lightning')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "  declare_namespace(parent)\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/lightning/pytorch/__init__.py:37: Deprecated call to `pkg_resources.declare_namespace('lightning.pytorch')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/pkg_resources/__init__.py:2350: DeprecationWarning: Deprecated call to `pkg_resources.declare_namespace('lightning')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "  declare_namespace(parent)\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/lightning_fabric/__init__.py:41: Deprecated call to `pkg_resources.declare_namespace('lightning_fabric')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/pytorch_lightning/__init__.py:37: Deprecated call to `pkg_resources.declare_namespace('pytorch_lightning')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "\u001b[2mCONFIG\u001b[0m\n",
      "\u001b[2m├── \u001b[0m\u001b[2malgo\u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m└── \u001b[0m\u001b[2;91;40mname\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mdreamer_v3                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mtotal_steps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m10000                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mper_rank_batch_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m32                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mrun_test\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcnn_keys\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mencoder\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mrgb                                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdecoder\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mrgb                                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mmlp_keys\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mencoder\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mstate                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdecoder\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mstate                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mworld_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40moptimizer\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.optim.Adam                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlr\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.0001                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0e-08                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mweight_decay\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mbetas\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.9                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.999                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdiscrete_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m32                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mstochastic_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m32                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkl_dynamic\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.5                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkl_representation\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.1                                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkl_free_nats\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkl_regularizer\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mcontinue_scale_factor\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mclip_gradients\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1000.0                                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdecoupled_rssm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mlearnable_initial_recurrent_state\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mencoder\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcnn_channels_multiplier\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m32                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcnn_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcnn_layer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNormChannelLast                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmlp_layer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mrecurrent_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mrecurrent_state_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mtransition_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mhidden_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mrepresentation_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mhidden_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mobservation_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcnn_channels_multiplier\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m32                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcnn_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcnn_layer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNormChannelLast                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmlp_layer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mreward_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mbins\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m255                                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdiscount_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlearnable\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mactor\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                  \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40moptimizer\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.optim.Adam                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlr\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m8.0e-05                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0e-05                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mweight_decay\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mbetas\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.9                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.999                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.algos.dreamer_v3.agent.Actor                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40ment_coef\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.0003                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mmin_std\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.1                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mmax_std\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40minit_std\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2.0                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mclip_gradients\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m100.0                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40munimix\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.01                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40maction_clip\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mmoments\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdecay\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.99                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmax\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mpercentile\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mlow\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.05                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mhigh\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.95                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcritic\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40moptimizer\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.optim.Adam                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlr\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m8.0e-05                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0e-05                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mweight_decay\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mbetas\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.9                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.999                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mper_rank_target_network_update_freq\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mtau\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.02                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mbins\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m255                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mclip_gradients\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m100.0                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mgamma\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.996996996996997                                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mlmbda\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.95                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mhorizon\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m15                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mreplay_ratio\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.5                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mlearning_starts\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1000                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mper_rank_pretrain_steps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mper_rank_sequence_length\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m64                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcnn_layer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNormChannelLast                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mmlp_layer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                                  \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcnn_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                                  \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40munimix\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.01                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mhafner_initialization\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mplayer\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdiscrete_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m32                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m                                                                        \u001b[0m\n",
      "\u001b[2m├── \u001b[0m\u001b[2mbuffer\u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m└── \u001b[0m\u001b[2;91;40msize\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m10000                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mmemmap\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mvalidate_args\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mfrom_numpy\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcheckpoint\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m                                                                        \u001b[0m\n",
      "\u001b[2m├── \u001b[0m\u001b[2mcheckpoint\u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m└── \u001b[0m\u001b[2;91;40mevery\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1000                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mresume_from\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mnull                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40msave_last\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mkeep_last\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m5                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m                                                                        \u001b[0m\n",
      "\u001b[2m├── \u001b[0m\u001b[2menv\u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m└── \u001b[0m\u001b[2;91;40mid\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mwalker_walk                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mnum_envs\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m4                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mframe_stack\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40msync_env\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mscreen_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m64                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40maction_repeat\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mgrayscale\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mclip_rewards\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcapture_video\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mframe_stack_dilation\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mmax_episode_steps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m-1                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mreward_as_observation\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mwrapper\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.envs.dmc.DMCWrapper                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdomain_name\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mwalker                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mtask_name\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mwalk                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mwidth\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m64                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mheight\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m64                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mseed\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mnull                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mfrom_pixels\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mfrom_vectors\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m                                                                        \u001b[0m\n",
      "\u001b[2m├── \u001b[0m\u001b[2mfabric\u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m└── \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mlightning.fabric.Fabric                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mdevices\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mnum_nodes\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mstrategy\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mauto                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40maccelerator\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mcuda                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mprecision\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mbf16-mixed                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcallbacks\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.utils.callback.CheckpointCallback                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkeep_last\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m5                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m                                                                        \u001b[0m\n",
      "\u001b[2m└── \u001b[0m\u001b[2mmetric\u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m└── \u001b[0m\u001b[2;91;40mlog_every\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m100                                                          \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mdisable_timer\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                    \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mlog_level\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1                                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                  \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40maggregator\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                             \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.utils.metric.MetricAggregator                       \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mraise_on_missing\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                               \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mmetrics\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mRewards/rew_avg\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                    \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mGame/ep_len_avg\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                    \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/world_model_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                              \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/value_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                    \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/policy_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                   \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/observation_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                              \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/reward_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                   \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/state_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                    \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/continue_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mState/kl\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                           \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mState/post_entropy\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mState/prior_entropy\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mGrads/world_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                  \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mGrads/actor\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                        \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mGrads/critic\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                       \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mlogger\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mlightning.fabric.loggers.TensorBoardLogger                  \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mname\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2024-05-05_19-47-22_dreamer_v3_walker_walk_5                    \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mroot_dir\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mlogs/runs/dreamer_v3/walker_walk                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mversion\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mnull                                                         \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdefault_hp_metric\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                               \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mprefix\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;93;40m'\u001b[0m\u001b[2;93;40m'\u001b[0m\u001b[2;40m                                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40msub_dir\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mnull                                                         \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m                                                                        \u001b[0m\n",
      "[2024-05-05 19:47:22,692][pytorch_lightning.utilities.rank_zero][INFO] - Using bfloat16 Automatic Mixed Precision (AMP)\n",
      "INFO: Seed set to 5\n",
      "[2024-05-05 19:47:22,711][lightning.fabric.utilities.seed][INFO] - Seed set to 5\n",
      "Log dir: logs/runs/dreamer_v3/walker_walk/2024-05-05_19-47-22_dreamer_v3_walker_walk_5/version_0\n",
      "Encoder CNN keys: ['rgb']\n",
      "Encoder MLP keys: ['state']\n",
      "Decoder CNN keys: ['rgb']\n",
      "Decoder MLP keys: ['state']\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/torchmetrics/utilities/prints.py:43: UserWarning: The ``compute`` method of metric MeanMetric was called before the ``update`` method which may lead to errors, as metric states have not yet been updated.\n",
      "  warnings.warn(*args, **kwargs)  # noqa: B028\n",
      "Rank-0: policy_step=2000, reward_env_0=29.94190788269043\n",
      "Rank-0: policy_step=2000, reward_env_1=39.607967376708984\n",
      "Rank-0: policy_step=2000, reward_env_2=45.71022033691406\n",
      "Rank-0: policy_step=2000, reward_env_3=42.15308380126953\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/torchmetrics/utilities/prints.py:43: UserWarning: The ``compute`` method of metric MeanMetric was called before the ``update`` method which may lead to errors, as metric states have not yet been updated.\n",
      "  warnings.warn(*args, **kwargs)  # noqa: B028\n",
      "Rank-0: policy_step=4000, reward_env_0=104.86756896972656\n",
      "Rank-0: policy_step=4000, reward_env_1=106.7025375366211\n",
      "Rank-0: policy_step=4000, reward_env_2=58.96342468261719\n",
      "Rank-0: policy_step=4000, reward_env_3=66.83173370361328\n",
      "Rank-0: policy_step=6000, reward_env_0=183.31201171875\n",
      "Rank-0: policy_step=6000, reward_env_1=210.99574279785156\n",
      "Rank-0: policy_step=6000, reward_env_2=158.48216247558594\n",
      "Rank-0: policy_step=6000, reward_env_3=239.71728515625\n",
      "Rank-0: policy_step=8000, reward_env_0=165.63050842285156\n",
      "Rank-0: policy_step=8000, reward_env_1=229.00682067871094\n",
      "Rank-0: policy_step=8000, reward_env_2=281.1435546875\n",
      "Rank-0: policy_step=8000, reward_env_3=156.20188903808594\n",
      "Rank-0: policy_step=10000, reward_env_0=484.4956359863281\n",
      "Rank-0: policy_step=10000, reward_env_1=433.30072021484375\n",
      "Rank-0: policy_step=10000, reward_env_2=355.4902038574219\n",
      "Rank-0: policy_step=10000, reward_env_3=434.70391845703125\n",
      "Test - Reward: 327.08649655482\n"
     ]
    }
   ],
   "source": [
    "!python sheeprl.py exp=dreamer_v3_dmc_walk_walk_grayscale_vector"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. Using vector values as observation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/lightning/fabric/__init__.py:41: Deprecated call to `pkg_resources.declare_namespace('lightning.fabric')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/pkg_resources/__init__.py:2350: DeprecationWarning: Deprecated call to `pkg_resources.declare_namespace('lightning')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "  declare_namespace(parent)\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/lightning/pytorch/__init__.py:37: Deprecated call to `pkg_resources.declare_namespace('lightning.pytorch')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/pkg_resources/__init__.py:2350: DeprecationWarning: Deprecated call to `pkg_resources.declare_namespace('lightning')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "  declare_namespace(parent)\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/lightning_fabric/__init__.py:41: Deprecated call to `pkg_resources.declare_namespace('lightning_fabric')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/pytorch_lightning/__init__.py:37: Deprecated call to `pkg_resources.declare_namespace('pytorch_lightning')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "\u001b[2mCONFIG\u001b[0m\n",
      "\u001b[2m├── \u001b[0m\u001b[2malgo\u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m└── \u001b[0m\u001b[2;91;40mname\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mdreamer_v3                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mtotal_steps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m10000                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mper_rank_batch_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m32                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mrun_test\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcnn_keys\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mencoder\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m[]                                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdecoder\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m[]                                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mmlp_keys\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mencoder\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mstate                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdecoder\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mstate                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mworld_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40moptimizer\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.optim.Adam                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlr\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.0001                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0e-08                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mweight_decay\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mbetas\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.9                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.999                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdiscrete_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m32                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mstochastic_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m32                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkl_dynamic\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.5                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkl_representation\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.1                                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkl_free_nats\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkl_regularizer\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mcontinue_scale_factor\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mclip_gradients\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1000.0                                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdecoupled_rssm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mlearnable_initial_recurrent_state\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mencoder\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcnn_channels_multiplier\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m32                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcnn_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcnn_layer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNormChannelLast                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmlp_layer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mrecurrent_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mrecurrent_state_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mtransition_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mhidden_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mrepresentation_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mhidden_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mobservation_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcnn_channels_multiplier\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m32                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcnn_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcnn_layer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNormChannelLast                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmlp_layer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mreward_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mbins\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m255                                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdiscount_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlearnable\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m        \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mactor\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                  \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40moptimizer\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.optim.Adam                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlr\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m8.0e-05                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0e-05                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mweight_decay\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mbetas\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.9                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.999                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.algos.dreamer_v3.agent.Actor                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40ment_coef\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.0003                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mmin_std\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.1                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mmax_std\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40minit_std\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2.0                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mclip_gradients\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m100.0                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40munimix\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.01                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40maction_clip\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mmoments\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mdecay\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.99                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mmax\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mpercentile\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mlow\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.05                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40mhigh\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.95                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcritic\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40moptimizer\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.optim.Adam                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mlr\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m8.0e-05                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1.0e-05                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mweight_decay\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mbetas\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.9                                                               \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.999                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mlayer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                      \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mper_rank_target_network_update_freq\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mtau\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.02                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mbins\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m255                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mclip_gradients\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m100.0                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mgamma\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.996996996996997                                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mlmbda\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.95                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mhorizon\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m15                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mreplay_ratio\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.5                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mlearning_starts\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1000                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mper_rank_pretrain_steps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mper_rank_sequence_length\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m64                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcnn_layer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNormChannelLast                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mmlp_layer_norm\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mcls\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.models.models.LayerNorm                                  \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkw\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40meps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.001                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mdense_units\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m512                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mmlp_layers\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                           \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mdense_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcnn_act\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorch.nn.SiLU                                                  \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40munimix\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m0.01                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mhafner_initialization\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mplayer\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdiscrete_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m32                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m                                                                        \u001b[0m\n",
      "\u001b[2m├── \u001b[0m\u001b[2mbuffer\u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m└── \u001b[0m\u001b[2;91;40msize\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m10000                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mmemmap\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mvalidate_args\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mfrom_numpy\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcheckpoint\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m                                                                        \u001b[0m\n",
      "\u001b[2m├── \u001b[0m\u001b[2mcheckpoint\u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m└── \u001b[0m\u001b[2;91;40mevery\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1000                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mresume_from\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mnull                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40msave_last\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mkeep_last\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m5                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m                                                                        \u001b[0m\n",
      "\u001b[2m├── \u001b[0m\u001b[2menv\u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m└── \u001b[0m\u001b[2;91;40mid\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mwalker_walk                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mnum_envs\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m4                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mframe_stack\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40msync_env\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mscreen_size\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m64                                                         \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40maction_repeat\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mgrayscale\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                        \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mclip_rewards\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcapture_video\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                     \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mframe_stack_dilation\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1                                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mmax_episode_steps\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m-1                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mreward_as_observation\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mwrapper\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.envs.dmc.DMCWrapper                                 \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdomain_name\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mwalker                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mtask_name\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mwalk                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mwidth\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m64                                                             \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mheight\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m64                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mseed\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mnull                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mfrom_pixels\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mfrom_vectors\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                                    \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m                                                                        \u001b[0m\n",
      "\u001b[2m├── \u001b[0m\u001b[2mfabric\u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m└── \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mlightning.fabric.Fabric                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mdevices\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mnum_nodes\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1                                                            \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mstrategy\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mauto                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40maccelerator\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mcuda                                                       \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mprecision\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mbf16-mixed                                                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mcallbacks\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m-\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.utils.callback.CheckpointCallback                   \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mkeep_last\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m5                                                          \u001b[0m\n",
      "\u001b[2m│   \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m                                                                        \u001b[0m\n",
      "\u001b[2m└── \u001b[0m\u001b[2mmetric\u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m└── \u001b[0m\u001b[2;91;40mlog_every\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m100                                                          \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mdisable_timer\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                    \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mlog_level\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m1                                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                                  \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40maggregator\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                             \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40msheeprl.utils.metric.MetricAggregator                       \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mraise_on_missing\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                               \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mmetrics\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                              \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mRewards/rew_avg\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                    \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mGame/ep_len_avg\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                    \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/world_model_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                              \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/value_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                    \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/policy_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                   \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/observation_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                              \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/reward_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                   \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/state_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                    \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mLoss/continue_loss\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mState/kl\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                           \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mState/post_entropy\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mState/prior_entropy\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mGrads/world_model\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                  \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mGrads/actor\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                        \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m    \u001b[0m\u001b[2;91;40mGrads/critic\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                       \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtorchmetrics.MeanMetric                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m      \u001b[0m\u001b[2;91;40msync_on_compute\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mfalse                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;91;40mlogger\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;40m                                                                 \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40m_target_\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mlightning.fabric.loggers.TensorBoardLogger                  \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mname\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40m2024-05-06_03-21-59_dreamer_v3_walker_walk_5                    \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mroot_dir\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mlogs/runs/dreamer_v3/walker_walk                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mversion\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mnull                                                         \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mdefault_hp_metric\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mtrue                                               \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40mprefix\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;93;40m'\u001b[0m\u001b[2;93;40m'\u001b[0m\u001b[2;40m                                                            \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;97;40m  \u001b[0m\u001b[2;91;40msub_dir\u001b[0m\u001b[2;97;40m:\u001b[0m\u001b[2;97;40m \u001b[0m\u001b[2;40mnull                                                         \u001b[0m\n",
      "\u001b[2m    \u001b[0m\u001b[2m    \u001b[0m\u001b[2;40m                                                                        \u001b[0m\n",
      "[2024-05-06 03:22:00,101][pytorch_lightning.utilities.rank_zero][INFO] - Using bfloat16 Automatic Mixed Precision (AMP)\n",
      "INFO: Seed set to 5\n",
      "[2024-05-06 03:22:00,174][lightning.fabric.utilities.seed][INFO] - Seed set to 5\n",
      "Log dir: logs/runs/dreamer_v3/walker_walk/2024-05-06_03-21-59_dreamer_v3_walker_walk_5/version_0\n",
      "Encoder CNN keys: []\n",
      "Encoder MLP keys: ['state']\n",
      "Decoder CNN keys: []\n",
      "Decoder MLP keys: ['state']\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/torchmetrics/utilities/prints.py:43: UserWarning: The ``compute`` method of metric MeanMetric was called before the ``update`` method which may lead to errors, as metric states have not yet been updated.\n",
      "  warnings.warn(*args, **kwargs)  # noqa: B028\n",
      "Rank-0: policy_step=2000, reward_env_0=32.44091033935547\n",
      "Rank-0: policy_step=2000, reward_env_1=41.8197021484375\n",
      "Rank-0: policy_step=2000, reward_env_2=36.76511001586914\n",
      "Rank-0: policy_step=2000, reward_env_3=43.39164352416992\n",
      "/home/jianheng/miniconda3/envs/sheeprl/lib/python3.10/site-packages/torchmetrics/utilities/prints.py:43: UserWarning: The ``compute`` method of metric MeanMetric was called before the ``update`` method which may lead to errors, as metric states have not yet been updated.\n",
      "  warnings.warn(*args, **kwargs)  # noqa: B028\n",
      "Rank-0: policy_step=4000, reward_env_0=49.986602783203125\n",
      "Rank-0: policy_step=4000, reward_env_1=69.43289947509766\n",
      "Rank-0: policy_step=4000, reward_env_2=114.92591857910156\n",
      "Rank-0: policy_step=4000, reward_env_3=114.09442138671875\n",
      "Rank-0: policy_step=6000, reward_env_0=129.74349975585938\n",
      "Rank-0: policy_step=6000, reward_env_1=149.32179260253906\n",
      "Rank-0: policy_step=6000, reward_env_2=126.85379791259766\n",
      "Rank-0: policy_step=6000, reward_env_3=111.9491958618164\n",
      "Rank-0: policy_step=8000, reward_env_0=164.39967346191406\n",
      "Rank-0: policy_step=8000, reward_env_1=174.74343872070312\n",
      "Rank-0: policy_step=8000, reward_env_2=185.4028778076172\n",
      "Rank-0: policy_step=8000, reward_env_3=160.4306640625\n",
      "Rank-0: policy_step=10000, reward_env_0=278.9344177246094\n",
      "Rank-0: policy_step=10000, reward_env_1=238.4259490966797\n",
      "Rank-0: policy_step=10000, reward_env_2=254.5703887939453\n",
      "Rank-0: policy_step=10000, reward_env_3=241.44131469726562\n",
      "Test - Reward: 282.6285009306698\n"
     ]
    }
   ],
   "source": [
    "!python sheeprl.py exp=dreamer_v3_dmc_walk_walk_vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorFlow installation not found - running with reduced feature set.\n",
      "\n",
      "NOTE: Using experimental fast data loading logic. To disable, pass\n",
      "    \"--load_fast=false\" and report issues on GitHub. More details:\n",
      "    https://github.com/tensorflow/tensorboard/issues/4784\n",
      "\n",
      "Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all\n",
      "TensorBoard 2.16.2 at http://localhost:6006/ (Press CTRL+C to quit)\n",
      "^C\n"
     ]
    }
   ],
   "source": [
    "!tensorboard --logdir logs"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sheeprl",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}