def step(self, aug_action):
    cur_idx = self.ori_action_dim
    action = aug_action[:self.ori_action_dim]
    if (self.anoise_type == 'action'):
        if (np.random.rand() < self.anoise_prob):
            action = (action + (np.random.randn(*action.shape) * self.anoise_std))
            action = np.clip(action, self.action_space.low[:len(action)], self.action_space.high[:len(action)])
    elif ((self.anoise_type is not None) and ('ext' in self.anoise_type)):
        if (np.random.rand() < self.anoise_prob):
            if (self.env.spec.id == 'Reacher-v2'):
                force = np.zeros(3)
                torque = (np.random.randn(3) * self.anoise_std)
                cur_info = torque
            else:
                force = (np.random.randn(3) * self.anoise_std)
                torque = np.zeros(3)
                cur_info = force
            if (self.anoise_type == 'ext_fpc'):
                self.cur_force = np.clip(cur_info, (- 1), 1)
            self.env.sim.data.xfrc_applied[self.env.sim.model._body_name2id[self.body_key], :] = np.r_[(force, torque)]
        else:
            self.env.sim.data.xfrc_applied[self.env.sim.model._body_name2id[self.body_key], :] = ([0] * 6)
    if ((self.max_d is not None) or (self.max_t is not None)):
        u = aug_action[cur_idx]
        cur_idx += 1
        norm_u = ((u + 1) / 2)
        u = norm_u
    else:
        u = None
    lambda_dt = self.lambda_dt
    total_reward = 0.0
    done = None
    cur_gamma = 1.0
    first_obs = self.cur_obs
    for i in range(100000000):
        (obs, reward, done, info) = self.env.step(action)
        if (self.anoise_type in ['ext_fpc']):
            obs = np.concatenate([obs, self.cur_force])
        if (not done):
            self._update_obs_estimate(obs[(np.newaxis, ...)])
            self.reset_update_obs_estimate = True
        total_reward += (reward * cur_gamma)
        cur_gamma *= self.gamma
        if done:
            break
        if ((self.max_d is None) and (self.max_t is None)):
            break
        if (self.max_t is not None):
            t_delta = ((i + 1) * self.env.dt)
        if (self.max_d is not None):
            norm_obs = ((obs - self.obs_rms.mean) / (np.sqrt(self.obs_rms.var) + self.epsilon_std))
            norm_first_obs = ((first_obs - self.obs_rms.mean) / (np.sqrt(self.obs_rms.var) + self.epsilon_std))
            d_delta = (np.linalg.norm((norm_obs - norm_first_obs), ord=1) / len(obs))
        if ((self.max_d is not None) and (self.max_t is not None)):
            if (lambda_dt is None):
                if (d_delta >= (u * self.max_d)):
                    break
                if (t_delta >= self.max_t):
                    break
            else:
                ori_t_delta = t_delta
                t_delta = (t_delta / self.max_t)
                d_delta = (d_delta / self.max_d)
                delta = ((lambda_dt * d_delta) + ((1 - lambda_dt) * t_delta))
                if (delta >= u):
                    break
                if (ori_t_delta >= self.max_t):
                    break
        elif (self.max_t is not None):
            if (t_delta >= (u * self.max_t)):
                break
        elif (self.max_d is not None):
            if (d_delta >= (u * self.max_d)):
                break
    self.cur_obs = obs
    info['w'] = (i + 1)
    info['t_diff'] = ((i + 1) * self.env.dt)
    if (u is not None):
        if ((self.max_d is not None) and (self.max_t is not None)):
            pass
        elif (self.max_t is not None):
            info['t'] = (u * self.max_t)
        elif (self.max_d is not None):
            info['d'] = (u * self.max_d)
        info['u'] = u
    if (lambda_dt is not None):
        info['lambda_dt'] = lambda_dt
    self.num_steps += 1
    return (self.cur_obs, total_reward, done, info)
def step(self, aug_action):
    cur_idx = self.ori_action_dim
    action = aug_action[:self.ori_action_dim]
    if (self.anoise_type == 'action'):
        if (np.random.rand() < self.anoise_prob):
            action = (action + (np.random.randn(*action.shape) * self.anoise_std))
            action = np.clip(action, self.action_space.low[:len(action)], self.action_space.high[:len(action)])
    elif ((self.anoise_type is not None) and ('ext' in self.anoise_type)):
        if (np.random.rand() < self.anoise_prob):
            if (self.env.spec.id == 'Reacher-v2'):
                force = np.zeros(3)
                torque = (np.random.randn(3) * self.anoise_std)
                cur_info = torque
            else:
                force = (np.random.randn(3) * self.anoise_std)
                torque = np.zeros(3)
                cur_info = force
            if (self.anoise_type == 'ext_fpc'):
                self.cur_force = np.clip(cur_info, (- 1), 1)
            self.env.sim.data.xfrc_applied[self.env.sim.model._body_name2id[self.body_key], :] = np.r_[(force, torque)]
        else:
            self.env.sim.data.xfrc_applied[self.env.sim.model._body_name2id[self.body_key], :] = ([0] * 6)
    if ((self.max_d is not None) or (self.max_t is not None)):
        u = aug_action[cur_idx]
        cur_idx += 1
        norm_u = ((u + 1) / 2)
        u = norm_u
    else:
        u = None
    lambda_dt = self.lambda_dt
    total_reward = 0.0
    done = None
    cur_gamma = 1.0
    first_obs = self.cur_obs
    for i in range(100000000):
        (obs, reward, done, info) = self.env.step(action)
        if (self.anoise_type in ['ext_fpc']):
            obs = np.concatenate([obs, self.cur_force])
        if (not done):
            self._update_obs_estimate(obs[(np.newaxis, ...)])
            self.reset_update_obs_estimate = True
        total_reward += (reward * cur_gamma)
        cur_gamma *= self.gamma
        if done:
            break
        if ((self.max_d is None) and (self.max_t is None)):
            break
        if (self.max_t is not None):
            t_delta = ((i + 1) * self.env.dt)
        if (self.max_d is not None):
            norm_obs = ((obs - self.obs_rms.mean) / (np.sqrt(self.obs_rms.var) + self.epsilon_std))
            norm_first_obs = ((first_obs - self.obs_rms.mean) / (np.sqrt(self.obs_rms.var) + self.epsilon_std))
            d_delta = (np.linalg.norm((norm_obs - norm_first_obs), ord=1) / len(obs))
        if ((self.max_d is not None) and (self.max_t is not None)):
            if (lambda_dt is None):
                if (d_delta >= (u * self.max_d)):
                    break
                if (t_delta >= self.max_t):
                    break
            else:
                ori_t_delta = t_delta
                t_delta = (t_delta / self.max_t)
                d_delta = (d_delta / self.max_d)
                delta = ((lambda_dt * d_delta) + ((1 - lambda_dt) * t_delta))
                if (delta >= u):
                    break
                if (ori_t_delta >= self.max_t):
                    break
        elif (self.max_t is not None):
            if (t_delta >= (u * self.max_t)):
                break
        elif (self.max_d is not None):
            if (d_delta >= (u * self.max_d)):
                break
    self.cur_obs = obs
    info['w'] = (i + 1)
    info['t_diff'] = ((i + 1) * self.env.dt)
    if (u is not None):
        if ((self.max_d is not None) and (self.max_t is not None)):
            pass
        elif (self.max_t is not None):
            info['t'] = (u * self.max_t)
        elif (self.max_d is not None):
            info['d'] = (u * self.max_d)
        info['u'] = u
    if (lambda_dt is not None):
        info['lambda_dt'] = lambda_dt
    self.num_steps += 1
    return (self.cur_obs, total_reward, done, info)