import tensorflow as tf

                ep_reward += sum(reward_n)
                obs_glob = [obs[0] + obs[1]]
                obs_glob_next = [obs_n[0] + obs_n[1]]
                self.store(obs_glob + obs + action + [sum(reward_n)] + obs_glob_next + obs_n + [all(done_n)])
                obs = obs_n
                self.learn()
            self.write_summary_scalar("ep_reward", ep_reward, self.learn_step_cnt)

    def write_summary_scalar(self, tag, value, iteration):
        self.summary_writer.add_summary(tf.Summary(value=[tf.Summary.Value(tag=tag, simple_value=value)]), iteration)


env = gym.make('Switch2-v0')
alg = QMix(env, env.observation_space[0].shape[0], env.action_space[0].n)
# alg = QMix(env, env.observation_space.shape[0], env.action_space.n)

alg.train()

