Reference : Data-Efficient Reinforcement Learning with Self-Predictive Representations Thanks to Repo: https://github.com/mila-iqia/spr.git
| 6 | |
| 7 | |
| 8 | class Transforms(object): |
| 9 | """ Reference : Data-Efficient Reinforcement Learning with Self-Predictive Representations |
| 10 | Thanks to Repo: https://github.com/mila-iqia/spr.git |
| 11 | """ |
| 12 | def __init__(self, augmentation, shift_delta=4, image_shape=(96, 96)): |
| 13 | self.augmentation = augmentation |
| 14 | |
| 15 | self.transforms = [] |
| 16 | for aug in self.augmentation: |
| 17 | if aug == "affine": |
| 18 | transformation = RandomAffine(5, (.14, .14), (.9, 1.1), (-5, 5)) |
| 19 | elif aug == "crop": |
| 20 | transformation = RandomCrop(image_shape) |
| 21 | elif aug == "rrc": |
| 22 | transformation = RandomResizedCrop((100, 100), (0.8, 1)) |
| 23 | elif aug == "blur": |
| 24 | transformation = GaussianBlur2d((5, 5), (1.5, 1.5)) |
| 25 | elif aug == "shift": |
| 26 | transformation = nn.Sequential(nn.ReplicationPad2d(shift_delta), RandomCrop(image_shape)) |
| 27 | elif aug == "intensity": |
| 28 | transformation = Intensity(scale=0.05) |
| 29 | elif aug == "none": |
| 30 | transformation = nn.Identity() |
| 31 | else: |
| 32 | raise NotImplementedError() |
| 33 | self.transforms.append(transformation) |
| 34 | |
| 35 | def apply_transforms(self, transforms, image): |
| 36 | for transform in transforms: |
| 37 | image = transform(image) |
| 38 | return image |
| 39 | |
| 40 | @torch.no_grad() |
| 41 | def transform(self, images): |
| 42 | # images = images.float() / 255. if images.dtype == torch.uint8 else images |
| 43 | flat_images = images.reshape(-1, *images.shape[-3:]) |
| 44 | processed_images = self.apply_transforms(self.transforms, flat_images) |
| 45 | |
| 46 | processed_images = processed_images.view(*images.shape[:-3], |
| 47 | *processed_images.shape[1:]) |
| 48 | return processed_images |
| 49 | |
| 50 | |
| 51 | class Intensity(nn.Module): |