UnsqueezeTransform

class torchrl.envs.transforms.UnsqueezeTransform(*args, **kwargs)

Inserts a dimension of size one at the specified position.

Parameters:

• unsqueeze_dim (int) – dimension to unsqueeze. Must be negative (or allow_positive_dim must be turned on).

• allow_positive_dim (bool, optional) – if True, positive dimensions are accepted. UnsqueezeTransform will map these to the n^th feature dimension (ie n^th dimension after batch size of parent env) of the input tensor, independently from the tensordict batch size (ie positive dims may be dangerous in contexts where tensordict of different batch dimension are passed). Defaults to False, ie. non-negative dimensions are not permitted.

transform_input_spec(input_spec)

Transforms the input spec such that the resulting spec matches transform mapping.

Parameters:

input_spec (TensorSpec) – spec before the transform

Returns:

expected spec after the transform

transform_observation_spec(observation_spec: TensorSpec) → TensorSpec

Transforms the observation spec such that the resulting spec matches transform mapping.

Parameters:

observation_spec (TensorSpec) – spec before the transform

Returns:

expected spec after the transform

transform_reward_spec(reward_spec: TensorSpec) → TensorSpec

Transforms the reward spec such that the resulting spec matches transform mapping.

Parameters:

reward_spec (TensorSpec) – spec before the transform

Returns:

expected spec after the transform