TODO files: about lagom HOT 1 CLOSED

class CategoricalPolicy(BasePolicy):
    def __init__(self, config, network, env_spec, device, learn_V=False, **kwargs):
        super().__init__(config=config, network=network, env_spec=env_spec, device=device, **kwargs)
        if self.recurrent:
            self.reset_rnn_states()
        
    def __call__(self, x, out_keys=['action'], info={}, **kwargs):
        if self.recurrent:
            if 'mask' in info:  # make the mask
                mask = np.logical_not(info['mask']).astype(np.float32)
                mask = torch.from_numpy(mask).unsqueeze(1).to(self.device)
            else:
                mask = None
                
            out_network = self.network(x=x, 
                                       hidden_states=self.rnn_states, 
                                       mask=mask)
            features = out_network['output']
            # Update the tracking of current RNN hidden states
            if 'rnn_state_no_update' not in info:
                self.rnn_states = out_network['hidden_states']
        
class GaussianPolicy(BasePolicy):
    def __init__(self,
                 config,
                 network, 
                 env_spec, 
                 device,
                 learn_V=False,
                 min_std=1e-6, 
                 std_style='exp', 
                 constant_std=None,
                 std_state_dependent=False,
                 init_std=1.0,
                 **kwargs):
        super().__init__(config=config, network=network, env_spec=env_spec, device=device, **kwargs)
        
        # Initialize and track the RNN hidden states
        if self.recurrent:
            self.reset_rnn_states()
    
    def __call__(self, x, out_keys=['action'], info={}, **kwargs):
        # Output dictionary
        out_policy = {}
        
        # Forward pass of feature networks to obtain features
        if self.recurrent:
            if 'mask' in info:  # make the mask
                mask = np.logical_not(info['mask']).astype(np.float32)
                mask = torch.from_numpy(mask).unsqueeze(1).to(self.device)
            else:
                mask = None
                
            out_network = self.network(x=x, 
                                       hidden_states=self.rnn_states, 
                                       mask=mask)
            features = out_network['output']
            # Update the tracking of current RNN hidden states
            if 'rnn_state_no_update' not in info:
                self.rnn_states = out_network['hidden_states']