ray/rllib/algorithms/ars/ars_torch_policy.py
2022-05-06 19:11:15 +02:00

20 lines
591 B
Python

# Code in this file is adapted from:
# https://github.com/openai/evolution-strategies-starter.
import ray
from ray.rllib.algorithms.es.es_torch_policy import (
after_init,
before_init,
make_model_and_action_dist,
)
from ray.rllib.policy.policy_template import build_policy_class
ARSTorchPolicy = build_policy_class(
name="ARSTorchPolicy",
framework="torch",
loss_fn=None,
get_default_config=lambda: ray.rllib.algorithms.ars.ars.DEFAULT_CONFIG,
before_init=before_init,
after_init=after_init,
make_model_and_action_dist=make_model_and_action_dist,
)