Ctrl+K

core.transform

core.transform#

Transformation module for all data transformations.

class soulsai.core.transform.Transform#

Base class for transformations.

update(x: TensorDict)#

Update the transformation parameters.

The base update is a no-op so that only transformations that require updates need to implement this method.

Parameters:

x – Input tensor.

serialize() → bytes#

Serialize the transformation into bytes for synchronization across nodes.

deserialize(serialization: bytes)#

Deserialize the transformation from bytes and load them into the transformation.

class soulsai.core.transform.Identity#

Identity transformation class for no transformation.

forward(x: TensorDict, keys_mapping: dict[NestedKey, NestedKey] | None = None) → TensorDict#

Return the input tensor unchanged.

Parameters:

  • x – Input tensor.

  • keys_mapping – Optional dictionary that maps keys to new keys. Present for compatibility.

Returns:

Input tensor.

class soulsai.core.transform.Chain(transforms: list[Transform | dict])#

Chain transformation class for chaining multiple transformations together.

forward(x: TensorDict, keys_mapping: dict[NestedKey, NestedKey] | None = None) → TensorDict#

Apply the transformations in the chain to the input tensor.

Parameters:

  • x – Input tensor.

  • keys_mapping – Optional dictionary that maps keys to new keys. This is useful when the transformation should be applied to a subset of the keys, or when the same parameters should be used for different keys, i.e. for ‘obs’ and ‘next_obs’.

Returns:

Transformed TensorDict.

class soulsai.core.transform.Normalize(keys: list[NestedKey], shapes: list[tuple[int]], indexes: list[list[int] | None] | None = None)#

Normalization class for normalizing tensors to have zero mean and unit variance.

Mean and standard deviation parameters are estimated from the data during the update step.

forward(x: TensorDict, keys_mapping: dict[NestedKey, NestedKey] | None = None) → TensorDict#

Normalize the input TensorDict to have zero mean and unit variance.

Parameters:

  • x – Sample TensorDict.

  • keys_mapping – Optional dictionary that maps keys to new keys. This is useful when the transformation should be applied to a subset of the keys, or when the same parameters should be used for different keys, i.e. for ‘obs’ and ‘next_obs’.

Returns:

Normalized TensorDict.

update(x: TensorDict)#

Update the keys’ mean and standard deviation estimate from a sample TensorDict batch.

Parameters:

x – Sample TensorDict.

class soulsai.core.transform.NormalizeImg(keys: list[NestedKey])#

Image normalization transformation class for normalizing images to the range [-1, 1].

forward(x: TensorDict, keys_mapping: dict[NestedKey, NestedKey] | None = None) → TensorDict#

Normalize the input tensor to the range [-1, 1].

Parameters:

  • x – Sample TensorDict.

  • keys_mapping – Optional dictionary that maps keys to new keys. This is useful when the normalization should be applied to a subset of the keys, or when the same parameters should be used for different keys, i.e. for ‘obs’ and ‘next_obs’.

Returns:

Normalized TensorDict.

class soulsai.core.transform.Mask(key: NestedKey, mask_key: NestedKey, mask_value: float = -inf)#

Mask the value tensor with -inf at the masked indices.

forward(x: TensorDict, keys_mapping: dict[NestedKey, NestedKey] | None = None) → TensorDict#

Mask the action value tensor with self._mask_value at the masked indices.

Parameters:

  • x – Input TensorDict.

  • keys_mapping – Optional dictionary that maps keys to new keys. This is useful when the transformation should be applied to a subset of the keys, or when the same parameters should be used for different keys, i.e. for ‘obs’ and ‘next_obs’.

Returns:

The masked TensorDict.

mask_tensor(x: Tensor, sample: TensorDict, auto_expand: bool = False) → Tensor#

Mask the input tensor with self._mask_value at the masked indices.

Parameters:

  • x – Input tensor.

  • sample – Input TensorDict containing self._mask_key.

  • auto_expand – Whether to automatically expand the mask to the right.

Returns:

The masked tensor.

class soulsai.core.transform.GreedyAction(value_key: str, action_key: str)#

Greedy action transformation class for selecting the action with the highest value.

forward(x: TensorDict, keys_mapping: dict[NestedKey, NestedKey] | None = None) → TensorDict#

Select the action with the highest value.

Parameters:

  • x – Input TensorDict.

  • keys_mapping – Optional dictionary that maps keys to new keys. This is useful when the transformation should be applied to a subset of the keys, or when the same parameters should be used for different keys, i.e. for ‘obs’ and ‘next_obs’.

Returns:

The input TensorDict with the action tensor set to the action with the highest value.

class soulsai.core.transform.ExponentialAction(value_key: NestedKey, action_key: NestedKey, scheduler: Scheduler | dict)#

Select an action using Boltzmann exploration with a temperature parameter.

forward(x: TensorDict, keys_mapping: dict[NestedKey, NestedKey] | None = None) → TensorDict#

Select an action using Boltzmann exploration with a temperature parameter.

Parameters:

  • x – Input TensorDict.

  • keys_mapping – Optional dictionary that maps keys to new keys. This is useful when the transformation should be applied to a subset of the keys, or when the same parameters should be used for different keys, i.e. for ‘obs’ and ‘next_obs’.

Returns:

The input TensorDict with the action tensor set to the action selected using Boltzmann exploration.

update(x: TensorDict)#

Update the temperature parameter.

class soulsai.core.transform.Choice(key: NestedKey, transforms: list[Transform | dict], probs: list[float])#

Choice transformation class for selecting one of several transforms at random.

forward(x: TensorDict, keys_mapping: dict[NestedKey, NestedKey] | None = None) → TensorDict#

Choose a transformation at random and apply it to the input tensor.

Parameters:

  • x – Sample TensorDict.

  • keys_mapping – Optional dictionary that maps keys to new keys. This is useful when the transformation should be applied to a subset of the keys, or when the same parameters should be used for different keys, i.e. for ‘obs’ and ‘next_obs’.

Returns:

Transformed TensorDict.

class soulsai.core.transform.ScheduledChoice(key: NestedKey, transforms: list[Transform | dict], scheduler: Scheduler | dict)#

Scheduled choice transformation class for selecting one of several transforms at random.

The chance of choosing a transformation is scheduled by a scheduler. The scheduler advances on each call to update. This allows us to implement behaviors like annealing exploration rates.

forward(x: TensorDict, keys_mapping: dict[NestedKey, NestedKey] | None = None) → TensorDict#

Choose a transformation at random and apply it to the input tensor.

Parameters:

  • x – Sample TensorDict.

  • keys_mapping – Optional dictionary that maps keys to new keys. This is useful when the transformation should be applied to a subset of the keys, or when the same parameters should be used for different keys, i.e. for ‘obs’ and ‘next_obs’.

Returns:

Transformed TensorDict.

update(_: TensorDict)#

Update the transformation parameters.

Parameters:

x – Input tensor.

class soulsai.core.transform.ReplaceWithNoise(key: NestedKey, noise: Noise | dict)#

Replace the input tensor with noise.

forward(x: TensorDict, keys_mapping: dict[NestedKey, NestedKey] | None = None) → Tensor#

Replace the input tensor with noise.

Parameters:

  • x – Sample TensorDict.

  • keys_mapping – Optional dictionary that maps keys to new keys. This is useful when the transformation should be applied to a subset of the keys, or when the same parameters should be used for different keys, i.e. for ‘obs’ and ‘next_obs’.

Returns:

Transformed TensorDict.

Contents