def __init__()

in ml3/mbrl_utils.py [0:0]

• 21 lines of code
• 2 McCabe index (conditional complexity)

    def __init__(self, config):
        super(PModel, self).__init__()
        if config["device"] == "gpu":
            self.device = config["gpu_name"]
        else:
            self.device = "cpu"
        self.input_sz = config['state_dim'] + config['action_dim']
        self.output_sz = config['output_dim']

        self.learning_rate = config["learning_rate"]
        self.display_epoch = config["display_epoch"]
        self.epochs = config["epochs"]

        w = config["hidden_size"]

        self.layers = nn.Sequential(
            nn.Linear(self.input_sz, w[0]),
            nn.Tanh(),
            nn.Linear(w[0], w[1]),
            nn.Tanh(),
        )

        self.mean = nn.Linear(w[1], self.output_sz)
        self.var = nn.Sequential(nn.Linear(w[1], self.output_sz), nn.Softplus())
        self.to(self.device)