rlmodule 0.1.6.2__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.
Files changed (56) hide show
  1. {rlmodule-0.1.6.2/rlmodule.egg-info → rlmodule-0.1.7}/PKG-INFO +19 -7
  2. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/README.md +9 -1
  3. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/pyproject.toml +22 -5
  4. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule/skrl/torch/__init__.py +1 -1
  5. rlmodule-0.1.7/rlmodule/skrl/torch/examples/gymnasium/cnn_experimental.py +152 -0
  6. {rlmodule-0.1.6.2/rlmodule/skrl/torch → rlmodule-0.1.7/rlmodule/skrl/torch/examples}/gymnasium/pendulum_ppo_mlp_separate_model.py +0 -3
  7. rlmodule-0.1.7/rlmodule/skrl/torch/network/__init__.py +59 -0
  8. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule/source/build.py +4 -4
  9. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule/source/model.py +58 -47
  10. rlmodule-0.1.7/rlmodule/source/network.py +1747 -0
  11. rlmodule-0.1.7/rlmodule/source/network_cfg.py +384 -0
  12. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule/source/nvidia_utils/configclass.py +18 -3
  13. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule/source/output_layer.py +53 -17
  14. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule/source/output_layer_cfg.py +21 -1
  15. rlmodule-0.1.7/rlmodule/source/rlmodel.py +170 -0
  16. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule/source/utils.py +38 -1
  17. {rlmodule-0.1.6.2 → rlmodule-0.1.7/rlmodule.egg-info}/PKG-INFO +19 -7
  18. rlmodule-0.1.7/rlmodule.egg-info/SOURCES.txt +49 -0
  19. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule.egg-info/requires.txt +3 -3
  20. rlmodule-0.1.7/tests/test_build_model_integration.py +234 -0
  21. rlmodule-0.1.7/tests/test_cnn_backbone.py +311 -0
  22. rlmodule-0.1.7/tests/test_graph_network.py +787 -0
  23. rlmodule-0.1.7/tests/test_head_mlp.py +93 -0
  24. rlmodule-0.1.7/tests/test_mlp_layer_norm.py +102 -0
  25. rlmodule-0.1.7/tests/test_residual_mlp.py +97 -0
  26. rlmodule-0.1.7/tests/test_rnn_memory_and_correctness.py +444 -0
  27. rlmodule-0.1.7/tests/test_skrl_startup.py +561 -0
  28. rlmodule-0.1.7/tests/test_skrl_train_smoke.py +403 -0
  29. rlmodule-0.1.6.2/rlmodule/skrl/torch/network/__init__.py +0 -20
  30. rlmodule-0.1.6.2/rlmodule/source/network.py +0 -697
  31. rlmodule-0.1.6.2/rlmodule/source/network_cfg.py +0 -67
  32. rlmodule-0.1.6.2/rlmodule/source/rlmodel.py +0 -88
  33. rlmodule-0.1.6.2/rlmodule.egg-info/SOURCES.txt +0 -39
  34. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/LICENSE +0 -0
  35. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule/__init__.py +0 -0
  36. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule/skrl/__init__.py +0 -0
  37. {rlmodule-0.1.6.2/rlmodule/skrl/torch → rlmodule-0.1.7/rlmodule/skrl/torch/examples}/gymnasium/pendulum_ppo_mlp_shared_model.py +0 -0
  38. {rlmodule-0.1.6.2/rlmodule/skrl/torch → rlmodule-0.1.7/rlmodule/skrl/torch/examples}/gymnasium/pendulum_ppo_rnn_forwarded_input_shared_model.py +0 -0
  39. {rlmodule-0.1.6.2/rlmodule/skrl/torch → rlmodule-0.1.7/rlmodule/skrl/torch/examples}/gymnasium/pendulum_ppo_rnn_separate_model.py +0 -0
  40. {rlmodule-0.1.6.2/rlmodule/skrl/torch → rlmodule-0.1.7/rlmodule/skrl/torch/examples}/gymnasium/pendulum_ppo_rnn_shared_model.py +0 -0
  41. {rlmodule-0.1.6.2/rlmodule/skrl/torch → rlmodule-0.1.7/rlmodule/skrl/torch/examples}/gymnasium/pendulum_sac_mlp_separate_model.py +0 -0
  42. {rlmodule-0.1.6.2/rlmodule/skrl/torch → rlmodule-0.1.7/rlmodule/skrl/torch/examples}/gymnasium/skrl_no_module.py +0 -0
  43. {rlmodule-0.1.6.2/rlmodule/skrl/torch → rlmodule-0.1.7/rlmodule/skrl/torch/examples}/gymnasium/skrl_no_module_rnn.py +0 -0
  44. {rlmodule-0.1.6.2/rlmodule/skrl/torch → rlmodule-0.1.7/rlmodule/skrl/torch/examples}/gymnasium/skrl_sac.py +0 -0
  45. {rlmodule-0.1.6.2/rlmodule/skrl/torch → rlmodule-0.1.7/rlmodule/skrl/torch/examples}/gymnasium/temp.py +0 -0
  46. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule/skrl/torch/output_layer/__init__.py +0 -0
  47. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule/source/__init__.py +0 -0
  48. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule/source/log_std.py +0 -0
  49. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule/source/log_std_cfg.py +0 -0
  50. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule/source/nvidia_utils/__init__.py +0 -0
  51. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule/source/nvidia_utils/dict.py +0 -0
  52. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule/source/nvidia_utils/string.py +0 -0
  53. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule/source/rlmodel_cfg.py +0 -0
  54. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule.egg-info/dependency_links.txt +0 -0
  55. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/rlmodule.egg-info/top_level.txt +0 -0
  56. {rlmodule-0.1.6.2 → rlmodule-0.1.7}/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.2
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.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,9 +81,17 @@ 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
- python3 rlmodule/skrl/torch/gymnasium/pendulum_ppo_mlp_separate_model.py
94
+ python3 rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_mlp_separate_model.py
83
95
  ```
84
96
 
85
97
  Optional: observe run results in Tensorboard
@@ -44,9 +44,17 @@ 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
- python3 rlmodule/skrl/torch/gymnasium/pendulum_ppo_mlp_separate_model.py
57
+ python3 rlmodule/skrl/torch/examples/gymnasium/pendulum_ppo_mlp_separate_model.py
50
58
  ```
51
59
 
52
60
  Optional: observe run results in Tensorboard
@@ -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.2"
7
+ version = "0.1.7"
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,152 @@
1
+ from typing import Tuple
2
+
3
+ import os
4
+ from datetime import datetime
5
+ import gymnasium as gym
6
+
7
+ # import the skrl components to build the RL system
8
+ from skrl.agents.torch.ppo import PPO, PPO_DEFAULT_CONFIG, PPO_RNN
9
+ from skrl.envs.wrappers.torch.gymnasium_envs import GymnasiumWrapper
10
+ from skrl.memories.torch import RandomMemory
11
+ from skrl.resources.preprocessors.torch import RunningStandardScaler
12
+ from skrl.resources.schedulers.torch import KLAdaptiveRL
13
+ from skrl.trainers.torch import SequentialTrainer
14
+ from skrl.utils import set_seed
15
+
16
+ import numpy as np
17
+ import torch.nn as nn
18
+
19
+ from rlmodule.skrl.torch import SharedRLModelCfg, build_model
20
+ from rlmodule.skrl.torch.network import CnnCfg
21
+ from rlmodule.skrl.torch.output_layer import DeterministicLayerCfg, GaussianLayerCfg, ParameterLogStdCfg
22
+
23
+
24
+ # CNN is in an experimental phase. Some combination of inputs may not be yet supported
25
+
26
+
27
+ class DummyEnv(gym.Env):
28
+ """Dummy env producing random observations of given shape."""
29
+
30
+ def __init__(self):
31
+ super(DummyEnv, self).__init__()
32
+ # Define observation space:
33
+ self.observation_space = gym.spaces.Box(low=0, high=255, shape=(169,), dtype=np.uint8)
34
+ # Define action space: Continuous values between -1.0 and 1.0, with 6 action dimensions
35
+ self.action_space = gym.spaces.Box(low=-1.0, high=1.0, shape=(6,), dtype=np.float32)
36
+
37
+ def reset(self) -> Tuple[np.ndarray, dict]:
38
+ # Reset environment and return initial random observation
39
+ observation = self._get_random_observation()
40
+ return observation, {}
41
+
42
+ def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, dict]:
43
+ # Create a new random observation
44
+ observation = self._get_random_observation()
45
+ reward = 0.0 # No meaningful reward in this dummy environment
46
+ terminated = False # This dummy environment never ends
47
+ truncated = False
48
+ info = {}
49
+ return observation, reward, terminated, truncated, info
50
+
51
+ def _get_random_observation(self) -> np.ndarray:
52
+ # Generate a random observation with values between 0 and 255
53
+ return np.random.randint(0, 256, 169, dtype=np.uint8)
54
+
55
+
56
+ def get_model(env):
57
+ """Instantiate the agent's models (function approximators)."""
58
+
59
+ net_cfg = CnnCfg(
60
+ input_states=(13, 13, 1),
61
+ layers=[
62
+ nn.Conv2d(in_channels=1, out_channels=64, kernel_size=3, stride=2),
63
+ nn.ReLU(),
64
+ nn.Flatten(),
65
+ ],
66
+ )
67
+
68
+ model = build_model(
69
+ SharedRLModelCfg(
70
+ network=net_cfg,
71
+ device=device,
72
+ policy_output_layer=GaussianLayerCfg(
73
+ output_size=env.action_space,
74
+ log_std=ParameterLogStdCfg(
75
+ min_log_std=-1.2,
76
+ max_log_std=2,
77
+ initial_log_std=0.0,
78
+ ),
79
+ ),
80
+ value_output_layer=DeterministicLayerCfg(),
81
+ )
82
+ )
83
+
84
+ print(model)
85
+ return {"policy": model, "value": model}
86
+
87
+
88
+ # set seed for reproducibility
89
+ seed = 42
90
+ set_seed(seed)
91
+
92
+ env = DummyEnv()
93
+ # env = gym.make_vec("ALE/Pong-v5", num_envs=4, vectorization_mode="sync")
94
+ env = GymnasiumWrapper(env)
95
+
96
+ device = env.device
97
+
98
+ # instantiate a memory as rollout buffer (any memory can be used for this)
99
+ memory = RandomMemory(memory_size=1024, num_envs=env.num_envs, device=device)
100
+
101
+ models = get_model(env)
102
+
103
+ # configure and instantiate the agent (visit its documentation to see all the options)
104
+ # https://skrl.readthedocs.io/en/latest/api/agents/ppo.html#configuration-and-hyperparameters
105
+ cfg = PPO_DEFAULT_CONFIG.copy()
106
+ cfg["rollouts"] = 1024 # memory_size
107
+ cfg["learning_epochs"] = 10
108
+ cfg["mini_batches"] = 32
109
+ cfg["discount_factor"] = 0.9
110
+ cfg["lambda"] = 0.95
111
+ cfg["learning_rate"] = 1e-3
112
+ cfg["learning_rate_scheduler"] = KLAdaptiveRL
113
+ cfg["learning_rate_scheduler_kwargs"] = {"kl_threshold": 0.008}
114
+ cfg["grad_norm_clip"] = 0.5
115
+ cfg["ratio_clip"] = 0.2
116
+ cfg["value_clip"] = 0.2
117
+ cfg["clip_predicted_values"] = False
118
+ cfg["entropy_loss_scale"] = 0.0
119
+ cfg["value_loss_scale"] = 0.5
120
+ cfg["kl_threshold"] = 0
121
+ cfg["state_preprocessor"] = RunningStandardScaler
122
+ cfg["state_preprocessor_kwargs"] = {"size": env.observation_space, "device": device}
123
+ cfg["value_preprocessor"] = RunningStandardScaler
124
+ cfg["value_preprocessor_kwargs"] = {"size": 1, "device": device}
125
+ # logging to TensorBoard and write checkpoints (in timesteps)
126
+ cfg["experiment"]["write_interval"] = 500
127
+ cfg["experiment"]["checkpoint_interval"] = 5000
128
+ cfg["experiment"]["directory"] = "runs/skrl/torch/"
129
+ cfg["experiment"]["experiment_name"] = (
130
+ os.path.splitext(os.path.basename(__file__))[0] + "_" + datetime.now().strftime("%Y-%m-%d_%H:%M:%S")
131
+ )
132
+
133
+ params = {
134
+ "models": models,
135
+ "memory": memory,
136
+ "cfg": cfg,
137
+ "observation_space": env.observation_space,
138
+ "action_space": env.action_space,
139
+ "device": device,
140
+ }
141
+
142
+ if models["policy"].is_rnn:
143
+ agent = PPO_RNN(**params)
144
+ else:
145
+ agent = PPO(**params)
146
+
147
+ # configure and instantiate the RL trainer
148
+ cfg_trainer = {"timesteps": 100000, "headless": True}
149
+ trainer = SequentialTrainer(cfg=cfg_trainer, env=env, agents=[agent])
150
+
151
+ # start training
152
+ trainer.train()
@@ -27,9 +27,6 @@ def get_model(env):
27
27
  activation=nn.ReLU,
28
28
  )
29
29
 
30
- print(env.action_space)
31
- print(env.observation_space)
32
-
33
30
  policy_model = build_model(
34
31
  RLModelCfg(
35
32
  network=net_cfg,
@@ -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
+ )
@@ -10,9 +10,9 @@ def build_model(cfg: BaseRLCfg):
10
10
  actions_as_input = hasattr(network_cfg, "input_actions") and network_cfg.input_actions is not None
11
11
 
12
12
  # process input size
13
- network_cfg.input_states = get_space_size(network_cfg.input_states)
13
+ network_cfg.input_size = get_space_size(network_cfg.input_states)
14
14
  if actions_as_input:
15
- network_cfg.input_states += get_space_size(network_cfg.input_actions)
15
+ network_cfg.input_size += get_space_size(network_cfg.input_actions)
16
16
 
17
17
  # build base network of function approximator
18
18
  net = network_cfg.module(network_cfg)
@@ -21,10 +21,10 @@ def build_model(cfg: BaseRLCfg):
21
21
  net.input_actions = True
22
22
 
23
23
  # get output size to be used as output layer input
24
- network_output_size = get_output_size(net, network_cfg.input_states)
24
+ network_output_size = get_output_size(net, network_cfg.input_size)
25
25
 
26
26
  def build_output_layer(layer_cfg):
27
- return layer_cfg.class_type(device=cfg.device, input_states=network_output_size, cfg=layer_cfg)
27
+ return layer_cfg.class_type(device=cfg.device, input_size=network_output_size, cfg=layer_cfg)
28
28
 
29
29
  # build function approximator
30
30
  if isinstance(cfg, RLModelCfg):
@@ -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
- 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
- ],
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
- "sizes": [(self._net.num_layers, self._net.num_envs, self._net.hidden_size)],
148
+ # hidden states (D * num_layers, N, Hout)
149
+ "sizes": [size],
131
150
  }
132
- } # hidden states (D ∗ num_layers, N, Hout)
133
- else:
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, None, Mapping[str, Union[torch.Tensor, Any]]]:
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
- - ``"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
220
+ Returns a 2-tuple ``(actions, outputs_dict)`` matching the
221
+ SKRL 2.x ``model.act`` contract.
222
+ """
209
223
 
210
- :raises NotImplementedError: Unsupported action space
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, (inputs["states"].shape[0], 1), device=self.device), None, {}
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=(inputs["states"].shape[0], self.num_actions)),
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 set_mode(self, mode: str) -> None:
376
- """Set the model mode (training or evaluation)
390
+ def enable_training_mode(self, enabled: bool = True) -> None:
391
+ """Enable or disable training mode.
377
392
 
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
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
- :raises ValueError: If the mode is not ``"train"`` or ``"eval"``
397
+ :param enabled: Set training mode if ``True``, evaluation mode if
398
+ ``False`` (default: ``True``).
384
399
  """
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")
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