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.
Files changed (57) hide show
  1. {rlmodule-0.1.6.3/rlmodule.egg-info → rlmodule-0.2.0}/PKG-INFO +18 -6
  2. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/README.md +8 -0
  3. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/pyproject.toml +22 -5
  4. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/__init__.py +1 -1
  5. rlmodule-0.2.0/rlmodule/skrl/torch/network/__init__.py +63 -0
  6. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/log_std.py +14 -7
  7. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/model.py +62 -48
  8. rlmodule-0.2.0/rlmodule/source/network.py +1902 -0
  9. rlmodule-0.2.0/rlmodule/source/network_cfg.py +413 -0
  10. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/nvidia_utils/configclass.py +18 -3
  11. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/output_layer.py +104 -11
  12. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/output_layer_cfg.py +39 -1
  13. rlmodule-0.2.0/rlmodule/source/rlmodel.py +173 -0
  14. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/utils.py +10 -2
  15. {rlmodule-0.1.6.3 → rlmodule-0.2.0/rlmodule.egg-info}/PKG-INFO +18 -6
  16. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule.egg-info/SOURCES.txt +12 -1
  17. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule.egg-info/requires.txt +3 -3
  18. rlmodule-0.2.0/tests/test_build_model_integration.py +234 -0
  19. rlmodule-0.2.0/tests/test_cnn_backbone.py +515 -0
  20. rlmodule-0.2.0/tests/test_combined_log_std.py +42 -0
  21. rlmodule-0.2.0/tests/test_graph_network.py +856 -0
  22. rlmodule-0.2.0/tests/test_head_mlp.py +93 -0
  23. rlmodule-0.2.0/tests/test_mlp_layer_norm.py +102 -0
  24. rlmodule-0.2.0/tests/test_residual_mlp.py +97 -0
  25. rlmodule-0.2.0/tests/test_rnn_memory_and_correctness.py +444 -0
  26. rlmodule-0.2.0/tests/test_skrl_startup.py +561 -0
  27. rlmodule-0.2.0/tests/test_skrl_train_smoke.py +403 -0
  28. rlmodule-0.2.0/tests/test_squashed_gaussian.py +118 -0
  29. rlmodule-0.1.6.3/rlmodule/skrl/torch/network/__init__.py +0 -29
  30. rlmodule-0.1.6.3/rlmodule/source/network.py +0 -646
  31. rlmodule-0.1.6.3/rlmodule/source/network_cfg.py +0 -82
  32. rlmodule-0.1.6.3/rlmodule/source/rlmodel.py +0 -88
  33. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/LICENSE +0 -0
  34. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/__init__.py +0 -0
  35. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/__init__.py +0 -0
  36. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/cnn_experimental.py +0 -0
  37. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_mlp_separate_model.py +0 -0
  38. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_mlp_shared_model.py +0 -0
  39. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_rnn_forwarded_input_shared_model.py +0 -0
  40. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_rnn_separate_model.py +0 -0
  41. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_rnn_shared_model.py +0 -0
  42. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/pendulum_sac_mlp_separate_model.py +0 -0
  43. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/skrl_no_module.py +0 -0
  44. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/skrl_no_module_rnn.py +0 -0
  45. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/skrl_sac.py +0 -0
  46. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/examples/gymnasium/temp.py +0 -0
  47. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/skrl/torch/output_layer/__init__.py +0 -0
  48. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/__init__.py +0 -0
  49. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/build.py +0 -0
  50. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/log_std_cfg.py +0 -0
  51. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/nvidia_utils/__init__.py +0 -0
  52. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/nvidia_utils/dict.py +0 -0
  53. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/nvidia_utils/string.py +0 -0
  54. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule/source/rlmodel_cfg.py +0 -0
  55. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule.egg-info/dependency_links.txt +0 -0
  56. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/rlmodule.egg-info/top_level.txt +0 -0
  57. {rlmodule-0.1.6.3 → rlmodule-0.2.0}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.1
1
+ Metadata-Version: 2.4
2
2
  Name: rlmodule
3
- Version: 0.1.6.3
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.6
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.3.0
31
+ Requires-Dist: skrl>=2.1
29
32
  Provides-Extra: skrl
30
- Requires-Dist: skrl>=1.3.0; extra == "skrl"
33
+ Requires-Dist: skrl>=2.1; extra == "skrl"
31
34
  Provides-Extra: all
32
- Requires-Dist: skrl>=1.3.0; extra == "all"
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.6.3"
7
+ version = "0.2.0"
4
8
  description = "Flexible reinforcement learning models instantiators library"
5
9
  readme = "README.md"
6
- requires-python = ">=3.6"
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.3.0",
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.3.0",
43
+ "skrl>=2.1",
37
44
  ]
38
45
 
39
46
  all = [
40
- "skrl>=1.3.0",
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__ = ["build_modelBaseRLCfg", "RLModelCfg", "SeparatedRLModelCfg", "SharedRLModelCfg"]
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
- module_cfg.class_type(device, input_size, output_size, module_cfg) for module_cfg in cfg.combined_modules
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._modules)))
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._modules) != len(self._combination_constants):
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._modules)}"
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) for module_id, module in enumerate(self._modules)
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
- return {
114
- "rnn": {
115
- "sequence_length": self._net.sequence_length,
116
- "sizes": [
117
- (
118
- self._net.num_layers,
119
- self._net.num_envs,
120
- self._net.hidden_size,
121
- ), # hidden states (D ∗ num_layers, N, Hout)
122
- (self._net.num_layers, self._net.num_envs, self._net.hidden_size),
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
- "sizes": [(self._net.num_layers, self._net.num_envs, self._net.hidden_size)],
151
+ # hidden states (D * num_layers, N, Hout)
152
+ "sizes": [size],
131
153
  }
132
- } # hidden states (D ∗ num_layers, N, Hout)
133
- else:
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, None, Mapping[str, Union[torch.Tensor, Any]]]:
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
- :param inputs: Model inputs. The most common keys are:
203
-
204
- - ``"states"``: state of the environment used to make the decision
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
- :raises NotImplementedError: Unsupported action space
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, (inputs["states"].shape[0], 1), device=self.device), None, {}
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=(inputs["states"].shape[0], self.num_actions)),
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 set_mode(self, mode: str) -> None:
376
- """Set the model mode (training or evaluation)
393
+ def enable_training_mode(self, enabled: bool = True) -> None:
394
+ """Enable or disable training mode.
377
395
 
378
- :param mode: Mode: ``"train"`` for training or ``"eval"`` for evaluation.
379
- See `torch.nn.Module.train
380
- <https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.train>`_
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
- :raises ValueError: If the mode is not ``"train"`` or ``"eval"``
400
+ :param enabled: Set training mode if ``True``, evaluation mode if
401
+ ``False`` (default: ``True``).
384
402
  """
385
- if mode == "train":
386
- self.train(True)
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