wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-09-12 12:51:15 +08:00
Code Issues Packages Projects Releases Wiki Activity

[rllib] ModelV2 API (#4926)

Browse Source
This commit is contained in:
Eric Liang
2019-07-03 15:59:47 -07:00
committed by GitHub
parent 9e0192bc0b
commit 34d054ff19
42 changed files with 1641 additions and 518 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/source/rllib-models.rst
+2 -2
View File
@@ -184,14 +184,14 @@ You can use ``tf.layers.batch_normalization(x, training=input_dict["is_training"
Custom Models (PyTorch)
-----------------------
Similarly, you can create and register custom PyTorch models for use with PyTorch-based algorithms (e.g., A2C, PG, QMIX). See these examples of `fully connected <https://github.com/ray-project/ray/blob/master/python/ray/rllib/models/pytorch/fcnet.py>`__, `convolutional <https://github.com/ray-project/ray/blob/master/python/ray/rllib/models/pytorch/visionnet.py>`__, and `recurrent <https://github.com/ray-project/ray/blob/master/python/ray/rllib/agents/qmix/model.py>`__ torch models.
Similarly, you can create and register custom PyTorch models for use with PyTorch-based algorithms (e.g., A2C, PG, QMIX). See these examples of `fully connected <https://github.com/ray-project/ray/blob/master/python/ray/rllib/models/torch/fcnet.py>`__, `convolutional <https://github.com/ray-project/ray/blob/master/python/ray/rllib/models/torch/visionnet.py>`__, and `recurrent <https://github.com/ray-project/ray/blob/master/python/ray/rllib/agents/qmix/model.py>`__ torch models.
.. code-block:: python
import ray
from ray.rllib.agents import a3c
from ray.rllib.models import ModelCatalog
from ray.rllib.models.pytorch.model import TorchModel
from ray.rllib.models.torch.model import TorchModel
class CustomTorchModel(TorchModel):