Examples of data augmentation for symmetry

[StephenWelch]

Hi! Just wanted to ask if there are any examples available for implementing the data_augmentation_func/get_symmetric_states for symmetry data augmentation.

[Evan-wyl]

import torch


def data_augmentation_func_g1(obs, actions, env, is_critic):
    if is_critic:
        obs_batch = flip_g1_critic_obs(obs)
    else:
        obs_batch = flip_g1_critic_obs(actions)

    mean_actions_batch = flip_g1_actions(actions)
    return obs_batch, mean_actions_batch


def flip_g1_actor_obs(self, obs):
    if obs is None:
        return obs

    flipped_obs = torch.zeros_like(obs)
    flipped_obs[..., :3] = obs[..., :3]
    flipped_obs[..., 3:6] = obs[..., 3:6]
    flipped_obs[..., 6:9] = obs[..., 6:9]

    flipped_obs[..., 9:15] = obs[..., 15:21]
    flipped_obs[..., 15:21] = obs[..., 9:15]

    flipped_obs[..., 21:27] = obs[..., 27:33]
    flipped_obs[..., 27:33] = obs[..., 21:27]

    flipped_obs[..., 33:39] = obs[..., 39:45]
    flipped_obs[..., 39:45] = obs[..., 33:39]

    flipped_obs[..., 45:47] = obs[..., 45:47]
    return flipped_obs


def flip_g1_critic_obs(self, obs):
    if obs is None:
        return obs

    flipped_obs = torch.zeros_like(obs)
    flipped_obs[..., :3] = obs[..., :3]
    flipped_obs[..., 3:6] = obs[..., 3:6]
    flipped_obs[..., 6:9] = obs[..., 6:9]

    flipped_obs[..., 9:15] = obs[..., 15:21]
    flipped_obs[..., 15:21] = obs[..., 9:15]

    flipped_obs[..., 21:27] = obs[..., 27:33]
    flipped_obs[..., 27:33] = obs[..., 21:27]

    flipped_obs[..., 33:39] = obs[..., 39:45]
    flipped_obs[..., 39:45] = obs[..., 33:39]

    flipped_obs[..., 45:47] = obs[..., 45:47]
    flipped_obs[..., 47:50] = obs[..., 47:50]
    flipped_obs[..., 50:51] = obs[..., 50:51]
    return flipped_obs


def flip_g1_actions(self, actions):
    if actions is None:
        return None

    flip_actions = torch.zeros_like(actions)
    flip_actions[..., :6] = actions[..., :6]
    flip_actions[..., 6:12] = actions[..., 6:12]
    return flip_actions

[hshi74]

Thanks for posting the example! Just to confirm, in the else block,
obs_batch = flip_g1_critic_obs(actions) should be obs_batch = flip_g1_actor_obs(obs)?