torchrl.trainers.algorithms.configs.modules.TensorDictModuleConfig

class torchrl.trainers.algorithms.configs.modules.TensorDictModuleConfig(_partial_: bool = False, in_keys: Optional[Any] = None, out_keys: Optional[Any] = None, module: MLPConfig = '???', _target_: str = 'tensordict.nn.TensorDictModule')

用于配置 TensorDictModule 的类。

示例

>>> cfg = TensorDictModuleConfig(module=MLPConfig(in_features=10, out_features=10, depth=2, num_cells=32), in_keys=["observation"], out_keys=["action"])
>>> module = instantiate(cfg)
>>> assert isinstance(module, TensorDictModule)
>>> assert module(observation=torch.randn(10, 10)).shape == (10, 10)

另请参阅

tensordict.nn.TensorDictModule