Spaces:
Sleeping
Sleeping
| # model.py | |
| import torch | |
| import torch.nn as nn | |
| class ModuleLayer(torch.nn.Module): | |
| """Class for the individual layer blocks.""" | |
| def __init__(self, intermediate_dim=32, dropout_rate=0.1): | |
| super().__init__() | |
| self.mod_linear = torch.nn.Linear(intermediate_dim, intermediate_dim) | |
| self.mod_norm = torch.nn.LayerNorm(normalized_shape=intermediate_dim) | |
| self.mod_relu = torch.nn.ReLU() | |
| self.dropout = torch.nn.Dropout(p=dropout_rate) | |
| def forward(self, x): | |
| """Forward pass of the layer block.""" | |
| residual = x | |
| x = self.mod_linear(x) | |
| x = self.mod_norm(x) | |
| x = self.mod_relu(x) | |
| x = self.dropout(x) | |
| x += residual | |
| return x | |
| class Agent(torch.nn.Module): | |
| """Class for Agent Structure using multiple Layer Blocks.""" | |
| def __init__(self, cfg): | |
| """Initializer for the 'Agent' class. | |
| Args: | |
| cfg (dict): Configuration dictionary containing model parameters. | |
| """ | |
| super().__init__() | |
| self.linear = torch.nn.Linear(in_features=cfg["in_dim"], out_features=cfg["intermediate_dim"]) | |
| self.layers = torch.nn.Sequential(*[ModuleLayer(intermediate_dim=cfg["intermediate_dim"], dropout_rate=cfg["dropout_rate"]) for _ in range(cfg["num_blocks"])]) | |
| self.out = torch.nn.Linear(in_features=cfg["intermediate_dim"], out_features=cfg["out_dim"]) | |
| def forward(self, x): | |
| """Forward pass of Agent's Layers.""" | |
| x = self.linear(x) | |
| x = self.layers(x) | |
| x = self.out(x) | |
| return x | |
| def get_prediction(self, features): | |
| """Get the deterministic prediction on a single observation or a batch of observations. | |
| Parameters: | |
| features (torch.tensor): the agent's input features. Expected shape is either `(num_features,)` for a single observation | |
| or `(batch_size, num_features)` for a batch of observations. | |
| Returns: | |
| action (int or torch.tensor): | |
| - If `features` is a single features (i.e., `features.dim() == 1`), returns a scalar `int` representing the chosen action. | |
| - If `features` is a batch of features (i.e., `features.dim() > 1`), | |
| returns a `torch.Tensor` of `int`s, where each element is the | |
| chosen action for the corresponding observation in the batch""" | |
| # Ensure single samples have a batch dimension | |
| if features.dim() == 1: | |
| features = features.unsqueeze(0) # Add a batch dimension if it's a single batch of features | |
| with torch.no_grad(): | |
| if not isinstance(features, torch.Tensor): # Check if features is not already a tensor | |
| features = torch.tensor(features, dtype=torch.float) | |
| prediction = self.forward(features) # Run a forward pass through the model | |
| if features.size(0) == 1: # This method checks if there is only 1 element in a 1-D tensor | |
| return prediction.item() # Returns a Python scalar for a single observation | |
| else: | |
| return prediction # Returns a tensor of predictions |