Tensor Specs: Typing Actions, Observations, and Rewards

from torchrl.data import Bounded, Categorical, Composite
import torch

# Define what an environment's spaces look like
obs_spec   = Bounded(low=-1.0, high=1.0, shape=(8,), dtype=torch.float32)
action_spec = Categorical(n=4, shape=())

# Sample a random observation and action — no environment needed
obs    = obs_spec.rand()     # tensor in [-1, 1]^8
action = action_spec.rand()  # tensor in {0, 1, 2, 3}

# Check membership
print(obs_spec.is_in(obs))     # True
print(action_spec.is_in(torch.tensor(10)))  # False — 10 ≥ n=4

# Project an out-of-domain value back into the spec
clipped = obs_spec.project(torch.tensor([2.0, -2.0, 0.0, 0.5, 0.5, 0.5, 0.5, 0.5]))

from torchrl.data import Bounded
import torch

# Continuous spec for a 4-dimensional observation in [-1, 1]
obs_spec = Bounded(low=-1.0, high=1.0, shape=(4,), dtype=torch.float32)
print(obs_spec.rand())   # tensor in [-1, 1]^4
print(obs_spec.is_in(torch.zeros(4)))   # True
print(obs_spec.is_in(torch.full((4,), 2.0)))  # False

# Per-dimension bounds via tensors
action_spec = Bounded(
    low=torch.tensor([-1.0, -0.5]),
    high=torch.tensor([1.0, 0.5]),
)

from torchrl.data import Unbounded, UnboundedContinuous, UnboundedDiscrete
import torch

reward_spec = Unbounded(shape=(1,), dtype=torch.float32)
# Equivalent to:
reward_spec = UnboundedContinuous(shape=(1,))

count_spec = UnboundedDiscrete(shape=(1,), dtype=torch.int64)
print(reward_spec.rand())  # draws from N(0, 1)

from torchrl.data import Categorical
import torch

# Discrete action space with 6 actions
action_spec = Categorical(n=6)
print(action_spec.rand())       # tensor in {0..5}
print(action_spec.rand((4,)))   # 4 independent samples

# Masking: only allow actions 0, 2, 4
mask = torch.tensor([True, False, True, False, True, False])
action_spec.update_mask(mask)
print(action_spec.rand())  # always 0, 2, or 4

from torchrl.data import OneHot
import torch

spec = OneHot(n=5, dtype=torch.float32)
sample = spec.rand()
print(sample)          # e.g. tensor([0., 0., 1., 0., 0.])
print(sample.sum())    # always 1.0

# Convert between one-hot and categorical:
categ = spec.to_categorical(sample)  # tensor(2)
back  = spec.to_one_hot(categ)

from torchrl.data import MultiOneHot
import torch

# Two independent sub-actions: 3 choices and 2 choices
spec = MultiOneHot(nvec=(3, 2))
sample = spec.rand()
print(sample)        # e.g. tensor([False, True, False, True, False])
print(sample.shape)  # torch.Size([5])  (= 3+2)

# Decode back to integer indices
indices = spec.to_categorical(sample)  # tensor([1, 0])

from torchrl.data import Binary
import torch

spec = Binary(n=4, shape=(2, 4))
print(spec.rand())
# e.g. tensor([[0, 1, 1, 0],
#               [1, 1, 1, 1]], dtype=torch.int8)

from torchrl.data import Composite, Bounded, Categorical, Unbounded
import torch

# Environment observation spec
obs_spec = Composite(
    pixels=Bounded(
        low=torch.zeros(3, 84, 84, dtype=torch.uint8),
        high=torch.full((3, 84, 84), 255, dtype=torch.uint8),
    ),
    proprioception=Bounded(low=-10.0, high=10.0, shape=(27,)),
)

# Nested composite specs are created automatically for tuple keys:
full_spec = Composite(
    {("next", "observation"): Bounded(low=-1.0, high=1.0, shape=(8,))},
    action=Categorical(n=4),
)

# Sampling returns a TensorDict:
td = obs_spec.rand()
print(type(td))          # tensordict.TensorDict
print(td["pixels"].shape)  # torch.Size([3, 84, 84])
print(td["proprioception"].shape)  # torch.Size([27])

# Iterate over keys:
for key, spec in obs_spec.items():
    print(key, spec.shape)

from torchrl.data import Bounded, Composite
import torch

spec = Bounded(low=-1.0, high=1.0, shape=(4,))

# --- Shape manipulation ---
batched = spec.expand(8, 4)    # shape (8, 4), same bounds
squeezed = spec.unsqueeze(0)   # shape (1, 4)

# --- Type and device ---
gpu_spec = spec.to("cuda")
f64_spec = spec.to(torch.float64)

# --- Sampling ---
single = spec.rand()           # shape (4,)
batch  = spec.rand((32,))      # shape (32, 4)
zeros  = spec.zero((8,))       # shape (8, 4)

# --- Validation ---
val = torch.randn(4) * 2       # may be out of bounds
print(spec.is_in(val))         # possibly False
clipped = spec.project(val)    # clamps to [-1, 1]
spec.assert_is_in(clipped)     # raises AssertionError if fails

# --- encode (numpy/list → tensor) ---
import numpy as np
arr = np.array([0.1, -0.2, 0.5, 0.3])
tensor = spec.encode(arr)      # torch.Tensor on spec.device

# --- Composite spec operations ---
comp = Composite(obs=spec, action=Bounded(low=0.0, high=1.0, shape=(2,)))
td = comp.rand()
comp.assert_is_in(td)

from torchrl.envs import GymEnv, TransformedEnv, ObservationNorm

env = GymEnv("HalfCheetah-v4")

# Inspect the observation spec
print(env.observation_spec)
print(env.action_spec)
print(env.reward_spec)

# Check that a rollout satisfies the specs
td = env.rollout(10)
env.check_env_specs()  # validates observation_spec, action_spec, reward_spec

# Specs propagate through transforms:
tenv = TransformedEnv(env, ObservationNorm(loc=0, scale=1))
print(tenv.observation_spec)  # updated by the transform