Stable Baselines の使いやすいインターフェイスを使ったままで、独自の方策ネットワークを定義することも可能です。
以下、Stable Baselinesドキュメントの該当箇所の翻訳です。
-----
本ドキュメントは、Stable Baselines3 ドキュメントの Custom Policy Network を翻訳したものです。
カスタム方策ネットワーク
Stable Baselines3 は、画像(CnnPolicy)及びその他タイプの入力機能(MlpPolicy)の方策ネットワークを提供します。
警告
A2CおよびPPOの場合、トレーニングやテスト中に（範囲外のエラーを回避するため）継続的な行動がクリップされます。SAC、DDPG、TD3は、`tanh()` 境界をより正確に処理する変換を使って、行動を押しつぶします。
カスタム方策アーキテクチャ
カスタム方策アーキテクチャをカスタマイズするひとつの方法として、モデルを生成する際の引数に policy_kwargs パラメータを渡すことです。
import gym
import torch as th
from stable_baselines3 import PPO
# 活性化関数にRelu を使う、サイズが32である2つのレイヤを備えたカスタムMLP(多層パーセプトロン)Policy
policy_kwargs = dict(activation_fn=th.nn.ReLU, net_arch=[32, 32])
# エージェントの生成
model = PPO("MlpPolicy", "CartPole-v1", policy_kwargs=policy_kwargs, verbose=1)
# 環境の取得
env = model.get_env()
# エージェントのトレーニング
model.learn(total_timesteps=100000)
# エージェントの保存
model.save("ppo-cartpole")
del model
# policy_kwargs も自動的にロードされる
model = PPO.load("ppo-cartpole")
方策（もしくは価値）ネットワークのカスタムアーキテクチャも簡単に定義することができます。
注意
カスタム方策クラスを定義することと、引数として policy_kwargs を渡すことは同じ意味をさしていますが、ポリシーに名前をつけることができるため、コード上では意図が明確に示しやすくなります。policy_kwargs は、ハイパーパラメータ検索を行う場合には便利です。
カスタム特徴抽出機(Feature Extractor)
(画像を使用する際のカスタムCNNなど)カスタム特徴抽出器が必要な場合は、 BaseFeatureExtractor を継承したクラスを定義して、トレーニング処理にてモデルを生成する差異に引数として渡します。
import gym
import torch as th
import torch.nn as nn
from stable_baselines3 import PPO
from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
class CustomCNN(BaseFeaturesExtractor):
"""
:param observation_space: (gym.Space) 観測空間
:param features_dim: (int) 抽出される特徴量の次元（最後のレイヤのユニット数に相当）
"""
def __init__(self, observation_space: gym.spaces.Box, features_dim: int = 256):
super(CustomCNN, self).__init__(observation_space, features_dim)
# (チャネルが最初にくる) C×H×W イメージを想定
# 並べ直しは事前処理もしくはラッパクラスにより実行されることを想定
n_input_channels = observation_space.shape[0]
self.cnn = nn.Sequential(
nn.Conv2d(n_input_channels, 32, kernel_size=8, stride=4, padding=0),
nn.ReLU(),
nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=0),
nn.ReLU(),
nn.Flatten(),
)
# 観測空間のサンプルを１件取得して形式を算出
with th.no_grad():
n_flatten = self.cnn(
th.as_tensor(observation_space.sample()[None]).float()
).shape[1]
self.linear = nn.Sequential(nn.Linear(n_flatten, features_dim), nn.ReLU())
def forward(self, observations: th.Tensor) -> th.Tensor:
return self.linear(self.cnn(observations))
policy_kwargs = dict(
features_extractor_class=CustomCNN,
features_extractor_kwargs=dict(features_dim=128),
)
model = PPO("CnnPolicy", "BreakoutNoFrameskip-v4", policy_kwargs=policy_kwargs, verbose=1)
model.learn(1000)
オンポリシーアルゴリズム
共有ネットワーク
A2CやPPOの方策パラメータ `net_arch` を使用すると、非表示レイヤ数やサイズ、方策ネットワークと価値ネットワーク間で共有されるレイヤ数を指定することができます。これは次のようなリストであると想定されています。
- 0を含む任意の共有レイヤのユニット数を指定する任意の長さをあらわす整数(0を含む)。intの数が0の場合、共有レイヤはありません。
- 価値ネットワークや方策ネットワークの非共有レイヤを指定するための辞書(オプション)。dict(vf=[<価値ネットワークレイヤサイズ>], pi=[<方策ネットワークレイヤサイズ>]) のような形式。pi や vf のいずれかのキーが欠落している場合、非共有レイヤ(空のリスト)は想定されません。
要約すると、 [<共有レイヤ数>, dict(vf=[<非共有価値ネットワークレイヤ数>], pi=[<非共有方策ネットワークレイヤ数>])] という形式のリストを指します。
例
サイズが128である2層の共有レイヤの場合：net_arch=[128, 128]
観測
|
<128>
|
<128>
/ \
行動 価値
方策ネットワークより価値ネットワークのほうが層が厚く、第１層を共有する場合：net_arch=[128, dict(vf=[256, 256])]
観測
|
<128>
/ \
行動 <256>
|
<256>
|
価値
第１層を共有し、第２層目から分岐する場合：net_arch=[128, dict(vf=[256], pi=[16])]
観測
|
<128>
/ \
<16> <256>
| |
行動 価値
高度な例
方策/価値アーキテクチャをさらに細かく制御する必要がある場合は、方策を直接再定義することができます。
from typing import Callable, Dict, List, Optional, Tuple, Type, Union
import gym
import torch as th
from torch import nn
from stable_baselines3 import PPO
from stable_baselines3.common.policies import ActorCriticPolicy
class CustomNetwork(nn.Module):
"""
方策および価値関数のカスタムネットワークをあらわすクラス。
特徴抽出器(Feature Extractor)によって抽出された特徴量を入力として受け取る。
:param feature_dim: feature_extractor で抽出された特徴量の次元（例：CNNにより抽出された特徴量）
:param last_layer_dim_pi: 方策ネットワーク最終層のユニット数(int)
:param last_layer_dim_vf: 価値ネットワーク最終層のユニット数(int)
"""
def __init__(
self,
feature_dim: int,
last_layer_dim_pi: int = 64,
last_layer_dim_vf: int = 64,
):
super(CustomNetwork, self).__init__()
# 重要：
# 確率分布の作成に使用される出力層次元数を保存
self.latent_dim_pi = last_layer_dim_pi
self.latent_dim_vf = last_layer_dim_vf
# 方策ネットワーク
self.policy_net = nn.Sequential(
nn.Linear(feature_dim, last_layer_dim_pi), nn.ReLU()
)
# 価値ネットワーク
self.value_net = nn.Sequential(
nn.Linear(feature_dim, last_layer_dim_vf), nn.ReLU()
)
def forward(self, features: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
"""
:return: (th.Tensor, th.Tensor) 指定ネットワークの latent_policy(th.Tensor形式), latent_value(th.Tensor形式) 。
すべての層が共有の場合、``latent_policy == latent_value`` となる。
"""
return self.policy_net(features), self.value_net(features)
class CustomActorCriticPolicy(ActorCriticPolicy):
def __init__(
self,
observation_space: gym.spaces.Space,
action_space: gym.spaces.Space,
lr_schedule: Callable[[float], float],
net_arch: Optional[List[Union[int, Dict[str, List[int]]]]] = None,
activation_fn: Type[nn.Module] = nn.Tanh,
*args,
**kwargs,
):
super(CustomActorCriticPolicy, self).__init__(
observation_space,
action_space,
lr_schedule,
net_arch,
activation_fn,
# 残りの引数群はベースクラスに渡される
*args,
**kwargs,
)
# 直行初期化の無効化
self.ortho_init = False
def _build_mlp_extractor(self) -> None:
self.mlp_extractor = CustomNetwork(self.features_dim)
model = PPO(CustomActorCriticPolicy, "CartPole-v1", verbose=1)
model.learn(5000)
オフポリシーアルゴリズム
SACやDDPG、TD3を使用する際に行動器(actor)と評価器(critic)が異なるネットワークアーキテクチャが必要な場合、次のような形式の辞書を引数に渡します： dict(qf=[<評価器ネットワークアーキテクチャ>], pi=[<行動器ネットワークアーキテクチャ>])
例えば、行動器ネットワーク（別名pi）と評価器ネットワーク（別名qf）を別のアーキテクチャにする場合、net_arch=dict(qf=[400, 300], pi=[64, 64]) のように指定します。
行動器ネットワークと評価器ネットワークを同一アーキテクチャを使う場合は、net_arch=[256, 256] のように指定します（ここでは256ユニットの2つの隠れ層を指定しています）。
注意
オンポリシーとは対照的に、行動器と評価器の間の（特徴抽出器以外の）共有ネットワークは許可されていません。
from stable_baselines3 import SAC
# それぞれ64ユニットの2層を持つカスタム行動器アーキテクチャ
# 400ユニットと300ユニットの2層を持つカスタム評価器アーキテクチャ
policy_kwargs = dict(net_arch=dict(pi=[64, 64], qf=[400, 300]))
# エージェントの生成
model = SAC("MlpPolicy", "Pendulum-v0", policy_kwargs=policy_kwargs, verbose=1)
model.learn(5000)-----
カスタムSIで独自モデルを作ることは殆どないとは思いますが、ご参考まで。
p.s.
Donkeycar のモデルを強化学習化するには、最後の「高度な例」セクションの手段を取れば可能になります。
..というか、gym-donkeycar使うか..ふつー..
