UnsqueezeTransform

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

Inserts a dimension of size one at the specified position.

Parameters:

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

Keyword Arguments:

• 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 of 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.

• in_keys (list of NestedKeys) – input entries (read).

• out_keys (list of NestedKeys) – input entries (write). Defaults to in_keys if not provided.

• in_keys_inv (list of NestedKeys) – input entries (read) during inv() calls.

• out_keys_inv (list of NestedKeys) – input entries (write) during inv() calls. Defaults to in_keys_in if not provided.

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