| 697 | |
| 698 | |
| 699 | class CosTransform(Transform): |
| 700 | def _apply_transform(self, obs: torch.Tensor) -> None: |
| 701 | return obs.cos() |
| 702 | |
| 703 | # The transform must also modify the data at reset time |
| 704 | def _reset( |
| 705 | self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase |
| 706 | ) -> TensorDictBase: |
| 707 | return self._call(tensordict_reset) |
| 708 | |
| 709 | # _apply_to_composite will execute the observation spec transform across all |
| 710 | # in_keys/out_keys pairs and write the result in the observation_spec which |
| 711 | # is of type ``Composite`` |
| 712 | @_apply_to_composite |
| 713 | def transform_observation_spec(self, observation_spec): |
| 714 | return BoundedTensorSpec( |
| 715 | low=-1, |
| 716 | high=1, |
| 717 | shape=observation_spec.shape, |
| 718 | dtype=observation_spec.dtype, |
| 719 | device=observation_spec.device, |
| 720 | ) |
| 721 | |
| 722 | |
| 723 | t_sin = SinTransform(in_keys=["th"], out_keys=["sin"]) |