mirror of
https://github.com/saymrwulf/stable-baselines3.git
synced 2026-09-05 20:30:42 +00:00
58 lines
2 KiB
Python
58 lines
2 KiB
Python
import torch.nn as nn
|
|
|
|
|
|
class BasePolicy(nn.Module):
|
|
"""
|
|
The base policy object
|
|
|
|
:param observation_space: (Gym Space) The observation space of the environment
|
|
:param action_space: (Gym Space) The action space of the environment
|
|
"""
|
|
|
|
def __init__(self, observation_space, action_space, device='cpu'):
|
|
super(BasePolicy, self).__init__()
|
|
self.observation_space = observation_space
|
|
self.action_space = action_space
|
|
self.device = device
|
|
|
|
|
|
_policy_registry = dict()
|
|
|
|
|
|
def get_policy_from_name(base_policy_type, name):
|
|
"""
|
|
returns the registed policy from the base type and name
|
|
|
|
:param base_policy_type: (BasePolicy) the base policy object
|
|
:param name: (str) the policy name
|
|
:return: (base_policy_type) the policy
|
|
"""
|
|
if base_policy_type not in _policy_registry:
|
|
raise ValueError("Error: the policy type {} is not registered!".format(base_policy_type))
|
|
if name not in _policy_registry[base_policy_type]:
|
|
raise ValueError("Error: unknown policy type {}, the only registed policy type are: {}!"
|
|
.format(name, list(_policy_registry[base_policy_type].keys())))
|
|
return _policy_registry[base_policy_type][name]
|
|
|
|
|
|
def register_policy(name, policy):
|
|
"""
|
|
returns the registed policy from the base type and name
|
|
|
|
:param name: (str) the policy name
|
|
:param policy: (subclass of BasePolicy) the policy
|
|
"""
|
|
sub_class = None
|
|
for cls in BasePolicy.__subclasses__():
|
|
if issubclass(policy, cls):
|
|
sub_class = cls
|
|
break
|
|
if sub_class is None:
|
|
raise ValueError("Error: the policy {} is not of any known subclasses of BasePolicy!".format(policy))
|
|
|
|
if sub_class not in _policy_registry:
|
|
_policy_registry[sub_class] = {}
|
|
if name in _policy_registry[sub_class]:
|
|
raise ValueError("Error: the name {} is alreay registered for a different policy, will not override."
|
|
.format(name))
|
|
_policy_registry[sub_class][name] = policy
|