cusrl 1.0.0__py3-none-any.whl

cusrl/hook/gae.py

@@ -0,0 +1,143 @@
+from typing import Any
+
+import torch
+
+from cusrl.template import ActorCritic, Hook
+from cusrl.utils import ExponentialMovingNormalizer
+
+__all__ = ["GAE"]
+
+
+@torch.jit.script
+def _generalized_advantage_estimation(
+    reward: torch.Tensor,
+    done: torch.Tensor,
+    value: torch.Tensor,
+    next_value: torch.Tensor,
+    gamma: float,
+    lamda: float,
+) -> torch.Tensor:
+    not_done = done.logical_not()
+    advantage = reward + next_value * gamma - value
+    for step in range(advantage.size(0) - 2, -1, -1):
+        advantage[step] += not_done[step] * (gamma * lamda) * advantage[step + 1]
+    return advantage
+
+
+class GAE(Hook[ActorCritic]):
+    """A hook that computes advantages and returns using Generalized Advantage Estimation (GAE).
+
+    GAE is described in:
+    "High-Dimensional Continuous Control Using Generalized Advantage Estimation",
+    https://arxiv.org/abs/1506.02438
+
+    Distinct lambda values can be enabled to individually control the bias-variance trade-offs
+    for policy and value function, described in:
+    "DNA: Proximal Policy Optimization with a Dual Network Architecture"
+    https://proceedings.neurips.cc/paper_files/paper/2022/hash/e95475f5fb8edb9075bf9e25670d4013-Abstract-Conference.html
+
+    PopArt normalization can be applied to the value function, described in:
+    "Learning values across many orders of magnitude",
+    https://proceedings.neurips.cc/paper/2016/hash/5227b6aaf294f5f027273aebf16015f2-Abstract.html
+
+    Args:
+        gamma (float, optional):
+            Discount factor for future rewards, in [0, 1). Defaults to 0.99.
+        lamda (float, optional):
+            Smoothing factor for advantage estimation, in [0, 1]. Defaults to 0.95.
+        lamda_value (float | None, optional):
+            Smoothing factor for value function calculation, in [0, 1].
+            If None, the same value as `lamda` is used. Defaults to None.
+        recompute (bool, optional):
+            If True, recompute advantages and returns after each update. Defaults to False.
+        popart_alpha (float | None, optional):
+            If not None, applies PopArt normalization to the value function with the specified
+            alpha. Defaults to None, which means no normalization is applied.
+    """
+
+    MODULES = ["value_rms"]
+    MUTABLE_ATTRS = ["gamma", "lamda"]
+
+    def __init__(
+        self,
+        gamma: float = 0.99,
+        lamda: float = 0.95,
+        lamda_value: float | None = None,
+        recompute: bool = False,
+        popart_alpha: float | None = None,
+    ):
+        if gamma < 0 or gamma >= 1:
+            raise ValueError(f"Invalid gamma value {gamma}, which should be in [0, 1).")
+        if lamda < 0 or lamda > 1:
+            raise ValueError(f"Invalid lambda value {lamda}, which should be in [0, 1].")
+        if lamda_value is not None and (lamda_value < 0 or lamda_value > 1):
+            raise ValueError(f"Invalid lambda value for value function {lamda_value}, which should be in [0, 1].")
+
+        self.gamma = gamma
+        self.lamda = lamda
+        self.lamda_value = lamda_value
+        self.recompute = recompute
+        self.popart_alpha = popart_alpha
+        self.value_rms: ExponentialMovingNormalizer | None = None
+
+    def init(self):
+        if self.popart_alpha is not None:
+            self.value_rms = self.__make_normalizer(self.agent.value_dim)
+            self.agent.critic.value_rms = self.__make_normalizer(self.agent.value_dim)
+
+    def __make_normalizer(self, num_channels: int):
+        return ExponentialMovingNormalizer(num_channels, alpha=self.popart_alpha).to(self.agent.device)
+
+    def pre_update(self, buffer):
+        if not self.recompute:
+            self._compute_advantage_and_return(buffer)
+
+    def objective(self, batch: dict[str, Any]):
+        if self.recompute:
+            self._compute_advantage_and_return(batch)
+
+    def post_update(self):
+        if self.value_rms is not None:
+            old_value_rms: ExponentialMovingNormalizer = self.agent.critic.value_rms
+            old_mean, old_std = old_value_rms.mean, old_value_rms.std
+            # Adjust value head weights and biases
+            new_mean, new_std = self.value_rms.mean, self.value_rms.std
+            value_head = self.agent.critic.value_head
+            with torch.no_grad():
+                value_head.weight.data.mul_(old_std / new_std)
+                value_head.bias.data.mul_(old_std).add_(old_mean).sub_(new_mean).div_(new_std)
+            old_value_rms.load_state_dict(self.value_rms.state_dict())
+
+    @torch.no_grad()
+    def _compute_advantage_and_return(self, data):
+        value = data["value"]
+        next_value = data["next_value"]
+        if (value_rms := self.agent.critic.value_rms) is not None:
+            value = value_rms.unnormalize(value)
+            next_value = value_rms.unnormalize(next_value)
+
+        data["advantage"] = _generalized_advantage_estimation(
+            reward=data["reward"],
+            done=data["done"],
+            value=value,
+            next_value=next_value,
+            gamma=self.gamma,
+            lamda=self.lamda,
+        )
+
+        data["return"] = value + (
+            data["advantage"]
+            if self.lamda_value is None
+            else _generalized_advantage_estimation(
+                reward=data["reward"],
+                done=data["done"],
+                value=value,
+                next_value=next_value,
+                gamma=self.gamma,
+                lamda=self.lamda_value,
+            )
+        )
+
+        if value_rms is not None:
+            self.value_rms.update(data["return"])
+            value_rms.normalize_(data["return"])

cusrl/hook/gradient.py

@@ -0,0 +1,48 @@
+from torch import nn
+
+from cusrl.template import Hook
+
+__all__ = ["GradientClipping"]
+
+
+class GradientClipping(Hook):
+    """A hook to clip gradients of model parameters before the optimizer step,
+    grouping parameters by name prefixes to handle varying gradient scales
+    across modules.
+
+    Args:
+        max_grad_norm (float | None, optional):
+            Default max norm for gradient clipping. If None, no clipping is
+            applied for the default group. Defaults to 1.0.
+        **groups (dict[str, float | None]):
+            Keyword arguments mapping parameter name prefixes to specific max
+            gradient norms.
+    """
+
+    def __init__(self, max_grad_norm: float | None = 1.0, **groups: float | None):
+        groups[""] = max_grad_norm
+        for prefix, max_grad_norm in groups.items():
+            if max_grad_norm is not None and max_grad_norm < 0:
+                raise ValueError(f"'max_grad_norm' for prefix '{prefix}' must be positive.")
+        # Sort by length of prefix (longest first for more specific matching)
+        self.groups = dict(sorted(groups.items(), key=lambda x: len(x[0]), reverse=True))
+
+    def pre_optim(self, optimizer):
+        prefixed_parameters = {prefix: [] for prefix in self.groups}
+        for param_group in optimizer.param_groups:
+            params = param_group["params"]
+            param_names = param_group.get("param_names", [""] * len(params))
+            for param, name in zip(params, param_names, strict=True):
+                prefix = self._match_prefix(name)
+                prefixed_parameters[prefix].append(param)
+        # Clip gradients for each group
+        for prefix, params in prefixed_parameters.items():
+            if params and (max_grad_norm := self.groups[prefix]) is not None:
+                nn.utils.clip_grad_norm_(params, max_grad_norm)
+
+    def _match_prefix(self, name):
+        # Find the longest matching prefix (most specific)
+        for prefix in self.groups:
+            if name == prefix or name.startswith(f"{prefix}."):
+                return prefix
+        return ""

cusrl/hook/initialization.py

@@ -0,0 +1,94 @@
+import itertools
+import math
+from typing import Literal
+
+from torch import nn
+
+from cusrl.template import ActorCritic, Hook
+
+__all__ = ["ModuleInitialization"]
+
+
+class ModuleInitialization(Hook[ActorCritic]):
+    def __init__(
+        self,
+        scale: float = math.sqrt(2),
+        scale_dist: float = math.sqrt(2) * 0.1,
+        zero_bias: bool = True,
+        conv_a: float = 0.0,
+        conv_mode: Literal["fan_in", "fan_out"] = "fan_in",
+        conv_nonlinearity: Literal["relu", "leaky_relu"] = "leaky_relu",
+        init_actor: bool = True,
+        init_critic: bool = True,
+        distribution_std: float | None = None,
+    ):
+        self.scale = scale
+        self.scale_dist = scale_dist
+        self.zero_bias = zero_bias
+        self.conv_a = conv_a
+        self.conv_mode = conv_mode
+        self.conv_nonlinearity = conv_nonlinearity
+        self.init_actor = init_actor
+        self.init_critic = init_critic
+        self.distribution_std = distribution_std
+
+    def init(self):
+        if self.init_actor:
+            for module in itertools.chain(self.agent.actor.modules()):
+                self._init_module(module, self.scale, self.zero_bias)
+            if self.scale_dist != self.scale:
+                self._init_linear(self.agent.actor.distribution.mean_head, self.scale_dist, self.zero_bias)
+        if self.distribution_std is not None:
+            self.agent.actor.set_distribution_std(self.distribution_std)
+        if self.init_critic:
+            for module in itertools.chain(self.agent.critic.modules()):
+                self._init_module(module, self.scale, self.zero_bias)
+
+    def _init_module(self, module: nn.Module, scale: float, zero_bias: bool):
+        if isinstance(module, nn.Linear):
+            self._init_linear(module, scale, zero_bias)
+        elif isinstance(module, (nn.LSTM, nn.GRU)):
+            self._init_gru_lstm(module, scale, zero_bias)
+        elif isinstance(module, nn.MultiheadAttention):
+            self._init_mha(module, scale, zero_bias)
+        elif isinstance(module, nn.Conv2d):
+            self._init_conv2d(module, zero_bias)
+
+    def _init_linear(self, module: nn.Linear, scale: float, zero_bias: bool):
+        nn.init.orthogonal_(module.weight, gain=scale)
+        if zero_bias and module.bias is not None:
+            nn.init.zeros_(module.bias)
+
+    def _init_gru_lstm(self, module: nn.GRU | nn.LSTM, scale: float, zero_bias: bool):
+        for i in range(module.num_layers):
+            nn.init.orthogonal_(getattr(module, f"weight_hh_l{i}"), gain=scale)
+            nn.init.orthogonal_(getattr(module, f"weight_ih_l{i}"), gain=scale)
+            if zero_bias and getattr(module, f"bias_hh_l{i}") is not None:
+                nn.init.zeros_(getattr(module, f"bias_hh_l{i}"))
+                nn.init.zeros_(getattr(module, f"bias_ih_l{i}"))
+
+    def _init_mha(self, module: nn.MultiheadAttention, scale: float, zero_bias: bool):
+        if module.in_proj_weight is not None:
+            nn.init.orthogonal_(module.in_proj_weight, gain=scale)
+        else:
+            nn.init.orthogonal_(module.q_proj_weight, gain=scale)
+            nn.init.orthogonal_(module.k_proj_weight, gain=scale)
+            nn.init.orthogonal_(module.v_proj_weight, gain=scale)
+
+        if zero_bias:
+            if module.in_proj_bias is not None:
+                nn.init.zeros_(module.in_proj_bias)
+            if module.bias_k is not None:
+                nn.init.zeros_(module.bias_k)
+            if module.bias_v is not None:
+                nn.init.zeros_(module.bias_v)
+
+    def _init_conv2d(self, module: nn.Conv2d, zero_bias: bool):
+        nn.init.kaiming_normal_(
+            module.weight,
+            a=self.conv_a,
+            mode=self.conv_mode,
+            nonlinearity=self.conv_nonlinearity,
+        )
+        if zero_bias and module.bias is not None:
+            nn.init.zeros_(module.bias)

cusrl/hook/lr_schedule.py

@@ -0,0 +1,178 @@
+import math
+
+import torch
+
+from cusrl.template import ActorCritic, Hook
+from cusrl.utils import distributed
+
+__all__ = ["AdaptiveLRSchedule", "MiniBatchWiseLRSchedule", "ThresholdLRSchedule"]
+
+
+class KLDivergenceBasedLRSchedule(Hook[ActorCritic]):
+    MUTABLE_ATTRS = ["desired_kl_divergence"]
+
+    def __init__(
+        self,
+        desired_kl_divergence: float = 0.01,
+        scale_all_params: bool = False,
+    ):
+        if desired_kl_divergence <= 0:
+            raise ValueError("'desired_kl_divergence' must be positive.")
+
+        self.desired_kl_divergence = desired_kl_divergence
+        self.scale_all_params = scale_all_params
+        self._lr_scale = 1.0
+
+    def post_update(self):
+        kl_divergence = self.agent.metrics["kl_divergence"].mean.clone()
+        distributed.reduce_mean_(kl_divergence)
+        scale = self._compute_scale(kl_divergence.item())
+        self._scale_lr_of_parameters(scale)
+        self.agent.record(lr_scale=self._lr_scale)
+
+    def _compute_scale(self, kl_divergence: float) -> float | None:
+        raise NotImplementedError
+
+    def _scale_lr_of_parameters(self, scale: float | None):
+        if scale is None or scale == 1.0:
+            return
+
+        self._lr_scale *= scale
+        for param_group in self.agent.optimizer.param_groups:
+            if self.scale_all_params or any(name.startswith("actor.") for name in param_group["param_names"]):
+                param_group["lr"] *= scale
+
+    def state_dict(self):
+        return {"lr_scale": self._lr_scale}
+
+    def load_state_dict(self, state_dict):
+        self._lr_scale = state_dict["lr_scale"]
+
+
+class ThresholdLRSchedule(KLDivergenceBasedLRSchedule):
+    """Adjusts the learning rate based on thresholded KL divergence.
+
+    Args:
+        desired_kl_divergence (float, optional):
+            Target KL divergence to maintain. Defaults to 0.01.
+        threshold (float, optional):
+            Ratio threshold (>1) for deciding when to adjust. Defaults to 1.2.
+        scale_factor (float, optional):
+            Multiplicative factor (>1) for scaling the LR. Defaults to 1.1.
+        scale_all_params (bool, optional):
+            If True, scales all optimizer parameter groups; otherwise only
+            scales actor parameter groups. Defaults to False.
+    """
+
+    MUTABLE_ATTRS = ["desired_kl_divergence"]
+
+    def __init__(
+        self,
+        desired_kl_divergence: float = 0.01,
+        threshold: float = 1.2,
+        scale_factor: float = 1.1,
+        scale_all_params: bool = False,
+    ):
+        super().__init__(desired_kl_divergence, scale_all_params)
+        if threshold <= 1:
+            raise ValueError("'threshold' must be greater than 1.")
+        if scale_factor <= 1:
+            raise ValueError("'scale_factor' must be greater than 1.")
+
+        self.threshold = threshold
+        self.scale_factor = scale_factor
+
+    def _compute_scale(self, kl_divergence: float):
+        if kl_divergence > self.desired_kl_divergence * self.threshold:
+            return 1 / self.scale_factor
+        if kl_divergence < self.desired_kl_divergence / self.threshold:
+            return self.scale_factor
+        return None
+
+
+class AdaptiveLRSchedule(KLDivergenceBasedLRSchedule):
+    """Adaptively adjusts the learning rate based on accumulated KL divergence error.
+
+    Args:
+        desired_kl_divergence (float, optional):
+            Target KL divergence to maintain. Defaults to 0.01.
+        threshold (float, optional):
+            Positive threshold for accumulated log-error before scaling. Defaults to 1.0.
+        scale_factor (float, optional):
+            Positive coefficient controlling adjustment magnitude. Defaults to 0.2.
+        scale_all_params (bool, optional):
+            If True, scales all optimizer parameter groups; otherwise only scales actor
+            parameter groups. Defaults to False.
+    """
+
+    MUTABLE_ATTRS = ["desired_kl_divergence"]
+
+    def __init__(
+        self,
+        desired_kl_divergence: float = 0.01,
+        threshold: float = 1.0,
+        scale_factor: float = 0.2,
+        scale_all_params: bool = False,
+    ):
+        super().__init__(desired_kl_divergence, scale_all_params)
+        if threshold <= 0:
+            raise ValueError("'threshold' must be positive.")
+        if scale_factor <= 0:
+            raise ValueError("'scale_factor' must be positive.")
+
+        self.threshold = threshold
+        self.scale_factor = scale_factor
+        self.accumulated_log_error = 0.0
+        self.count = 0
+
+    def _compute_scale(self, kl_divergence: float):
+        self.accumulated_log_error += math.log(kl_divergence / self.desired_kl_divergence)
+        self.count += 1
+        if self.threshold > self.accumulated_log_error > -self.threshold:
+            return None
+        average_log_error = self.accumulated_log_error / self.count
+        scale = math.exp(-min(max(average_log_error, -1.0), 1.0) * self.scale_factor)
+        self.accumulated_log_error = 0.0
+        self.count = 0
+        return scale
+
+
+class MiniBatchWiseLRSchedule(ThresholdLRSchedule):
+    """Applies a threshold-based LR schedule on a per-mini-batch KL divergence.
+    Modified from (RSL-RL)[https://github.com/leggedrobotics/rsl_rl].
+
+    Args:
+        desired_kl_divergence (float, optional):
+            Target KL divergence per mini-batch. Defaults to 0.01.
+        threshold (float, optional):
+            Ratio threshold for deciding scaling per batch. Defaults to 2.0.
+        scale_factor (float, optional):
+            Multiplicative factor for scaling the LR. Defaults to 1.5.
+    """
+
+    def __init__(
+        self,
+        desired_kl_divergence: float = 0.01,
+        threshold: float = 2.0,
+        scale_factor: float = 1.5,
+    ):
+        super().__init__(
+            desired_kl_divergence,
+            threshold=threshold,
+            scale_factor=scale_factor,
+            scale_all_params=True,
+        )
+
+    def post_init(self):
+        self.agent.hook["OnPolicyPreparation"].calculate_kl_divergence = True
+
+    def post_update(self):
+        pass
+
+    def objective(self, batch):
+        with torch.inference_mode():
+            kl_divergence = batch["kl_divergence"].mean()
+        distributed.reduce_mean_(kl_divergence)
+        scale = self._compute_scale(kl_divergence.item())
+        self._scale_lr_of_parameters(scale)
+        self.agent.record(lr_scale=self._lr_scale)

cusrl/hook/normalization.py

@@ -0,0 +1,194 @@
+from collections.abc import Sequence
+from typing import Any
+
+import numpy as np
+import torch
+
+from cusrl.hook.symmetry import SymmetryDef
+from cusrl.template import ActorCritic, Hook
+from cusrl.utils import RunningMeanStd, mean_var_count
+from cusrl.utils.export import ExportSpec
+from cusrl.utils.typing import Slice
+
+__all__ = ["ObservationNormalization"]
+
+
+class ObservationNormalization(Hook[ActorCritic]):
+    """Normalizes observations and states using a running mean and standard deviation.
+
+    This hook maintains a running estimate of the mean and standard deviation for
+    observations and, if available, states. It intercepts transitions during data
+    collection to normalize the `observation`, `state`, `next_observation`, and
+    `next_state` fields. The original, unnormalized values are preserved under
+    keys with an "original_" prefix.
+
+    The running statistics are updated with new data from each step, unless the
+    agent is in inference mode or the `frozen` attribute is set to `True`.
+
+    The hook also handles scenarios where the observation is a subset of the
+    state or where there is symmetry in the observations or states. It correctly
+    synchronizes statistics across distributed processes. During model export,
+    it attaches forward pre-hooks to the actor and critic models to ensure that
+    inputs are automatically normalized.
+
+    Args:
+        max_count:
+            The maximum count for the running statistics to prevent numerical
+            overflow. Defaults to None.
+        defer_synchronization:
+            If True, synchronization of running statistics in a distributed setting
+            is deferred until the end of a rollout. This can improve performance
+            by reducing the frequency of synchronization. Defaults to False.
+    """
+
+    observation_rms: RunningMeanStd
+    state_rms: RunningMeanStd | None = None
+    _mirror_observation: SymmetryDef | None = None
+    _mirror_state: SymmetryDef | None = None
+    _observation_is_subset_of_state: Slice | torch.Tensor | None = None
+
+    MODULES = ["observation_rms", "state_rms"]
+    MUTABLE_ATTRS = ["frozen"]
+
+    def __init__(self, max_count: int | None = None, defer_synchronization: bool = False):
+        if max_count is not None and max_count <= 0:
+            raise ValueError("'max_count' must be positive or None.")
+        self.max_count = max_count
+        self.frozen: bool = False
+        self.defer_synchronization = defer_synchronization
+        self._last_done: torch.Tensor | None = None
+
+    def init(self):
+        # Retrieve and normalize the subset index spec
+        env_spec = self.agent.environment_spec
+        observation_is_subset_of_state = env_spec.observation_is_subset_of_state
+        if observation_is_subset_of_state is not None:
+            if not self.agent.has_state:
+                raise ValueError("'observation_is_subset_of_state' is set but state is not defined.")
+            # Convert numpy or list indices to a tensor for consistent indexing
+            if isinstance(observation_is_subset_of_state, (np.ndarray, Sequence)):
+                observation_is_subset_of_state = self.agent.to_tensor(np.asarray(observation_is_subset_of_state))
+        self._observation_is_subset_of_state = observation_is_subset_of_state
+
+        if self._observation_is_subset_of_state is not None:
+            self.observation_rms = self.__make_rms(self.agent.observation_dim)
+        else:
+            self.observation_rms = self.__make_rms(
+                self.agent.observation_dim, self.max_count, env_spec.observation_stat_groups
+            )
+        if self.agent.has_state:
+            self.state_rms = self.__make_rms(self.agent.state_dim, self.max_count, env_spec.state_stat_groups)
+        self._mirror_observation = env_spec.mirror_observation
+        self._mirror_state = env_spec.mirror_state
+
+    def pre_act(self, transition: dict):
+        observation, state = transition["observation"], transition.get("state")
+        if self._last_done is None or not self.agent.environment_spec.final_state_is_missing:
+            self.__update_rms(observation, state, self._last_done)
+
+        transition["original_observation"] = observation
+        transition["observation"] = self.observation_rms.normalize(observation)
+        if state is not None:
+            transition["original_state"] = state
+            transition["state"] = self.state_rms.normalize(state)
+
+    def post_step(self, transition: dict):
+        next_observation, next_state = transition["next_observation"], transition.get("next_state")
+        self.__update_rms(next_observation, next_state)
+        self._last_done = transition["done"].squeeze(-1)
+
+        transition["original_next_observation"] = next_observation
+        transition["next_observation"] = self.observation_rms.normalize(next_observation)
+        if next_state is not None:
+            transition["original_next_state"] = next_state
+            transition["next_state"] = self.state_rms.normalize(next_state)
+
+    def __make_rms(
+        self,
+        num_channels: int,
+        max_count: int | None = None,
+        stat_groups: tuple[tuple[int, int], ...] = (),
+    ):
+        normalizer = RunningMeanStd(num_channels, max_count=max_count).to(self.agent.device)
+        for group in stat_groups:
+            normalizer.register_stat_group(*group)
+        return normalizer
+
+    def __update_rms(
+        self,
+        observation: torch.Tensor,
+        state: torch.Tensor | None,
+        indices: torch.Tensor | None = None,
+    ):
+        if self.agent.inference_mode or self.frozen:
+            # Do not update the statistics during inference or if frozen
+            return
+
+        if state is not None:
+            self.__update_rms_impl(state, self.state_rms, self._mirror_state, indices)
+        if self._observation_is_subset_of_state is not None:
+            self.observation_rms.mean.copy_(self.state_rms.mean[self._observation_is_subset_of_state])
+            self.observation_rms.var.copy_(self.state_rms.var[self._observation_is_subset_of_state])
+            self.observation_rms.std.copy_(self.state_rms.std[self._observation_is_subset_of_state])
+            self.observation_rms.count = self.state_rms.count
+        else:
+            self.__update_rms_impl(observation, self.observation_rms, self._mirror_observation, indices)
+
+    def __update_rms_impl(
+        self,
+        observation: torch.Tensor,
+        rms: RunningMeanStd,
+        mirror: SymmetryDef | None = None,
+        indices: torch.Tensor | None = None,
+    ):
+        if indices is not None:
+            observation = observation[indices]
+        mean, var, count = mean_var_count(observation)
+        if mirror is not None:
+            mirrored_mean = mirror(mean)
+            mirrored_var = abs(mirror(var))
+            var = (var + mirrored_var) / 2 + (mean - mirrored_mean) ** 2 / 4
+            mean = (mean + mirrored_mean) / 2
+        rms.update_from_stats(mean, var, count, synchronize=not self.defer_synchronization)
+
+    def pre_update(self, buffer):
+        if self.defer_synchronization:
+            if self.state_rms is not None:
+                self.state_rms.synchronize()
+            if self._observation_is_subset_of_state is not None:
+                self.observation_rms.mean.copy_(self.state_rms.mean[self._observation_is_subset_of_state])
+                self.observation_rms.var.copy_(self.state_rms.var[self._observation_is_subset_of_state])
+                self.observation_rms.std.copy_(self.state_rms.std[self._observation_is_subset_of_state])
+                self.observation_rms.count = self.state_rms.count
+            else:
+                self.observation_rms.synchronize()
+
+    def export(self, export_data: dict[str, ExportSpec]):
+        export_data["actor"].module.register_forward_pre_hook(
+            self.__normalize_observation_forward_pre_hook, with_kwargs=True
+        )
+        if "critic" in export_data:
+            export_data["critic"].module.register_forward_pre_hook(
+                (
+                    self.__normalize_state_forward_pre_hook
+                    if self.agent.has_state
+                    else self.__normalize_observation_forward_pre_hook
+                ),
+                with_kwargs=True,
+            )
+
+    @torch.no_grad()
+    def __normalize_observation_forward_pre_hook(self, module, args: tuple, kwargs: dict[str, Any]):
+        if "observation" in kwargs:
+            kwargs["observation"] = self.observation_rms.normalize(kwargs["observation"])
+            return args, kwargs
+        normalized_observation = self.observation_rms.normalize(args[0])
+        return (normalized_observation, *args[1:]), kwargs
+
+    @torch.no_grad()
+    def __normalize_state_forward_pre_hook(self, module, args: tuple, kwargs: dict[str, Any]):
+        if "state" in kwargs:
+            kwargs["state"] = self.state_rms.normalize(kwargs["state"])
+            return args, kwargs
+        normalized_state = self.state_rms.normalize(args[0])
+        return (normalized_state, *args[1:]), kwargs