From 3c5213b42203a00b74c25c0ca3e12b3ad21bae7e Mon Sep 17 00:00:00 2001 From: Johan Obando Ceron Date: Fri, 15 Sep 2023 20:41:14 +0000 Subject: [PATCH] Adding logging in ImpalaDQNNetwork(). PiperOrigin-RevId: 565767089 --- dopamine/jax/networks.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/dopamine/jax/networks.py b/dopamine/jax/networks.py index d6fb30b0..b665f54b 100644 --- a/dopamine/jax/networks.py +++ b/dopamine/jax/networks.py @@ -17,7 +17,7 @@ import itertools import time from typing import Optional, Sequence, Tuple, Union - +from absl import logging from dopamine.discrete_domains import atari_lib from flax import linen as nn import gin @@ -92,6 +92,11 @@ class ImpalaEncoder(nn.Module): num_blocks: int = 2 def setup(self): + logging.info('\t Creating %s ...', self.__class__.__name__) + logging.info('\t Creating ImpalaDQNNetwork ...') + logging.info('\t num_blocks: %s', self.num_blocks) + logging.info('\t nn_scale: %s', self.nn_scale) + logging.info('\t stack_sizes: %s', self.stack_sizes) self._stacks = [ Stack(num_ch=stack_size * self.nn_scale, num_blocks=self.num_blocks) for stack_size in self.stack_sizes