reference/generated/torchrl.modules.OneStepModel

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

OneStepModel

class torchrl.modules.OneStepModel(network: Module, action_dim: int, *, low: float | Tensor = -1.0, high: float | Tensor = 1.0)

Tensor-only network for one-step flow distillation.

Parameters:

• network (nn.Module) – maps concatenated observation and Gaussian noise directly to an action, without an output activation.

• action_dim (int) – number of action coordinates.

Keyword Arguments:

• low (float or Tensor, optional) – lower action bound, broadcast over actions. Must be strictly less than high. Defaults to -1.

• high (float or Tensor, optional) – upper action bound, broadcast over actions. Defaults to 1.

OneStepPolicy provides the TensorDict interface. Outputs are clipped to [low, high]. Both models sample Gaussian noise even under deterministic exploration; pass explicit noise for repeatability.

forward(observation: Tensor, noise: Tensor | None = None, *, clamp: bool = True) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.