Missing fields for checkpoint in DQNAgent ? about rlcard HOT 9 CLOSED

def checkpoint_attributes(self):
    '''
    Return the current checkpoint attributes (dict)
    Checkpoint attributes are used to save and restore the model in the middle of training
    Saves the model state dict, optimizer state dict, and all other instance variables
    '''
    
    return {
        'agent_type': 'DQNAgent',
        'q_estimator': self.q_estimator.checkpoint_attributes(),
        'memory': self.memory.checkpoint_attributes(),
        'total_t': self.total_t,
        'train_t': self.train_t,
        'epsilon_start': self.epsilons.min(),
        'epsilon_end': self.epsilons.max(),
        'epsilon_decay_steps': self.epsilon_decay_steps,
        'discount_factor': self.discount_factor,
        'update_target_estimator_every': self.update_target_estimator_every,
        'batch_size': self.batch_size,
        'num_actions': self.num_actions,
        'train_every': self.train_every,
        'device': self.device,
        # add the following:
        'use_raw': self.use_raw,
        'replay_memory_init_size': self.replay_memory_init_size,
        'model_name': self.model_name,
        'save_path': self.save_path,
        'save_every': self.save_every
    }

@classmethod
def from_checkpoint(cls, checkpoint):
    '''
    Restore the model from a checkpoint
    
    Args:
        checkpoint (dict): the checkpoint attributes generated by checkpoint_attributes()
    '''
    
    print("\nINFO - Restoring model from checkpoint...")
    agent_instance = cls(
        replay_memory_size=checkpoint['memory']['memory_size'],
        update_target_estimator_every=checkpoint['update_target_estimator_every'],
        discount_factor=checkpoint['discount_factor'],
        epsilon_start=checkpoint['epsilon_start'],
        epsilon_end=checkpoint['epsilon_end'],
        epsilon_decay_steps=checkpoint['epsilon_decay_steps'],
        batch_size=checkpoint['batch_size'],
        num_actions=checkpoint['num_actions'], 
        device=checkpoint['device'], 
        state_shape=checkpoint['q_estimator']['state_shape'],
        mlp_layers=checkpoint['q_estimator']['mlp_layers'],
        train_every=checkpoint['train_every'],
        # add the following:
        replay_memory_init_size=checkpoint['replay_memory_init_size'],
        save_every=checkpoint['save_every'],
        learning_rate=checkpoint['q_estimator']['learning_rate'],
        model_name=checkpoint['model_name'],
        save_path=checkpoint['save_path'],
        use_raw=checkpoint['use_raw']
    )

    etc.

tmp = self.total_t - self.replay_memory_init_size
        if tmp>=0 #...