src/runners/episode_runner.py
def run(self, test_mode=False):
        '''
        Items that have been added by me:
        * current_episode_agent_returns
        '''
        self.reset()
        # this clears the agent returns list every episode. 
        # TODO: Works
        current_episode_agent_returns = [0]*self.args.n_agents
        terminated = False
        episode_return = 0
        self.mac.init_hidden(batch_size=self.batch_size)

        while not terminated:

            pre_transition_data = {
                "state": [self.env.get_state()],
                "avail_actions": [self.env.get_avail_actions()],
                "obs": [self.env.get_obs()]
            }

            self.batch.update(pre_transition_data, ts=self.t)

            # Pass the entire batch of experiences up till now to the agents
            # Receive the actions for each agent at this timestep in a batch of size 1
            actions = self.mac.select_actions(self.batch, t_ep=self.t, t_env=self.t_env, test_mode=test_mode)

            reward, terminated, env_info = self.env.step(actions[0])
            episode_return += reward
            # TODO: works 
            for i, value in enumerate(env_info['agent_rewards']): 
                    current_episode_agent_returns[i] += value
            post_transition_data = {
                "actions": actions,
                "reward": [(reward,)],
                "terminated": [(terminated != env_info.get("episode_limit", False),)],
            }

            self.batch.update(post_transition_data, ts=self.t)

            self.t += 1

        last_data = {
            "state": [self.env.get_state()],
            "avail_actions": [self.env.get_avail_actions()],
            "obs": [self.env.get_obs()]
        }
        self.batch.update(last_data, ts=self.t)

        # Select actions in the last stored state
        actions = self.mac.select_actions(self.batch, t_ep=self.t, t_env=self.t_env, test_mode=test_mode)
        self.batch.update({"actions": actions}, ts=self.t)

        cur_stats = self.test_stats if test_mode else self.train_stats
        cur_returns = self.test_returns if test_mode else self.train_returns
        log_prefix = "test_" if test_mode else ""

        # this is where I can append current_episode_agent_returns
        # TODO: works
        self.agent_returns.append(current_episode_agent_returns)
        env_info.pop('agent_rewards')


        cur_stats.update({k: cur_stats.get(k, 0) + env_info.get(k, 0) for k in set(cur_stats) | set(env_info)})
        cur_stats["n_episodes"] = 1 + cur_stats.get("n_episodes", 0)
        cur_stats["ep_length"] = self.t + cur_stats.get("ep_length", 0)

        if not test_mode:
            self.t_env += self.t

        cur_returns.append(episode_return)
       

        if test_mode and (len(self.test_returns) == self.args.test_nepisode):
            self._log(cur_returns, cur_stats, log_prefix)
            self._log_info(log_prefix)
        elif self.t_env - self.log_train_stats_t >= self.args.runner_log_interval:
            self._log(cur_returns, cur_stats, log_prefix)
            # time to log 
            self._log_info(log_prefix)

            if hasattr(self.mac.action_selector, "epsilon"):
                self.logger.log_stat("epsilon", self.mac.action_selector.epsilon, self.t_env)
            self.log_train_stats_t = self.t_env

        return self.batch

        Items that have been added by me:
        * current_episode_agent_returns
        
Items that have been added by me:
current_episode_agent_returns
Items that have been added by me : current_episode_agent_returns
def run(self, test_mode=False):
        self.reset()
        current_episode_agent_returns = [0]*self.args.n_agents
        terminated = False
        episode_return = 0
        self.mac.init_hidden(batch_size=self.batch_size)
        while not terminated:
            pre_transition_data = {
                "state": [self.env.get_state()],
                "avail_actions": [self.env.get_avail_actions()],
                "obs": [self.env.get_obs()]
            }
            self.batch.update(pre_transition_data, ts=self.t)
            actions = self.mac.select_actions(self.batch, t_ep=self.t, t_env=self.t_env, test_mode=test_mode)
            reward, terminated, env_info = self.env.step(actions[0])
            episode_return += reward
            for i, value in enumerate(env_info['agent_rewards']): 
                    current_episode_agent_returns[i] += value
            post_transition_data = {
                "actions": actions,
                "reward": [(reward,)],
                "terminated": [(terminated != env_info.get("episode_limit", False),)],
            }
            self.batch.update(post_transition_data, ts=self.t)
            self.t += 1
        last_data = {
            "state": [self.env.get_state()],
            "avail_actions": [self.env.get_avail_actions()],
            "obs": [self.env.get_obs()]
        }
        self.batch.update(last_data, ts=self.t)
        actions = self.mac.select_actions(self.batch, t_ep=self.t, t_env=self.t_env, test_mode=test_mode)
        self.batch.update({"actions": actions}, ts=self.t)
        cur_stats = self.test_stats if test_mode else self.train_stats
        cur_returns = self.test_returns if test_mode else self.train_returns
        log_prefix = "test_" if test_mode else ""
        self.agent_returns.append(current_episode_agent_returns)
        env_info.pop('agent_rewards')
        cur_stats.update({k: cur_stats.get(k, 0) + env_info.get(k, 0) for k in set(cur_stats) | set(env_info)})
        cur_stats["n_episodes"] = 1 + cur_stats.get("n_episodes", 0)
        cur_stats["ep_length"] = self.t + cur_stats.get("ep_length", 0)
        if not test_mode:
            self.t_env += self.t
        cur_returns.append(episode_return)
        if test_mode and (len(self.test_returns) == self.args.test_nepisode):
            self._log(cur_returns, cur_stats, log_prefix)
            self._log_info(log_prefix)
        elif self.t_env - self.log_train_stats_t >= self.args.runner_log_interval:
            self._log(cur_returns, cur_stats, log_prefix)
            self._log_info(log_prefix)
            if hasattr(self.mac.action_selector, "epsilon"):
                self.logger.log_stat("epsilon", self.mac.action_selector.epsilon, self.t_env)
            self.log_train_stats_t = self.t_env
        return self.batch
def run ( self , test_mode = False ) : self . reset ( ) current_episode_agent_returns = [ 0 ] * self . args . n_agents terminated = False episode_return = 0 self . mac . init_hidden ( batch_size = self . batch_size ) while not terminated : pre_transition_data = { "state" : [ self . env . get_state ( ) ] , "avail_actions" : [ self . env . get_avail_actions ( ) ] , "obs" : [ self . env . get_obs ( ) ] } self . batch . update ( pre_transition_data , ts = self . t ) actions = self . mac . select_actions ( self . batch , t_ep = self . t , t_env = self . t_env , test_mode = test_mode ) reward , terminated , env_info = self . env . step ( actions [ 0 ] ) episode_return += reward for i , value in enumerate ( env_info [ 'agent_rewards' ] ) : current_episode_agent_returns [ i ] += value post_transition_data = { "actions" : actions , "reward" : [ ( reward , ) ] , "terminated" : [ ( terminated != env_info . get ( "episode_limit" , False ) , ) ] , } self . batch . update ( post_transition_data , ts = self . t ) self . t += 1 last_data = { "state" : [ self . env . get_state ( ) ] , "avail_actions" : [ self . env . get_avail_actions ( ) ] , "obs" : [ self . env . get_obs ( ) ] } self . batch . update ( last_data , ts = self . t ) actions = self . mac . select_actions ( self . batch , t_ep = self . t , t_env = self . t_env , test_mode = test_mode ) self . batch . update ( { "actions" : actions } , ts = self . t ) cur_stats = self . test_stats if test_mode else self . train_stats cur_returns = self . test_returns if test_mode else self . train_returns log_prefix = "test_" if test_mode else "" self . agent_returns . append ( current_episode_agent_returns ) env_info . pop ( 'agent_rewards' ) cur_stats . update ( { k : cur_stats . get ( k , 0 ) + env_info . get ( k , 0 ) for k in set ( cur_stats ) | set ( env_info ) } ) cur_stats [ "n_episodes" ] = 1 + cur_stats . get ( "n_episodes" , 0 ) cur_stats [ "ep_length" ] = self . t + cur_stats . get ( "ep_length" , 0 ) if not test_mode : self . t_env += self . t cur_returns . append ( episode_return ) if test_mode and ( len ( self . test_returns ) == self . args . test_nepisode ) : self . _log ( cur_returns , cur_stats , log_prefix ) self . _log_info ( log_prefix ) elif self . t_env - self . log_train_stats_t >= self . args . runner_log_interval : self . _log ( cur_returns , cur_stats , log_prefix ) self . _log_info ( log_prefix ) if hasattr ( self . mac . action_selector , "epsilon" ) : self . logger . log_stat ( "epsilon" , self . mac . action_selector . epsilon , self . t_env ) self . log_train_stats_t = self . t_env return self . batch
Items that have been added by me:
 current_episode_agent_returns
Items that have been added by me : current_episode_agent_returns
'''
        Items that have been added by me:
        * current_episode_agent_returns
        ''' # this clears the agent returns list every episode.  # TODO: Works # Pass the entire batch of experiences up till now to the agents # Receive the actions for each agent at this timestep in a batch of size 1 # TODO: works  # Select actions in the last stored state # this is where I can append current_episode_agent_returns # TODO: works # time to log 
self
test_mode