# Buggy code
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]


# 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 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]