# Buggy code
def __init__(self, obs_space, action_space, config):
    _validate(obs_space, action_space)
    config = dict(ray.rllib.agents.qmix.qmix.DEFAULT_CONFIG, **config)
    self.framework = "torch"
    super().__init__(obs_space, action_space, config)
    self.n_agents = len(obs_space.original_space.spaces)
    self.n_actions = action_space.spaces[0].n
    self.h_size = config["model"]["lstm_cell_size"]
    self.has_env_global_state = False
    self.has_action_mask = False
    self.device = (
        torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    )

    agent_obs_space = obs_space.original_space.spaces[0]
    if isinstance(agent_obs_space, Dict):
        space_keys = set(agent_obs_space.spaces.keys())
        if "obs" not in space_keys:
            raise ValueError("Dict obs space must have subspace labeled `obs`")
        self.obs_size = _get_size(agent_obs_space.spaces["obs"])
        if "action_mask" in space_keys:
            mask_shape = tuple(agent_obs_space.spaces["action_mask"].shape)
            if mask_shape != (self.n_actions,):
                raise ValueError(
                    "Action mask shape must be {}, got {}".format(
                        (self.n_actions,), mask_shape
                    )
                )
            self.has_action_mask = True
        if ENV_STATE in space_keys:
            self.env_global_state_shape = _get_size(agent_obs_space.spaces[ENV_STATE])
            self.has_env_global_state = True
        else:
            self.env_global_state_shape = (self.obs_size, self.n_agents)
        # The real agent obs space is nested inside the dict
        config["model"]["full_obs_space"] = agent_obs_space
        agent_obs_space = agent_obs_space.spaces["obs"]
    else:
        self.obs_size = _get_size(agent_obs_space)

    self.model = ModelCatalog.get_model_v2(
        agent_obs_space,
        action_space.spaces[0],
        self.n_actions,
        config["model"],
        framework="torch",
        name="model",
        default_model=RNNModel,
    ).to(self.device)

    self.target_model = ModelCatalog.get_model_v2(
        agent_obs_space,
        action_space.spaces[0],
        self.n_actions,
        config["model"],
        framework="torch",
        name="target_model",
        default_model=RNNModel,
    ).to(self.device)

    self.exploration = self._create_exploration()

    # Setup the mixer network.
    if config["mixer"] is None:
        self.mixer = None
        self.target_mixer = None
    elif config["mixer"] == "qmix":
        self.mixer = QMixer(
            self.n_agents, self.env_global_state_shape, config["mixing_embed_dim"]
        ).to(self.device)
        self.target_mixer = QMixer(
            self.n_agents, self.env_global_state_shape, config["mixing_embed_dim"]
        ).to(self.device)
    elif config["mixer"] == "vdn":
        self.mixer = VDNMixer().to(self.device)
        self.target_mixer = VDNMixer().to(self.device)
    else:
        raise ValueError("Unknown mixer type {}".format(config["mixer"]))

    self.cur_epsilon = 1.0
    self.update_target()  # initial sync

    # Setup optimizer
    self.params = list(self.model.parameters())
    if self.mixer:
        self.params += list(self.mixer.parameters())
    self.loss = QMixLoss(
        self.model,
        self.target_model,
        self.mixer,
        self.target_mixer,
        self.n_agents,
        self.n_actions,
        self.config["double_q"],
        self.config["gamma"],
    )
    from torch.optim import RMSprop

    self.optimiser = RMSprop(
        params=self.params,
        lr=config["lr"],
        alpha=config["optim_alpha"],
        eps=config["optim_eps"],
    )


# Traceback (AttributeError)
Failure # 1 (occurred at 2020-06-01_14-46-39)
Traceback (most recent call last):
File "/home/ravi/dev/gym-rl/qmix/lib/python3.6/site-packages/ray/tune/trial_runner.py", line 467, in _process_trial
result = self.trial_executor.fetch_result(trial)
File "/home/ravi/dev/gym-rl/qmix/lib/python3.6/site-packages/ray/tune/ray_trial_executor.py", line 430, in fetch_result
result = ray.get(trial_future[0], DEFAULT_GET_TIMEOUT)
File "/home/ravi/dev/gym-rl/qmix/lib/python3.6/site-packages/ray/worker.py", line 1522, in get
raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(AttributeError): �[36mray::QMIX.train()�[39m (pid=3884, ip=10.0.2.15)
File "python/ray/_raylet.pyx", line 421, in ray._raylet.execute_task
File "python/ray/_raylet.pyx", line 456, in ray._raylet.execute_task
File "python/ray/_raylet.pyx", line 459, in ray._raylet.execute_task
File "python/ray/_raylet.pyx", line 460, in ray._raylet.execute_task
File "python/ray/_raylet.pyx", line 414, in ray._raylet.execute_task.function_executor
File "/home/ravi/dev/gym-rl/qmix/lib/python3.6/site-packages/ray/rllib/agents/trainer_template.py", line 90, in __init__
Trainer.__init__(self, config, env, logger_creator)
File "/home/ravi/dev/gym-rl/qmix/lib/python3.6/site-packages/ray/rllib/agents/trainer.py", line 444, in __init__
super().__init__(config, logger_creator)
File "/home/ravi/dev/gym-rl/qmix/lib/python3.6/site-packages/ray/tune/trainable.py", line 174, in __init__
self._setup(copy.deepcopy(self.config))
File "/home/ravi/dev/gym-rl/qmix/lib/python3.6/site-packages/ray/rllib/agents/trainer.py", line 613, in _setup
self._init(self.config, self.env_creator)
File "/home/ravi/dev/gym-rl/qmix/lib/python3.6/site-packages/ray/rllib/agents/trainer_template.py", line 115, in _init
self.config["num_workers"])
File "/home/ravi/dev/gym-rl/qmix/lib/python3.6/site-packages/ray/rllib/agents/trainer.py", line 684, in _make_workers
logdir=self.logdir)
File "/home/ravi/dev/gym-rl/qmix/lib/python3.6/site-packages/ray/rllib/evaluation/worker_set.py", line 59, in __init__
RolloutWorker, env_creator, policy, 0, self._local_config)
File "/home/ravi/dev/gym-rl/qmix/lib/python3.6/site-packages/ray/rllib/evaluation/worker_set.py", line 282, in _make_worker
extra_python_environs=extra_python_environs)
File "/home/ravi/dev/gym-rl/qmix/lib/python3.6/site-packages/ray/rllib/evaluation/rollout_worker.py", line 393, in __init__
policy_dict, policy_config)
File "/home/ravi/dev/gym-rl/qmix/lib/python3.6/site-packages/ray/rllib/evaluation/rollout_worker.py", line 932, in _build_policy_map
policy_map[name] = cls(obs_space, act_space, merged_conf)
File "/home/ravi/dev/gym-rl/qmix/lib/python3.6/site-packages/ray/rllib/agents/qmix/qmix_policy.py", line 223, in __init__
self.mixer = QMixer(self.n_agents, self.env_global_state_shape,
AttributeError: 'QMixTorchPolicy' object has no attribute 'env_global_state_shape'

# Fixed code
def __init__(self, obs_space, action_space, config):
    _validate(obs_space, action_space)
    config = dict(ray.rllib.agents.qmix.qmix.DEFAULT_CONFIG, **config)
    self.framework = "torch"
    super().__init__(obs_space, action_space, config)
    self.n_agents = len(obs_space.original_space.spaces)
    self.n_actions = action_space.spaces[0].n
    self.h_size = config["model"]["lstm_cell_size"]
    self.has_env_global_state = False
    self.has_action_mask = False
    self.device = (
        torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    )

    agent_obs_space = obs_space.original_space.spaces[0]
    if isinstance(agent_obs_space, Dict):
        space_keys = set(agent_obs_space.spaces.keys())
        if "obs" not in space_keys:
            raise ValueError("Dict obs space must have subspace labeled `obs`")
        self.obs_size = _get_size(agent_obs_space.spaces["obs"])
        if "action_mask" in space_keys:
            mask_shape = tuple(agent_obs_space.spaces["action_mask"].shape)
            if mask_shape != (self.n_actions,):
                raise ValueError(
                    "Action mask shape must be {}, got {}".format(
                        (self.n_actions,), mask_shape
                    )
                )
            self.has_action_mask = True
        if ENV_STATE in space_keys:
            self.env_global_state_shape = _get_size(agent_obs_space.spaces[ENV_STATE])
            self.has_env_global_state = True
        else:
            self.env_global_state_shape = (self.obs_size, self.n_agents)
        # The real agent obs space is nested inside the dict
        config["model"]["full_obs_space"] = agent_obs_space
        agent_obs_space = agent_obs_space.spaces["obs"]
    else:
        self.obs_size = _get_size(agent_obs_space)
        self.env_global_state_shape = (self.obs_size, self.n_agents)

    self.model = ModelCatalog.get_model_v2(
        agent_obs_space,
        action_space.spaces[0],
        self.n_actions,
        config["model"],
        framework="torch",
        name="model",
        default_model=RNNModel,
    ).to(self.device)

    self.target_model = ModelCatalog.get_model_v2(
        agent_obs_space,
        action_space.spaces[0],
        self.n_actions,
        config["model"],
        framework="torch",
        name="target_model",
        default_model=RNNModel,
    ).to(self.device)

    self.exploration = self._create_exploration()

    # Setup the mixer network.
    if config["mixer"] is None:
        self.mixer = None
        self.target_mixer = None
    elif config["mixer"] == "qmix":
        self.mixer = QMixer(
            self.n_agents, self.env_global_state_shape, config["mixing_embed_dim"]
        ).to(self.device)
        self.target_mixer = QMixer(
            self.n_agents, self.env_global_state_shape, config["mixing_embed_dim"]
        ).to(self.device)
    elif config["mixer"] == "vdn":
        self.mixer = VDNMixer().to(self.device)
        self.target_mixer = VDNMixer().to(self.device)
    else:
        raise ValueError("Unknown mixer type {}".format(config["mixer"]))

    self.cur_epsilon = 1.0
    self.update_target()  # initial sync

    # Setup optimizer
    self.params = list(self.model.parameters())
    if self.mixer:
        self.params += list(self.mixer.parameters())
    self.loss = QMixLoss(
        self.model,
        self.target_model,
        self.mixer,
        self.target_mixer,
        self.n_agents,
        self.n_actions,
        self.config["double_q"],
        self.config["gamma"],
    )
    from torch.optim import RMSprop

    self.optimiser = RMSprop(
        params=self.params,
        lr=config["lr"],
        alpha=config["optim_alpha"],
        eps=config["optim_eps"],
    )