metastable-baselines/sb3_trl/trl_pg/policies.py

8 lines
304 B
Python

# This file is here just to define MlpPolicy/CnnPolicy
# that work for TRL_PG
from stable_baselines3.common.policies import ActorCriticCnnPolicy, ActorCriticPolicy, MultiInputActorCriticPolicy
MlpPolicy = ActorCriticPolicy
CnnPolicy = ActorCriticCnnPolicy
MultiInputPolicy = MultiInputActorCriticPolicy