1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273 | from abc import ABC, abstractmethod
from collections.abc import Callable
from typing import TYPE_CHECKING, Any, Optional
import lightning as L # noqa: N812
from lit_sb3.vec_env import VecEnv
if TYPE_CHECKING:
from lit_sb3.algorithms.base import BaseAlgorithm
class BaseCallback(L.Callback, ABC):
"""Base class for callback.
Inherits from ``lightning.Callback`` so instances can be passed directly
to a ``lightning.Trainer`` via its ``callbacks`` argument. When training
runs through :meth:`BaseAlgorithm.learn` the SB3 lifecycle methods
(``on_training_start`` / ``on_training_end``) are called manually before
``Trainer.fit``; the Lightning shims below detect that and skip to avoid
double-calls. When a callback is passed directly to a ``Trainer`` without
going through ``learn``, the shims invoke the SB3 lifecycle instead.
Args:
verbose: Verbosity level: 0 for no output, 1 for info messages, 2 for debug messages
"""
# The RL model
# Type hint as string to avoid circular import
model: "BaseAlgorithm"
def __init__(self, verbose: int = 0) -> None:
super().__init__()
# Number of time the callback was called
self.n_calls = 0 # type: int
# n_envs * n times env.step() was called
self.num_timesteps = 0 # type: int
self.verbose = verbose
self.locals: dict[str, Any] = {}
self.globals: dict[str, Any] = {}
# Sometimes, for event callback, it is useful
# to have access to the parent object
self.parent = None # type: Optional[BaseCallback]
# Guard flags to prevent double-calls when used via learn()
self._training_started = False
self._training_ended = False
@property
def training_env(self) -> VecEnv:
training_env = self.model.get_env()
assert training_env is not None, (
"`model.get_env()` returned None, you must initialize the model with an environment to use callbacks"
)
return training_env
@property
def summary_writer(self) -> Any:
return self.model.get_summary_writer()
# Type hint as string to avoid circular import
def init_callback(self, model: "BaseAlgorithm") -> None:
"""Initialize the callback by saving references to the.
RL model and the training environment for convenience.
"""
self.model = model
self._training_started = False
self._training_ended = False
self._init_callback()
def _init_callback(self) -> None: # noqa: B027
pass
def on_training_start(self, locals_: dict[str, Any], globals_: dict[str, Any]) -> None:
self._training_started = True
self._training_ended = False
# Those are reference and will be updated automatically
self.locals = locals_
self.globals = globals_
# Update num_timesteps in case training was done before
self.num_timesteps = self.model.num_timesteps
self._on_training_start()
def _on_training_start(self) -> None: # noqa: B027
pass
def on_rollout_start(self) -> None:
self._on_rollout_start()
def _on_rollout_start(self) -> None: # noqa: B027
pass
@abstractmethod
def _on_step(self) -> bool:
"""Returns.
If the callback returns False, training is aborted early.
"""
return True
def on_step(self) -> bool:
"""This method will be called by the model after each call to ``env.step()``.
For child callback (of an ``EventCallback``), this will be called
when the event is triggered.
Returns:
If the callback returns False, training is aborted early.
"""
self.n_calls += 1
self.num_timesteps = self.model.num_timesteps
return self._on_step()
def on_training_end(self) -> None:
self._training_ended = True
self._on_training_end()
def _on_training_end(self) -> None: # noqa: B027
pass
def on_rollout_end(self) -> None:
self._on_rollout_end()
def _on_rollout_end(self) -> None: # noqa: B027
pass
def update_locals(self, locals_: dict[str, Any]) -> None:
"""Update the references to the local variables.
Args:
locals_: the local variables during rollout collection
"""
self.locals.update(locals_)
self.update_child_locals(locals_)
def update_child_locals(self, locals_: dict[str, Any]) -> None: # noqa: B027
"""Update the references to the local variables on sub callbacks.
Args:
locals_: the local variables during rollout collection
"""
pass
# ------------------------------------------------------------------
# Lightning shims — bridge fit-level hooks to SB3 lifecycle
# ------------------------------------------------------------------
def on_fit_start(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:
"""Lightning hook: called at the very beginning of ``Trainer.fit``.
When invoked through :meth:`BaseAlgorithm.learn`, ``on_training_start``
has already been called manually, so this is a no-op. When a callback
is passed directly to ``Trainer(callbacks=[cb])`` this initialises the
SB3 side of the callback instead.
"""
if self._training_started:
return
self.init_callback(pl_module) # type: ignore[arg-type]
self.on_training_start(
{"total_timesteps": getattr(pl_module, "_total_timesteps", 0)},
{},
)
def on_fit_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:
"""Lightning hook: called at the very end of ``Trainer.fit``.
When invoked through :meth:`BaseAlgorithm.learn`, ``on_training_end``
has already been called by ``LightningModule.on_train_end``, so this
is a no-op. When a callback is passed directly to ``Trainer`` this
tears down the SB3 side instead.
"""
if self._training_ended:
return
self.on_training_end()
class EventCallback(BaseCallback):
"""Base class for triggering callback on event.
Args:
callback: Callback that will be called
when an event is triggered.
verbose: Verbosity level: 0 for no output, 1 for info messages, 2 for debug messages
"""
def __init__(self, callback: "BaseCallback | None" = None, verbose: int = 0) -> None:
super().__init__(verbose=verbose)
self.callback = callback
# Give access to the parent
if callback is not None:
assert self.callback is not None
self.callback.parent = self
def init_callback(self, model: "BaseAlgorithm") -> None:
super().init_callback(model)
if self.callback is not None:
self.callback.init_callback(self.model)
def _on_training_start(self) -> None:
if self.callback is not None:
self.callback.on_training_start(self.locals, self.globals)
def _on_event(self) -> bool:
if self.callback is not None:
return self.callback.on_step()
return True
def _on_step(self) -> bool:
return True
def update_child_locals(self, locals_: dict[str, Any]) -> None:
"""Update the references to the local variables.
Args:
locals_: the local variables during rollout collection
"""
if self.callback is not None:
self.callback.update_locals(locals_)
class CallbackList(BaseCallback):
"""Class for chaining callbacks.
Args:
callbacks: A list of callbacks that will be called
sequentially.
"""
def __init__(self, callbacks: list[BaseCallback]) -> None:
super().__init__()
assert isinstance(callbacks, list)
self.callbacks = callbacks
def _init_callback(self) -> None:
for callback in self.callbacks:
callback.init_callback(self.model)
# Fix for https://github.com/DLR-RM/stable-baselines3/issues/1791
# pass through the parent callback to all children
callback.parent = self.parent
def _on_training_start(self) -> None:
for callback in self.callbacks:
callback.on_training_start(self.locals, self.globals)
def _on_rollout_start(self) -> None:
for callback in self.callbacks:
callback.on_rollout_start()
def _on_step(self) -> bool:
continue_training = True
for callback in self.callbacks:
# Return False (stop training) if at least one callback returns False
continue_training = callback.on_step() and continue_training
return continue_training
def _on_rollout_end(self) -> None:
for callback in self.callbacks:
callback.on_rollout_end()
def _on_training_end(self) -> None:
for callback in self.callbacks:
callback.on_training_end()
def update_child_locals(self, locals_: dict[str, Any]) -> None:
"""Update the references to the local variables.
Args:
locals_: the local variables during rollout collection
"""
for callback in self.callbacks:
callback.update_locals(locals_)
|