Potential Predictor Module for Corrected Actions.py

class PredictorModule(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.ReLU(),
            nn.Linear(128, output_dim)
        )
    
    def forward(self, x):
        return self.network(x)

predictor = PredictorModule(input_dim=10, output_dim=3)