def __init__()

in lerobot/common/policies/tdmpc/modeling_tdmpc.py [0:0]

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

    def __init__(self, config: TDMPCConfig):
        super().__init__()
        self.config = config
        self._encoder = TDMPCObservationEncoder(config)
        self._dynamics = nn.Sequential(
            nn.Linear(config.latent_dim + config.action_feature.shape[0], config.mlp_dim),
            nn.LayerNorm(config.mlp_dim),
            nn.Mish(),
            nn.Linear(config.mlp_dim, config.mlp_dim),
            nn.LayerNorm(config.mlp_dim),
            nn.Mish(),
            nn.Linear(config.mlp_dim, config.latent_dim),
            nn.LayerNorm(config.latent_dim),
            nn.Sigmoid(),
        )
        self._reward = nn.Sequential(
            nn.Linear(config.latent_dim + config.action_feature.shape[0], config.mlp_dim),
            nn.LayerNorm(config.mlp_dim),
            nn.Mish(),
            nn.Linear(config.mlp_dim, config.mlp_dim),
            nn.LayerNorm(config.mlp_dim),
            nn.Mish(),
            nn.Linear(config.mlp_dim, 1),
        )
        self._pi = nn.Sequential(
            nn.Linear(config.latent_dim, config.mlp_dim),
            nn.LayerNorm(config.mlp_dim),
            nn.Mish(),
            nn.Linear(config.mlp_dim, config.mlp_dim),
            nn.LayerNorm(config.mlp_dim),
            nn.Mish(),
            nn.Linear(config.mlp_dim, config.action_feature.shape[0]),
        )
        self._Qs = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Linear(config.latent_dim + config.action_feature.shape[0], config.mlp_dim),
                    nn.LayerNorm(config.mlp_dim),
                    nn.Tanh(),
                    nn.Linear(config.mlp_dim, config.mlp_dim),
                    nn.ELU(),
                    nn.Linear(config.mlp_dim, 1),
                )
                for _ in range(config.q_ensemble_size)
            ]
        )
        self._V = nn.Sequential(
            nn.Linear(config.latent_dim, config.mlp_dim),
            nn.LayerNorm(config.mlp_dim),
            nn.Tanh(),
            nn.Linear(config.mlp_dim, config.mlp_dim),
            nn.ELU(),
            nn.Linear(config.mlp_dim, 1),
        )
        self._init_weights()