rlmodule 0.1.6.3__tar.gz → 0.1.7__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.1.7}/PKG-INFO +18 -6
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/README.md +8 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/pyproject.toml +22 -5
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/skrl/torch/__init__.py +1 -1
- rlmodule-0.1.7/rlmodule/skrl/torch/network/__init__.py +59 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/source/model.py +58 -47
- rlmodule-0.1.7/rlmodule/source/network.py +1747 -0
- rlmodule-0.1.7/rlmodule/source/network_cfg.py +384 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/source/nvidia_utils/configclass.py +18 -3
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/source/output_layer.py +45 -9
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/source/output_layer_cfg.py +21 -1
- rlmodule-0.1.7/rlmodule/source/rlmodel.py +170 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/source/utils.py +10 -2
- {rlmodule-0.1.6.3 → rlmodule-0.1.7/rlmodule.egg-info}/PKG-INFO +18 -6
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule.egg-info/SOURCES.txt +10 -1
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule.egg-info/requires.txt +3 -3
- rlmodule-0.1.7/tests/test_build_model_integration.py +234 -0
- rlmodule-0.1.7/tests/test_cnn_backbone.py +311 -0
- rlmodule-0.1.7/tests/test_graph_network.py +787 -0
- rlmodule-0.1.7/tests/test_head_mlp.py +93 -0
- rlmodule-0.1.7/tests/test_mlp_layer_norm.py +102 -0
- rlmodule-0.1.7/tests/test_residual_mlp.py +97 -0
- rlmodule-0.1.7/tests/test_rnn_memory_and_correctness.py +444 -0
- rlmodule-0.1.7/tests/test_skrl_startup.py +561 -0
- rlmodule-0.1.7/tests/test_skrl_train_smoke.py +403 -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.1.7}/LICENSE +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/__init__.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/skrl/__init__.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/skrl/torch/examples/gymnasium/cnn_experimental.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_mlp_separate_model.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_mlp_shared_model.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_rnn_forwarded_input_shared_model.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_rnn_separate_model.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_rnn_shared_model.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/skrl/torch/examples/gymnasium/pendulum_sac_mlp_separate_model.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/skrl/torch/examples/gymnasium/skrl_no_module.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/skrl/torch/examples/gymnasium/skrl_no_module_rnn.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/skrl/torch/examples/gymnasium/skrl_sac.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/skrl/torch/examples/gymnasium/temp.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/skrl/torch/output_layer/__init__.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/source/__init__.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/source/build.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/source/log_std.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/source/log_std_cfg.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/source/nvidia_utils/__init__.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/source/nvidia_utils/dict.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/source/nvidia_utils/string.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule/source/rlmodel_cfg.py +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule.egg-info/dependency_links.txt +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/rlmodule.egg-info/top_level.txt +0 -0
- {rlmodule-0.1.6.3 → rlmodule-0.1.7}/setup.cfg +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
2
|
Name: rlmodule
|
|
3
|
-
Version: 0.1.
|
|
3
|
+
Version: 0.1.7
|
|
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.1.
|
|
7
|
+
version = "0.1.7"
|
|
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,59 @@
|
|
|
1
|
+
__all__ = [
|
|
2
|
+
# modules:
|
|
3
|
+
"MLP",
|
|
4
|
+
"RNN",
|
|
5
|
+
"GRU",
|
|
6
|
+
"LSTM",
|
|
7
|
+
"RnnBase",
|
|
8
|
+
"RnnMlp",
|
|
9
|
+
"RnnMlpWithForwardedInput",
|
|
10
|
+
"ResidualMlp",
|
|
11
|
+
"CnnBackbone",
|
|
12
|
+
"CnnMlp",
|
|
13
|
+
"CnnProprioMlp",
|
|
14
|
+
"GraphNetwork",
|
|
15
|
+
# configs:
|
|
16
|
+
"NetworkCfg",
|
|
17
|
+
"MlpCfg",
|
|
18
|
+
"ResidualMlpCfg",
|
|
19
|
+
"RnnBaseCfg",
|
|
20
|
+
"RnnCfg",
|
|
21
|
+
"GruCfg",
|
|
22
|
+
"LstmCfg",
|
|
23
|
+
"RnnMlpCfg",
|
|
24
|
+
"CnnCfg",
|
|
25
|
+
"CnnBackboneCfg",
|
|
26
|
+
"CnnMlpCfg",
|
|
27
|
+
"CnnProprioMlpCfg",
|
|
28
|
+
"GraphNetworkCfg",
|
|
29
|
+
]
|
|
30
|
+
from rlmodule.source.network import ( # noqa: F401
|
|
31
|
+
CNN,
|
|
32
|
+
GRU,
|
|
33
|
+
LSTM,
|
|
34
|
+
MLP,
|
|
35
|
+
RNN,
|
|
36
|
+
CnnBackbone,
|
|
37
|
+
CnnMlp,
|
|
38
|
+
CnnProprioMlp,
|
|
39
|
+
GraphNetwork,
|
|
40
|
+
ResidualMlp,
|
|
41
|
+
RnnBase,
|
|
42
|
+
RnnMlp,
|
|
43
|
+
RnnMlpWithForwardedInput,
|
|
44
|
+
)
|
|
45
|
+
from rlmodule.source.network_cfg import ( # noqa: F401
|
|
46
|
+
CnnBackboneCfg,
|
|
47
|
+
CnnCfg,
|
|
48
|
+
CnnMlpCfg,
|
|
49
|
+
CnnProprioMlpCfg,
|
|
50
|
+
GraphNetworkCfg,
|
|
51
|
+
GruCfg,
|
|
52
|
+
LstmCfg,
|
|
53
|
+
MlpCfg,
|
|
54
|
+
NetworkCfg,
|
|
55
|
+
ResidualMlpCfg,
|
|
56
|
+
RnnBaseCfg,
|
|
57
|
+
RnnCfg,
|
|
58
|
+
RnnMlpCfg,
|
|
59
|
+
)
|
|
@@ -109,29 +109,47 @@ class Model(torch.nn.Module):
|
|
|
109
109
|
>>> model.get_specification()
|
|
110
110
|
{'rnn': {'sizes': [(1, 4, 64), (1, 4, 64)]}}
|
|
111
111
|
"""
|
|
112
|
-
if self._lstm:
|
|
113
|
-
|
|
114
|
-
|
|
115
|
-
|
|
116
|
-
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
|
|
121
|
-
|
|
122
|
-
|
|
123
|
-
|
|
112
|
+
if self._lstm or self._rnn:
|
|
113
|
+
num_layers = self._net.num_layers
|
|
114
|
+
num_envs = self._net.num_envs
|
|
115
|
+
hidden_size = self._net.hidden_size
|
|
116
|
+
# ``num_envs`` is only used here to size the SKRL-side hidden-state
|
|
117
|
+
# buffer; it MUST match the trainer's vector-env count. If it's
|
|
118
|
+
# misconfigured (<= 0 or inconsistent), SKRL silently allocates a
|
|
119
|
+
# wrong-sized tensor and the mismatch only surfaces later inside
|
|
120
|
+
# the LSTM/RNN forward as an opaque shape error. Catch it here.
|
|
121
|
+
if num_envs is None or num_envs <= 0:
|
|
122
|
+
raise ValueError(
|
|
123
|
+
f"RnnBaseCfg.num_envs must be a positive integer matching the trainer's "
|
|
124
|
+
f"vector-env count; got {num_envs!r}."
|
|
125
|
+
)
|
|
126
|
+
if num_layers is None or num_layers <= 0:
|
|
127
|
+
raise ValueError(
|
|
128
|
+
f"RnnBaseCfg.num_layers must be a positive integer; got {num_layers!r}."
|
|
129
|
+
)
|
|
130
|
+
if hidden_size is None or hidden_size <= 0:
|
|
131
|
+
raise ValueError(
|
|
132
|
+
f"RnnBaseCfg.hidden_size must be a positive integer; got {hidden_size!r}."
|
|
133
|
+
)
|
|
134
|
+
|
|
135
|
+
size = (num_layers, num_envs, hidden_size)
|
|
136
|
+
if self._lstm:
|
|
137
|
+
return {
|
|
138
|
+
"rnn": {
|
|
139
|
+
"sequence_length": self._net.sequence_length,
|
|
140
|
+
# hidden states (D * num_layers, N, Hout)
|
|
141
|
+
# cell states (D * num_layers, N, Hcell)
|
|
142
|
+
"sizes": [size, size],
|
|
143
|
+
}
|
|
124
144
|
}
|
|
125
|
-
} # cell states (D ∗ num_layers, N, Hcell)
|
|
126
|
-
elif self._rnn:
|
|
127
145
|
return {
|
|
128
146
|
"rnn": {
|
|
129
147
|
"sequence_length": self._net.sequence_length,
|
|
130
|
-
|
|
148
|
+
# hidden states (D * num_layers, N, Hout)
|
|
149
|
+
"sizes": [size],
|
|
131
150
|
}
|
|
132
|
-
}
|
|
133
|
-
|
|
134
|
-
return {}
|
|
151
|
+
}
|
|
152
|
+
return {}
|
|
135
153
|
|
|
136
154
|
def tensor_to_space(
|
|
137
155
|
self, tensor: torch.Tensor, space: Union[gym.Space, gymnasium.Space], start: int = 0
|
|
@@ -196,27 +214,25 @@ class Model(torch.nn.Module):
|
|
|
196
214
|
|
|
197
215
|
def random_act(
|
|
198
216
|
self, inputs: Mapping[str, Union[torch.Tensor, Any]], role: str = ""
|
|
199
|
-
) -> Tuple[torch.Tensor,
|
|
200
|
-
"""Act randomly according to the action space
|
|
201
|
-
|
|
202
|
-
:param inputs: Model inputs. The most common keys are:
|
|
217
|
+
) -> Tuple[torch.Tensor, Mapping[str, Union[torch.Tensor, Any]]]:
|
|
218
|
+
"""Act randomly according to the action space.
|
|
203
219
|
|
|
204
|
-
|
|
205
|
-
|
|
206
|
-
|
|
207
|
-
:param role: Role play by the model (default: ``""``)
|
|
208
|
-
:type role: str, optional
|
|
220
|
+
Returns a 2-tuple ``(actions, outputs_dict)`` matching the
|
|
221
|
+
SKRL 2.x ``model.act`` contract.
|
|
222
|
+
"""
|
|
209
223
|
|
|
210
|
-
|
|
224
|
+
# SKRL 2.x agents pass the actor input under ``observations``;
|
|
225
|
+
# fall back to ``states`` when a model is invoked directly.
|
|
226
|
+
anchor = inputs.get("observations")
|
|
227
|
+
if anchor is None:
|
|
228
|
+
anchor = inputs["states"]
|
|
229
|
+
batch_size = anchor.shape[0]
|
|
211
230
|
|
|
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
231
|
# discrete action space (Discrete)
|
|
216
232
|
if issubclass(type(self.action_space), gym.spaces.Discrete) or issubclass(
|
|
217
233
|
type(self.action_space), gymnasium.spaces.Discrete
|
|
218
234
|
):
|
|
219
|
-
return torch.randint(self.action_space.n, (
|
|
235
|
+
return torch.randint(self.action_space.n, (batch_size, 1), device=self.device), {}
|
|
220
236
|
# continuous action space (Box)
|
|
221
237
|
elif issubclass(type(self.action_space), gym.spaces.Box) or issubclass(
|
|
222
238
|
type(self.action_space), gymnasium.spaces.Box
|
|
@@ -228,8 +244,7 @@ class Model(torch.nn.Module):
|
|
|
228
244
|
)
|
|
229
245
|
|
|
230
246
|
return (
|
|
231
|
-
self._random_distribution.sample(sample_shape=(
|
|
232
|
-
None,
|
|
247
|
+
self._random_distribution.sample(sample_shape=(batch_size, self.num_actions)),
|
|
233
248
|
{},
|
|
234
249
|
)
|
|
235
250
|
else:
|
|
@@ -372,22 +387,18 @@ class Model(torch.nn.Module):
|
|
|
372
387
|
"""
|
|
373
388
|
return self.forward(inputs, role)
|
|
374
389
|
|
|
375
|
-
def
|
|
376
|
-
"""
|
|
390
|
+
def enable_training_mode(self, enabled: bool = True) -> None:
|
|
391
|
+
"""Enable or disable training mode.
|
|
377
392
|
|
|
378
|
-
|
|
379
|
-
|
|
380
|
-
|
|
381
|
-
:type mode: str
|
|
393
|
+
SKRL 2.x agents call this from
|
|
394
|
+
``Agent.enable_models_training_mode`` to toggle the underlying
|
|
395
|
+
:py:meth:`torch.nn.Module.train` flag.
|
|
382
396
|
|
|
383
|
-
:
|
|
397
|
+
:param enabled: Set training mode if ``True``, evaluation mode if
|
|
398
|
+
``False`` (default: ``True``).
|
|
384
399
|
"""
|
|
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")
|
|
400
|
+
|
|
401
|
+
self.train(enabled)
|
|
391
402
|
|
|
392
403
|
def save(self, path: str, state_dict: Optional[dict] = None) -> None:
|
|
393
404
|
"""Save the model to the specified path
|