rlmodule 0.1.6.3__tar.gz → 0.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {rlmodule-0.1.6.3/rlmodule.egg-info → rlmodule-0.2.0}/PKG-INFO +18 -6
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/README.md +8 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/pyproject.toml +22 -5
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/__init__.py +1 -1
- rlmodule-0.2.0/rlmodule/skrl/torch/network/__init__.py +63 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/log_std.py +14 -7
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/model.py +62 -48
- rlmodule-0.2.0/rlmodule/source/network.py +1902 -0
- rlmodule-0.2.0/rlmodule/source/network_cfg.py +413 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/nvidia_utils/configclass.py +18 -3
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/output_layer.py +104 -11
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/output_layer_cfg.py +39 -1
- rlmodule-0.2.0/rlmodule/source/rlmodel.py +173 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/utils.py +10 -2
- {rlmodule-0.1.6.3 → rlmodule-0.2.0/rlmodule.egg-info}/PKG-INFO +18 -6
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule.egg-info/SOURCES.txt +12 -1
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule.egg-info/requires.txt +3 -3
- rlmodule-0.2.0/tests/test_build_model_integration.py +234 -0
- rlmodule-0.2.0/tests/test_cnn_backbone.py +515 -0
- rlmodule-0.2.0/tests/test_combined_log_std.py +42 -0
- rlmodule-0.2.0/tests/test_graph_network.py +856 -0
- rlmodule-0.2.0/tests/test_head_mlp.py +93 -0
- rlmodule-0.2.0/tests/test_mlp_layer_norm.py +102 -0
- rlmodule-0.2.0/tests/test_residual_mlp.py +97 -0
- rlmodule-0.2.0/tests/test_rnn_memory_and_correctness.py +444 -0
- rlmodule-0.2.0/tests/test_skrl_startup.py +561 -0
- rlmodule-0.2.0/tests/test_skrl_train_smoke.py +403 -0
- rlmodule-0.2.0/tests/test_squashed_gaussian.py +118 -0
- rlmodule-0.1.6.3/rlmodule/skrl/torch/network/__init__.py +0 -29
- rlmodule-0.1.6.3/rlmodule/source/network.py +0 -646
- rlmodule-0.1.6.3/rlmodule/source/network_cfg.py +0 -82
- rlmodule-0.1.6.3/rlmodule/source/rlmodel.py +0 -88
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/LICENSE +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/__init__.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/__init__.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/cnn_experimental.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_mlp_separate_model.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_mlp_shared_model.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_rnn_forwarded_input_shared_model.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_rnn_separate_model.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_rnn_shared_model.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/pendulum_sac_mlp_separate_model.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/skrl_no_module.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/skrl_no_module_rnn.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/skrl_sac.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/temp.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/output_layer/__init__.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/__init__.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/build.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/log_std_cfg.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/nvidia_utils/__init__.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/nvidia_utils/dict.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/nvidia_utils/string.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/rlmodel_cfg.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule.egg-info/dependency_links.txt +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule.egg-info/top_level.txt +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.2.0}/setup.cfg +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
2
|
Name: rlmodule
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.2.0
|
|
4
4
|
Summary: Flexible reinforcement learning models instantiators library
|
|
5
5
|
Author: Lopatovsky
|
|
6
6
|
Author-email: lopatovsky@gmail.com
|
|
@@ -18,18 +18,22 @@ Classifier: Intended Audience :: Science/Research
|
|
|
18
18
|
Classifier: Topic :: Scientific/Engineering
|
|
19
19
|
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
20
20
|
Classifier: Programming Language :: Python :: 3
|
|
21
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
22
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
23
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
21
24
|
Classifier: Operating System :: OS Independent
|
|
22
|
-
Requires-Python: >=3.
|
|
25
|
+
Requires-Python: >=3.10
|
|
23
26
|
Description-Content-Type: text/markdown
|
|
24
27
|
License-File: LICENSE
|
|
25
28
|
Requires-Dist: gym
|
|
26
29
|
Requires-Dist: gymnasium==0.29.0
|
|
27
30
|
Requires-Dist: torch>=1.9
|
|
28
|
-
Requires-Dist: skrl>=1
|
|
31
|
+
Requires-Dist: skrl>=2.1
|
|
29
32
|
Provides-Extra: skrl
|
|
30
|
-
Requires-Dist: skrl>=1
|
|
33
|
+
Requires-Dist: skrl>=2.1; extra == "skrl"
|
|
31
34
|
Provides-Extra: all
|
|
32
|
-
Requires-Dist: skrl>=1
|
|
35
|
+
Requires-Dist: skrl>=2.1; extra == "all"
|
|
36
|
+
Dynamic: license-file
|
|
33
37
|
|
|
34
38
|
# rlmodule
|
|
35
39
|
Flexible reinforcement learning models instantiators library
|
|
@@ -77,6 +81,14 @@ Note: sometimes installation may fail, if there is a run/ dir present, you may n
|
|
|
77
81
|
rm -rf runs
|
|
78
82
|
```
|
|
79
83
|
|
|
84
|
+
### Deep policies: LayerNorm and Residual MLPs
|
|
85
|
+
|
|
86
|
+
`rlmodule` supports per-layer LayerNorm on `MlpCfg` and a SimBa-style residual
|
|
87
|
+
MLP (`ResidualMlpCfg`). These work for both separated and shared
|
|
88
|
+
architectures, with independent per-head granularity via the new
|
|
89
|
+
`OutputLayerCfg.head` field. See [docs/layer_norm_and_residual.md](docs/layer_norm_and_residual.md)
|
|
90
|
+
for motivation, references, and copy-pasteable cfg snippets.
|
|
91
|
+
|
|
80
92
|
### Run chosen example
|
|
81
93
|
```
|
|
82
94
|
python3 rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_mlp_separate_model.py
|
|
@@ -44,6 +44,14 @@ Note: sometimes installation may fail, if there is a run/ dir present, you may n
|
|
|
44
44
|
rm -rf runs
|
|
45
45
|
```
|
|
46
46
|
|
|
47
|
+
### Deep policies: LayerNorm and Residual MLPs
|
|
48
|
+
|
|
49
|
+
`rlmodule` supports per-layer LayerNorm on `MlpCfg` and a SimBa-style residual
|
|
50
|
+
MLP (`ResidualMlpCfg`). These work for both separated and shared
|
|
51
|
+
architectures, with independent per-head granularity via the new
|
|
52
|
+
`OutputLayerCfg.head` field. See [docs/layer_norm_and_residual.md](docs/layer_norm_and_residual.md)
|
|
53
|
+
for motivation, references, and copy-pasteable cfg snippets.
|
|
54
|
+
|
|
47
55
|
### Run chosen example
|
|
48
56
|
```
|
|
49
57
|
python3 rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_mlp_separate_model.py
|
|
@@ -1,9 +1,13 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=61.0", "wheel"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
1
5
|
[project]
|
|
2
6
|
name = "rlmodule"
|
|
3
|
-
version = "0.
|
|
7
|
+
version = "0.2.0"
|
|
4
8
|
description = "Flexible reinforcement learning models instantiators library"
|
|
5
9
|
readme = "README.md"
|
|
6
|
-
requires-python = ">=3.
|
|
10
|
+
requires-python = ">=3.10"
|
|
7
11
|
license = {text = "MIT License"}
|
|
8
12
|
authors = [
|
|
9
13
|
{name = "Lopatovsky"},
|
|
@@ -20,24 +24,27 @@ classifiers = [
|
|
|
20
24
|
"Topic :: Scientific/Engineering",
|
|
21
25
|
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
22
26
|
"Programming Language :: Python :: 3",
|
|
27
|
+
"Programming Language :: Python :: 3.10",
|
|
28
|
+
"Programming Language :: Python :: 3.11",
|
|
29
|
+
"Programming Language :: Python :: 3.12",
|
|
23
30
|
"Operating System :: OS Independent",
|
|
24
31
|
]
|
|
25
32
|
dependencies = [
|
|
26
33
|
"gym",
|
|
27
34
|
"gymnasium==0.29.0",
|
|
28
35
|
"torch>=1.9",
|
|
29
|
-
"skrl>=1
|
|
36
|
+
"skrl>=2.1",
|
|
30
37
|
]
|
|
31
38
|
|
|
32
39
|
[project.optional-dependencies]
|
|
33
40
|
|
|
34
41
|
# TODO(ll) only examples should depend on skrl, not the main code.
|
|
35
42
|
skrl = [
|
|
36
|
-
"skrl>=1
|
|
43
|
+
"skrl>=2.1",
|
|
37
44
|
]
|
|
38
45
|
|
|
39
46
|
all = [
|
|
40
|
-
"skrl>=1
|
|
47
|
+
"skrl>=2.1",
|
|
41
48
|
]
|
|
42
49
|
|
|
43
50
|
# urls
|
|
@@ -49,6 +56,10 @@ all = [
|
|
|
49
56
|
"Say Thanks!" = "https://github.com/lopatovsky"
|
|
50
57
|
"Source" = "https://github.com/FabricaAI/rlmodule"
|
|
51
58
|
|
|
59
|
+
[tool.setuptools.packages.find]
|
|
60
|
+
include = ["rlmodule*"]
|
|
61
|
+
exclude = ["tests*", "docs*"]
|
|
62
|
+
|
|
52
63
|
[tool.yapf]
|
|
53
64
|
# run: yapf -p -m -i -r <folder>
|
|
54
65
|
based_on_style = "pep8"
|
|
@@ -94,3 +105,9 @@ sections = [
|
|
|
94
105
|
]
|
|
95
106
|
no_lines_before = "THIRDPARTY"
|
|
96
107
|
skip = ["docs"]
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
[tool.pytest.ini_options]
|
|
111
|
+
markers = [
|
|
112
|
+
"slow: marks tests as slow (skip with: -m 'not slow')",
|
|
113
|
+
]
|
|
@@ -1,3 +1,3 @@
|
|
|
1
|
-
__all__ = ["
|
|
1
|
+
__all__ = ["build_model", "BaseRLCfg", "RLModelCfg", "SeparatedRLModelCfg", "SharedRLModelCfg"]
|
|
2
2
|
from rlmodule.source.build import build_model # noqa: F401
|
|
3
3
|
from rlmodule.source.rlmodel_cfg import BaseRLCfg, RLModelCfg, SeparatedRLModelCfg, SharedRLModelCfg # noqa: F401
|
|
@@ -0,0 +1,63 @@
|
|
|
1
|
+
__all__ = [
|
|
2
|
+
# modules:
|
|
3
|
+
"MLP",
|
|
4
|
+
"RNN",
|
|
5
|
+
"GRU",
|
|
6
|
+
"LSTM",
|
|
7
|
+
"RnnBase",
|
|
8
|
+
"RnnMlp",
|
|
9
|
+
"RnnMlpWithForwardedInput",
|
|
10
|
+
"ResidualMlp",
|
|
11
|
+
"CnnBackbone",
|
|
12
|
+
"Cnn3dBackbone",
|
|
13
|
+
"CnnMlp",
|
|
14
|
+
"CnnProprioMlp",
|
|
15
|
+
"GraphNetwork",
|
|
16
|
+
# configs:
|
|
17
|
+
"NetworkCfg",
|
|
18
|
+
"MlpCfg",
|
|
19
|
+
"ResidualMlpCfg",
|
|
20
|
+
"RnnBaseCfg",
|
|
21
|
+
"RnnCfg",
|
|
22
|
+
"GruCfg",
|
|
23
|
+
"LstmCfg",
|
|
24
|
+
"RnnMlpCfg",
|
|
25
|
+
"CnnCfg",
|
|
26
|
+
"CnnBackboneCfg",
|
|
27
|
+
"Cnn3dBackboneCfg",
|
|
28
|
+
"CnnMlpCfg",
|
|
29
|
+
"CnnProprioMlpCfg",
|
|
30
|
+
"GraphNetworkCfg",
|
|
31
|
+
]
|
|
32
|
+
from rlmodule.source.network import ( # noqa: F401
|
|
33
|
+
CNN,
|
|
34
|
+
GRU,
|
|
35
|
+
LSTM,
|
|
36
|
+
MLP,
|
|
37
|
+
RNN,
|
|
38
|
+
Cnn3dBackbone,
|
|
39
|
+
CnnBackbone,
|
|
40
|
+
CnnMlp,
|
|
41
|
+
CnnProprioMlp,
|
|
42
|
+
GraphNetwork,
|
|
43
|
+
ResidualMlp,
|
|
44
|
+
RnnBase,
|
|
45
|
+
RnnMlp,
|
|
46
|
+
RnnMlpWithForwardedInput,
|
|
47
|
+
)
|
|
48
|
+
from rlmodule.source.network_cfg import ( # noqa: F401
|
|
49
|
+
Cnn3dBackboneCfg,
|
|
50
|
+
CnnBackboneCfg,
|
|
51
|
+
CnnCfg,
|
|
52
|
+
CnnMlpCfg,
|
|
53
|
+
CnnProprioMlpCfg,
|
|
54
|
+
GraphNetworkCfg,
|
|
55
|
+
GruCfg,
|
|
56
|
+
LstmCfg,
|
|
57
|
+
MlpCfg,
|
|
58
|
+
NetworkCfg,
|
|
59
|
+
ResidualMlpCfg,
|
|
60
|
+
RnnBaseCfg,
|
|
61
|
+
RnnCfg,
|
|
62
|
+
RnnMlpCfg,
|
|
63
|
+
)
|
|
@@ -102,18 +102,24 @@ class CombinedLogStd(LogStd):
|
|
|
102
102
|
def __init__(self, device: Union[str, torch.device], input_size: int, output_size: int, cfg):
|
|
103
103
|
super().__init__(device, input_size, output_size, cfg)
|
|
104
104
|
|
|
105
|
-
self._modules
|
|
106
|
-
|
|
107
|
-
|
|
105
|
+
# NOTE: must NOT be ``self._modules`` — that name is ``nn.Module``'s
|
|
106
|
+
# internal child-module registry. Assigning a plain list there
|
|
107
|
+
# silently breaks parameter registration (``.parameters()`` /
|
|
108
|
+
# ``.to()`` / ``state_dict()``), so the sub-LogStd params never reach
|
|
109
|
+
# the optimizer and their std stays frozen. Use a registered
|
|
110
|
+
# ``nn.ModuleList`` under a distinct attribute instead.
|
|
111
|
+
self._sub_log_stds = nn.ModuleList(
|
|
112
|
+
[module_cfg.class_type(device, input_size, output_size, module_cfg) for module_cfg in cfg.combined_modules]
|
|
113
|
+
)
|
|
108
114
|
|
|
109
115
|
if cfg.combination_constants is None:
|
|
110
|
-
self._combination_constants = list(repeat(1.0, len(self.
|
|
116
|
+
self._combination_constants = list(repeat(1.0, len(self._sub_log_stds)))
|
|
111
117
|
else:
|
|
112
118
|
self._combination_constants = cfg.combination_constants
|
|
113
|
-
if len(self.
|
|
119
|
+
if len(self._sub_log_stds) != len(self._combination_constants):
|
|
114
120
|
raise ValueError(
|
|
115
121
|
f"Number of combination constants {len(self._combination_constants)} does not match number of"
|
|
116
|
-
f" modules {len(self.
|
|
122
|
+
f" modules {len(self._sub_log_stds)}"
|
|
117
123
|
)
|
|
118
124
|
|
|
119
125
|
if cfg.combination_method == "max":
|
|
@@ -128,6 +134,7 @@ class CombinedLogStd(LogStd):
|
|
|
128
134
|
def forward(self, input: torch.Tensor) -> torch.Tensor:
|
|
129
135
|
"""Compute standard deviation as a combination from multiple Std modules."""
|
|
130
136
|
collected_stds = [
|
|
131
|
-
self._combination_constants[module_id] * module(input)
|
|
137
|
+
self._combination_constants[module_id] * module(input)
|
|
138
|
+
for module_id, module in enumerate(self._sub_log_stds)
|
|
132
139
|
]
|
|
133
140
|
return self.clip_std(reduce(self._combination_method, collected_stds))
|
|
@@ -3,14 +3,17 @@ from typing import Any, Mapping, Optional, Tuple, Union
|
|
|
3
3
|
import collections
|
|
4
4
|
import gym
|
|
5
5
|
import gymnasium
|
|
6
|
+
import logging
|
|
6
7
|
|
|
7
8
|
import torch
|
|
8
9
|
import torch.nn as nn
|
|
9
10
|
|
|
10
|
-
from rlmodule import logger
|
|
11
11
|
from rlmodule.source.utils import contains_rnn_module
|
|
12
12
|
|
|
13
13
|
|
|
14
|
+
logger = logging.getLogger("rlmodule")
|
|
15
|
+
|
|
16
|
+
|
|
14
17
|
class Model(torch.nn.Module):
|
|
15
18
|
def __init__(self, device: Union[str, torch.device], network: nn.Module) -> None:
|
|
16
19
|
"""Base class representing a function approximator"""
|
|
@@ -109,29 +112,47 @@ class Model(torch.nn.Module):
|
|
|
109
112
|
>>> model.get_specification()
|
|
110
113
|
{'rnn': {'sizes': [(1, 4, 64), (1, 4, 64)]}}
|
|
111
114
|
"""
|
|
112
|
-
if self._lstm:
|
|
113
|
-
|
|
114
|
-
|
|
115
|
-
|
|
116
|
-
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
|
|
121
|
-
|
|
122
|
-
|
|
123
|
-
|
|
115
|
+
if self._lstm or self._rnn:
|
|
116
|
+
num_layers = self._net.num_layers
|
|
117
|
+
num_envs = self._net.num_envs
|
|
118
|
+
hidden_size = self._net.hidden_size
|
|
119
|
+
# ``num_envs`` is only used here to size the SKRL-side hidden-state
|
|
120
|
+
# buffer; it MUST match the trainer's vector-env count. If it's
|
|
121
|
+
# misconfigured (<= 0 or inconsistent), SKRL silently allocates a
|
|
122
|
+
# wrong-sized tensor and the mismatch only surfaces later inside
|
|
123
|
+
# the LSTM/RNN forward as an opaque shape error. Catch it here.
|
|
124
|
+
if num_envs is None or num_envs <= 0:
|
|
125
|
+
raise ValueError(
|
|
126
|
+
f"RnnBaseCfg.num_envs must be a positive integer matching the trainer's "
|
|
127
|
+
f"vector-env count; got {num_envs!r}."
|
|
128
|
+
)
|
|
129
|
+
if num_layers is None or num_layers <= 0:
|
|
130
|
+
raise ValueError(
|
|
131
|
+
f"RnnBaseCfg.num_layers must be a positive integer; got {num_layers!r}."
|
|
132
|
+
)
|
|
133
|
+
if hidden_size is None or hidden_size <= 0:
|
|
134
|
+
raise ValueError(
|
|
135
|
+
f"RnnBaseCfg.hidden_size must be a positive integer; got {hidden_size!r}."
|
|
136
|
+
)
|
|
137
|
+
|
|
138
|
+
size = (num_layers, num_envs, hidden_size)
|
|
139
|
+
if self._lstm:
|
|
140
|
+
return {
|
|
141
|
+
"rnn": {
|
|
142
|
+
"sequence_length": self._net.sequence_length,
|
|
143
|
+
# hidden states (D * num_layers, N, Hout)
|
|
144
|
+
# cell states (D * num_layers, N, Hcell)
|
|
145
|
+
"sizes": [size, size],
|
|
146
|
+
}
|
|
124
147
|
}
|
|
125
|
-
} # cell states (D ∗ num_layers, N, Hcell)
|
|
126
|
-
elif self._rnn:
|
|
127
148
|
return {
|
|
128
149
|
"rnn": {
|
|
129
150
|
"sequence_length": self._net.sequence_length,
|
|
130
|
-
|
|
151
|
+
# hidden states (D * num_layers, N, Hout)
|
|
152
|
+
"sizes": [size],
|
|
131
153
|
}
|
|
132
|
-
}
|
|
133
|
-
|
|
134
|
-
return {}
|
|
154
|
+
}
|
|
155
|
+
return {}
|
|
135
156
|
|
|
136
157
|
def tensor_to_space(
|
|
137
158
|
self, tensor: torch.Tensor, space: Union[gym.Space, gymnasium.Space], start: int = 0
|
|
@@ -196,27 +217,25 @@ class Model(torch.nn.Module):
|
|
|
196
217
|
|
|
197
218
|
def random_act(
|
|
198
219
|
self, inputs: Mapping[str, Union[torch.Tensor, Any]], role: str = ""
|
|
199
|
-
) -> Tuple[torch.Tensor,
|
|
200
|
-
"""Act randomly according to the action space
|
|
220
|
+
) -> Tuple[torch.Tensor, Mapping[str, Union[torch.Tensor, Any]]]:
|
|
221
|
+
"""Act randomly according to the action space.
|
|
201
222
|
|
|
202
|
-
|
|
203
|
-
|
|
204
|
-
|
|
205
|
-
- ``"taken_actions"``: actions taken by the policy for the given states
|
|
206
|
-
:type inputs: dict where the values are typically torch.Tensor
|
|
207
|
-
:param role: Role play by the model (default: ``""``)
|
|
208
|
-
:type role: str, optional
|
|
223
|
+
Returns a 2-tuple ``(actions, outputs_dict)`` matching the
|
|
224
|
+
SKRL 2.x ``model.act`` contract.
|
|
225
|
+
"""
|
|
209
226
|
|
|
210
|
-
|
|
227
|
+
# SKRL 2.x agents pass the actor input under ``observations``;
|
|
228
|
+
# fall back to ``states`` when a model is invoked directly.
|
|
229
|
+
anchor = inputs.get("observations")
|
|
230
|
+
if anchor is None:
|
|
231
|
+
anchor = inputs["states"]
|
|
232
|
+
batch_size = anchor.shape[0]
|
|
211
233
|
|
|
212
|
-
:return: Model output. The first component is the action to be taken by the agent
|
|
213
|
-
:rtype: tuple of torch.Tensor, None, and dict
|
|
214
|
-
"""
|
|
215
234
|
# discrete action space (Discrete)
|
|
216
235
|
if issubclass(type(self.action_space), gym.spaces.Discrete) or issubclass(
|
|
217
236
|
type(self.action_space), gymnasium.spaces.Discrete
|
|
218
237
|
):
|
|
219
|
-
return torch.randint(self.action_space.n, (
|
|
238
|
+
return torch.randint(self.action_space.n, (batch_size, 1), device=self.device), {}
|
|
220
239
|
# continuous action space (Box)
|
|
221
240
|
elif issubclass(type(self.action_space), gym.spaces.Box) or issubclass(
|
|
222
241
|
type(self.action_space), gymnasium.spaces.Box
|
|
@@ -228,8 +247,7 @@ class Model(torch.nn.Module):
|
|
|
228
247
|
)
|
|
229
248
|
|
|
230
249
|
return (
|
|
231
|
-
self._random_distribution.sample(sample_shape=(
|
|
232
|
-
None,
|
|
250
|
+
self._random_distribution.sample(sample_shape=(batch_size, self.num_actions)),
|
|
233
251
|
{},
|
|
234
252
|
)
|
|
235
253
|
else:
|
|
@@ -372,22 +390,18 @@ class Model(torch.nn.Module):
|
|
|
372
390
|
"""
|
|
373
391
|
return self.forward(inputs, role)
|
|
374
392
|
|
|
375
|
-
def
|
|
376
|
-
"""
|
|
393
|
+
def enable_training_mode(self, enabled: bool = True) -> None:
|
|
394
|
+
"""Enable or disable training mode.
|
|
377
395
|
|
|
378
|
-
|
|
379
|
-
|
|
380
|
-
|
|
381
|
-
:type mode: str
|
|
396
|
+
SKRL 2.x agents call this from
|
|
397
|
+
``Agent.enable_models_training_mode`` to toggle the underlying
|
|
398
|
+
:py:meth:`torch.nn.Module.train` flag.
|
|
382
399
|
|
|
383
|
-
:
|
|
400
|
+
:param enabled: Set training mode if ``True``, evaluation mode if
|
|
401
|
+
``False`` (default: ``True``).
|
|
384
402
|
"""
|
|
385
|
-
|
|
386
|
-
|
|
387
|
-
elif mode == "eval":
|
|
388
|
-
self.train(False)
|
|
389
|
-
else:
|
|
390
|
-
raise ValueError("Invalid mode. Use 'train' for training or 'eval' for evaluation")
|
|
403
|
+
|
|
404
|
+
self.train(enabled)
|
|
391
405
|
|
|
392
406
|
def save(self, path: str, state_dict: Optional[dict] = None) -> None:
|
|
393
407
|
"""Save the model to the specified path
|