# Buggy code
def __init__(
    self,
    obs_space,
    action_space,
    num_outputs,
    model_config,
    name,
    q_hiddens=(256,),
    dueling=False,
    num_atoms=1,
    use_noisy=False,
    v_min=-10.0,
    v_max=10.0,
    sigma0=0.5,
    parameter_noise=False,
):
    """Initialize variables of this model.

    Extra model kwargs:
        q_hiddens (list): defines size of hidden layers for the q head.
            These will be used to postprocess the model output for the
            purposes of computing Q values.
        dueling (bool): whether to build the state value head for DDQN
        num_atoms (int): if >1, enables distributional DQN
        use_noisy (bool): use noisy nets
        v_min (float): min value support for distributional DQN
        v_max (float): max value support for distributional DQN
        sigma0 (float): initial value of noisy nets
        parameter_noise (bool): enable layer norm for param noise

    Note that the core layers for forward() are not defined here, this
    only defines the layers for the Q head. Those layers for forward()
    should be defined in subclasses of DistributionalQModel.
    """

    super(DistributionalQModel, self).__init__(
        obs_space, action_space, num_outputs, model_config, name
    )

    # setup the Q head output (i.e., model for get_q_values)
    self.model_out = tf.keras.layers.Input(shape=(num_outputs,), name="model_out")

    def build_action_value(model_out):
        if q_hiddens:
            action_out = model_out
            for i in range(len(q_hiddens)):
                if use_noisy:
                    action_out = self._noisy_layer(
                        "hidden_%d" % i, action_out, q_hiddens[i], sigma0
                    )
                elif parameter_noise:
                    action_out = tf.keras.layers.Dense(
                        units=q_hiddens[i],
                        activation_fn=tf.nn.relu,
                        normalizer_fn=tf.keras.layers.LayerNormalization,
                    )(action_out)
                else:
                    action_out = tf.keras.layers.Dense(
                        units=q_hiddens[i], activation=tf.nn.relu, name="hidden_%d" % i
                    )(action_out)
        else:
            # Avoid postprocessing the outputs. This enables custom models
            # to be used for parametric action DQN.
            action_out = model_out
        if use_noisy:
            action_scores = self._noisy_layer(
                "output",
                action_out,
                self.action_space.n * num_atoms,
                sigma0,
                non_linear=False,
            )
        elif q_hiddens:
            action_scores = tf.keras.layers.Dense(
                units=self.action_space.n * num_atoms, activation=None
            )(action_out)
        else:
            action_scores = model_out
        if num_atoms > 1:
            # Distributional Q-learning uses a discrete support z
            # to represent the action value distribution
            z = tf.range(num_atoms, dtype=tf.float32)
            z = v_min + z * (v_max - v_min) / float(num_atoms - 1)
            support_logits_per_action = tf.reshape(
                tensor=action_scores, shape=(-1, self.action_space.n, num_atoms)
            )
            support_prob_per_action = tf.nn.softmax(logits=support_logits_per_action)
            action_scores = tf.reduce_sum(
                input_tensor=z * support_prob_per_action, axis=-1
            )
            logits = support_logits_per_action
            dist = support_prob_per_action
            return [action_scores, z, support_logits_per_action, logits, dist]
        else:
            logits = tf.expand_dims(tf.ones_like(action_scores), -1)
            dist = tf.expand_dims(tf.ones_like(action_scores), -1)
            return [action_scores, logits, dist]

    def build_state_score(model_out):
        state_out = model_out
        for i in range(len(q_hiddens)):
            if use_noisy:
                state_out = self._noisy_layer(
                    "dueling_hidden_%d" % i, state_out, q_hiddens[i], sigma0
                )
            elif parameter_noise:
                state_out = tf.keras.layers.Dense(
                    units=q_hiddens[i],
                    activation_fn=tf.nn.relu,
                    normalizer_fn=tf.contrib.layers.layer_norm,
                )(state_out)
            else:
                state_out = tf.keras.layers.Dense(
                    units=q_hiddens[i], activation=tf.nn.relu
                )(state_out)
        if use_noisy:
            state_score = self._noisy_layer(
                "dueling_output", state_out, num_atoms, sigma0, non_linear=False
            )
        else:
            state_score = tf.keras.layers.Dense(units=num_atoms, activation=None)(
                state_out
            )
        return state_score

    if tf.executing_eagerly():
        from tensorflow.python.ops import variable_scope

        # Have to use a variable store to reuse variables in eager mode
        store = variable_scope.EagerVariableStore()

        # Save the scope objects, since in eager we will execute this
        # path repeatedly and there is no guarantee it will always be run
        # in the same original scope.
        with tf.variable_scope(name + "/action_value") as action_scope:
            pass
        with tf.variable_scope(name + "/state_value") as state_scope:
            pass

        def build_action_value_in_scope(model_out):
            with store.as_default():
                with tf.variable_scope(action_scope, reuse=tf.AUTO_REUSE):
                    return build_action_value(model_out)

        def build_state_score_in_scope(model_out):
            with store.as_default():
                with tf.variable_scope(state_scope, reuse=tf.AUTO_REUSE):
                    return build_state_score(model_out)
    else:

        def build_action_value_in_scope(model_out):
            with tf.variable_scope(name + "/action_value", reuse=tf.AUTO_REUSE):
                return build_action_value(model_out)

        def build_state_score_in_scope(model_out):
            with tf.variable_scope(name + "/state_value", reuse=tf.AUTO_REUSE):
                return build_state_score(model_out)

    q_out = build_action_value_in_scope(self.model_out)
    self.q_value_head = tf.keras.Model(self.model_out, q_out)
    self.register_variables(self.q_value_head.variables)

    if dueling:
        state_out = build_state_score_in_scope(self.model_out)
        self.state_value_head = tf.keras.Model(self.model_out, state_out)
        self.register_variables(self.state_value_head.variables)


# Traceback (AttributeError)
Traceback (most recent call last):
File "/usr/local/lib/python3.7/site-packages/ray/tune/trial_runner.py", line 459, in _process_trial
result = self.trial_executor.fetch_result(trial)
File "/usr/local/lib/python3.7/site-packages/ray/tune/ray_trial_executor.py", line 377, in fetch_result
result = ray.get(trial_future[0], DEFAULT_GET_TIMEOUT)
File "/usr/local/lib/python3.7/site-packages/ray/worker.py", line 1504, in get
raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(AttributeError): ray::DQN.__init__() (pid=96308, ip=192.168.120.74)
File "python/ray/_raylet.pyx", line 437, in ray._raylet.execute_task
File "python/ray/_raylet.pyx", line 449, in ray._raylet.execute_task
File "python/ray/_raylet.pyx", line 450, in ray._raylet.execute_task
File "python/ray/_raylet.pyx", line 452, in ray._raylet.execute_task
File "python/ray/_raylet.pyx", line 430, in ray._raylet.execute_task.function_executor
File "/usr/local/lib/python3.7/site-packages/ray/rllib/agents/trainer_template.py", line 86, in __init__
Trainer.__init__(self, config, env, logger_creator)
File "/usr/local/lib/python3.7/site-packages/ray/rllib/agents/trainer.py", line 447, in __init__
super().__init__(config, logger_creator)
File "/usr/local/lib/python3.7/site-packages/ray/tune/trainable.py", line 172, in __init__
self._setup(copy.deepcopy(self.config))
File "/usr/local/lib/python3.7/site-packages/ray/rllib/agents/trainer.py", line 591, in _setup
self._init(self.config, self.env_creator)
File "/usr/local/lib/python3.7/site-packages/ray/rllib/agents/trainer_template.py", line 105, in _init
self.config["num_workers"])
File "/usr/local/lib/python3.7/site-packages/ray/rllib/agents/trainer.py", line 658, in _make_workers
logdir=self.logdir)
File "/usr/local/lib/python3.7/site-packages/ray/rllib/evaluation/worker_set.py", line 60, in __init__
RolloutWorker, env_creator, policy, 0, self._local_config)
File "/usr/local/lib/python3.7/site-packages/ray/rllib/evaluation/worker_set.py", line 262, in _make_worker
_fake_sampler=config.get("_fake_sampler", False))
File "/usr/local/lib/python3.7/site-packages/ray/rllib/evaluation/rollout_worker.py", line 355, in __init__
self._build_policy_map(policy_dict, policy_config)
File "/usr/local/lib/python3.7/site-packages/ray/rllib/evaluation/rollout_worker.py", line 820, in _build_policy_map
policy_map[name] = cls(obs_space, act_space, merged_conf)
File "/usr/local/lib/python3.7/site-packages/ray/rllib/policy/tf_policy_template.py", line 138, in __init__
obs_include_prev_action_reward=obs_include_prev_action_reward)
File "/usr/local/lib/python3.7/site-packages/ray/rllib/policy/dynamic_tf_policy.py", line 137, in __init__
self.model = make_model(self, obs_space, action_space, config)
File "/usr/local/lib/python3.7/site-packages/ray/rllib/agents/dqn/dqn_policy.py", line 183, in build_q_model
parameter_noise=config["parameter_noise"])
File "/usr/local/lib/python3.7/site-packages/ray/rllib/models/catalog.py", line 349, in get_model_v2
name, **model_kwargs)
File "/usr/local/lib/python3.7/site-packages/ray/rllib/agents/dqn/distributional_q_model.py", line 185, in __init__
q_out = build_action_value_in_scope(self.model_out)
File "/usr/local/lib/python3.7/site-packages/ray/rllib/agents/dqn/distributional_q_model.py", line 178, in build_action_value_in_scope
return build_action_value(model_out)
File "/usr/local/lib/python3.7/site-packages/ray/rllib/agents/dqn/distributional_q_model.py", line 68, in build_action_value
"hidden_%d" % i, action_out, q_hiddens[i], sigma0)
File "/usr/local/lib/python3.7/site-packages/ray/rllib/agents/dqn/distributional_q_model.py", line 259, in _noisy_layer
initializer=tf.initializers.GlorotUniform())
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/util/module_wrapper.py", line 193, in __getattr__
attr = getattr(self._tfmw_wrapped_module, name)
AttributeError: module 'tensorflow._api.v1.compat.v1.initializers' has no attribute 'GlorotUniform'

# Fixed code
def __init__(
    self,
    obs_space,
    action_space,
    num_outputs,
    model_config,
    name,
    q_hiddens=(256,),
    dueling=False,
    num_atoms=1,
    use_noisy=False,
    v_min=-10.0,
    v_max=10.0,
    sigma0=0.5,
    parameter_noise=False,
):
    """Initialize variables of this model.

    Extra model kwargs:
        q_hiddens (list): defines size of hidden layers for the q head.
            These will be used to postprocess the model output for the
            purposes of computing Q values.
        dueling (bool): whether to build the state value head for DDQN
        num_atoms (int): if >1, enables distributional DQN
        use_noisy (bool): use noisy nets
        v_min (float): min value support for distributional DQN
        v_max (float): max value support for distributional DQN
        sigma0 (float): initial value of noisy nets
        parameter_noise (bool): enable layer norm for param noise

    Note that the core layers for forward() are not defined here, this
    only defines the layers for the Q head. Those layers for forward()
    should be defined in subclasses of DistributionalQModel.
    """

    super(DistributionalQModel, self).__init__(
        obs_space, action_space, num_outputs, model_config, name
    )

    # setup the Q head output (i.e., model for get_q_values)
    self.model_out = tf.keras.layers.Input(shape=(num_outputs,), name="model_out")

    def build_action_value(model_out):
        if q_hiddens:
            action_out = model_out
            for i in range(len(q_hiddens)):
                if use_noisy:
                    action_out = self._noisy_layer(
                        "hidden_%d" % i, action_out, q_hiddens[i], sigma0
                    )
                elif parameter_noise:
                    action_out = tf.keras.layers.Dense(
                        units=q_hiddens[i],
                        activation_fn=tf.nn.relu,
                        normalizer_fn=tf.keras.layers.LayerNormalization,
                    )(action_out)
                else:
                    action_out = tf.keras.layers.Dense(
                        units=q_hiddens[i], activation=tf.nn.relu, name="hidden_%d" % i
                    )(action_out)
        else:
            # Avoid postprocessing the outputs. This enables custom models
            # to be used for parametric action DQN.
            action_out = model_out
        if use_noisy:
            action_scores = self._noisy_layer(
                "output",
                action_out,
                self.action_space.n * num_atoms,
                sigma0,
                non_linear=False,
            )
        elif q_hiddens:
            action_scores = tf.keras.layers.Dense(
                units=self.action_space.n * num_atoms, activation=None
            )(action_out)
        else:
            action_scores = model_out

        if num_atoms > 1:
            # Distributional Q-learning uses a discrete support z
            # to represent the action value distribution
            z = tf.range(num_atoms, dtype=tf.float32)
            z = v_min + z * (v_max - v_min) / float(num_atoms - 1)

            def _layer(x):
                support_logits_per_action = tf.reshape(
                    tensor=x, shape=(-1, self.action_space.n, num_atoms)
                )
                support_prob_per_action = tf.nn.softmax(
                    logits=support_logits_per_action
                )
                x = tf.reduce_sum(input_tensor=z * support_prob_per_action, axis=-1)
                logits = support_logits_per_action
                dist = support_prob_per_action
                return [x, z, support_logits_per_action, logits, dist]

            return tf.keras.layers.Lambda(_layer)(action_scores)
        else:
            logits = tf.expand_dims(tf.ones_like(action_scores), -1)
            dist = tf.expand_dims(tf.ones_like(action_scores), -1)
            return [action_scores, logits, dist]

    def build_state_score(model_out):
        state_out = model_out
        for i in range(len(q_hiddens)):
            if use_noisy:
                state_out = self._noisy_layer(
                    "dueling_hidden_%d" % i, state_out, q_hiddens[i], sigma0
                )
            elif parameter_noise:
                state_out = tf.keras.layers.Dense(
                    units=q_hiddens[i],
                    activation_fn=tf.nn.relu,
                    normalizer_fn=tf.contrib.layers.layer_norm,
                )(state_out)
            else:
                state_out = tf.keras.layers.Dense(
                    units=q_hiddens[i], activation=tf.nn.relu
                )(state_out)
        if use_noisy:
            state_score = self._noisy_layer(
                "dueling_output", state_out, num_atoms, sigma0, non_linear=False
            )
        else:
            state_score = tf.keras.layers.Dense(units=num_atoms, activation=None)(
                state_out
            )
        return state_score

    if tf.executing_eagerly():
        from tensorflow.python.ops import variable_scope

        # Have to use a variable store to reuse variables in eager mode
        store = variable_scope.EagerVariableStore()

        # Save the scope objects, since in eager we will execute this
        # path repeatedly and there is no guarantee it will always be run
        # in the same original scope.
        with tf.variable_scope(name + "/action_value") as action_scope:
            pass
        with tf.variable_scope(name + "/state_value") as state_scope:
            pass

        def build_action_value_in_scope(model_out):
            with store.as_default():
                with tf.variable_scope(action_scope, reuse=tf.AUTO_REUSE):
                    return build_action_value(model_out)

        def build_state_score_in_scope(model_out):
            with store.as_default():
                with tf.variable_scope(state_scope, reuse=tf.AUTO_REUSE):
                    return build_state_score(model_out)
    else:

        def build_action_value_in_scope(model_out):
            with tf.variable_scope(name + "/action_value", reuse=tf.AUTO_REUSE):
                return build_action_value(model_out)

        def build_state_score_in_scope(model_out):
            with tf.variable_scope(name + "/state_value", reuse=tf.AUTO_REUSE):
                return build_state_score(model_out)

    q_out = build_action_value_in_scope(self.model_out)
    self.q_value_head = tf.keras.Model(self.model_out, q_out)
    self.register_variables(self.q_value_head.variables)

    if dueling:
        state_out = build_state_score_in_scope(self.model_out)
        self.state_value_head = tf.keras.Model(self.model_out, state_out)
        self.register_variables(self.state_value_head.variables)