Several implementations of TruncatedNormal?

adrische · December 1, 2025, 4:23pm

Hi,

There seem to be two implementations of TruncatedNormal,

torchrl/modules/distributions/truncated_normal.py

main

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
# from https://github.com/toshas/torch_truncnorm
from __future__ import annotations

import math
from numbers import Number

import torch
from torch.distributions import constraints, Distribution
from torch.distributions.utils import broadcast_all

CONST_SQRT_2 = math.sqrt(2)
CONST_INV_SQRT_2PI = 1 / math.sqrt(2 * math.pi)
CONST_INV_SQRT_2 = 1 / math.sqrt(2)
CONST_LOG_INV_SQRT_2PI = math.log(CONST_INV_SQRT_2PI)
CONST_LOG_SQRT_2PI_E = 0.5 * math.log(2 * math.pi * math.e)

This file has been truncated. show original

github.com/pytorch/rl

torchrl/modules/distributions/continuous.py

main

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from __future__ import annotations

import weakref
from collections.abc import Sequence
from numbers import Number

import numpy as np
import torch
from packaging import version
from torch import distributions as D, nn
from torch.distributions import constraints
from torch.distributions.transforms import _InverseTransform

from torchrl._utils import safe_is_current_stream_capturing
from torchrl.modules.distributions.truncated_normal import (
    TruncatedNormal as _TruncatedNormal,

This file has been truncated. show original

Both are similar, but they are not the same.

Why is this, and which should I use?

Thank you!