[RLlib] Fix use_lstm flag for ModelV2 (w/o ModelV1 wrapping) and add it for PyTorch. (#8734)

This commit is contained in:
Sven Mika
2020-06-05 15:40:30 +02:00
committed by GitHub
parent d78757623d
commit c74dc58f8b
21 changed files with 331 additions and 85 deletions
+1 -1
View File
@@ -121,7 +121,7 @@ Once implemented, the model can then be registered and used in place of a built-
from ray.rllib.models import ModelCatalog
from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
class CustomTorchModel(nn.Module, TorchModelV2):
class CustomTorchModel(TorchModelV2):
def __init__(self, obs_space, action_space, num_outputs, model_config, name): ...
def forward(self, input_dict, state, seq_lens): ...
def value_function(self): ...