|
|
from abc import ABC, abstractmethod |
|
|
from collections.abc import Sequence |
|
|
from typing import Generic, TypeVar |
|
|
|
|
|
from tianshou.highlevel.env import Environments |
|
|
from tianshou.highlevel.module.core import TDevice |
|
|
from tianshou.highlevel.module.intermediate import IntermediateModuleFactory |
|
|
from tianshou.highlevel.optim import OptimizerFactory |
|
|
from tianshou.policy import BasePolicy, ICMPolicy |
|
|
from tianshou.utils.net.discrete import IntrinsicCuriosityModule |
|
|
from tianshou.utils.string import ToStringMixin |
|
|
|
|
|
TPolicyOut = TypeVar("TPolicyOut", bound=BasePolicy) |
|
|
|
|
|
|
|
|
class PolicyWrapperFactory(Generic[TPolicyOut], ToStringMixin, ABC): |
|
|
@abstractmethod |
|
|
def create_wrapped_policy( |
|
|
self, |
|
|
policy: BasePolicy, |
|
|
envs: Environments, |
|
|
optim_factory: OptimizerFactory, |
|
|
device: TDevice, |
|
|
) -> TPolicyOut: |
|
|
pass |
|
|
|
|
|
|
|
|
class PolicyWrapperFactoryIntrinsicCuriosity( |
|
|
PolicyWrapperFactory[ICMPolicy], |
|
|
): |
|
|
def __init__( |
|
|
self, |
|
|
*, |
|
|
feature_net_factory: IntermediateModuleFactory, |
|
|
hidden_sizes: Sequence[int], |
|
|
lr: float, |
|
|
lr_scale: float, |
|
|
reward_scale: float, |
|
|
forward_loss_weight: float, |
|
|
): |
|
|
self.feature_net_factory = feature_net_factory |
|
|
self.hidden_sizes = hidden_sizes |
|
|
self.lr = lr |
|
|
self.lr_scale = lr_scale |
|
|
self.reward_scale = reward_scale |
|
|
self.forward_loss_weight = forward_loss_weight |
|
|
|
|
|
def create_wrapped_policy( |
|
|
self, |
|
|
policy: BasePolicy, |
|
|
envs: Environments, |
|
|
optim_factory: OptimizerFactory, |
|
|
device: TDevice, |
|
|
) -> ICMPolicy: |
|
|
feature_net = self.feature_net_factory.create_intermediate_module(envs, device) |
|
|
action_dim = envs.get_action_shape() |
|
|
if not isinstance(action_dim, int): |
|
|
raise ValueError(f"Environment action shape must be an integer, got {action_dim}") |
|
|
feature_dim = feature_net.output_dim |
|
|
icm_net = IntrinsicCuriosityModule( |
|
|
feature_net.module, |
|
|
feature_dim, |
|
|
action_dim, |
|
|
hidden_sizes=self.hidden_sizes, |
|
|
device=device, |
|
|
) |
|
|
icm_optim = optim_factory.create_optimizer(icm_net, lr=self.lr) |
|
|
return ICMPolicy( |
|
|
policy=policy, |
|
|
model=icm_net, |
|
|
optim=icm_optim, |
|
|
action_space=envs.get_action_space(), |
|
|
lr_scale=self.lr_scale, |
|
|
reward_scale=self.reward_scale, |
|
|
forward_loss_weight=self.forward_loss_weight, |
|
|
).to(device) |
|
|
|
