def get_action(self, obs, env, timestep):
    cur_stage = (- 1)
    if (not self.milestone_0_complete):
        if self.milestone_0_cond(obs):
            self.milestone_0_complete = True
            cur_stage = 1
        else:
            cur_stage = 0
    else:
        if (not self.milestone_1_complete):
            if self.milestone_1_cond(obs):
                self.milestone_1_complete = True
                self.first_time_all_complete = timestep
                print('solved')
        cur_stage = 1
    if (cur_stage == 0):
        grip_pos = env.sim.data.get_site_xpos('robot0:grip')
        action = [0, 0, 0, 0]
        pos_act = get_linear_pos_act(grip_pos, self.waypoint)
        pos_act += np.random.uniform(0.0, ACT_NOISE_SCALE, 3)
        for i in range(len(pos_act)):
            action[i] = pos_act[i]
        action[(len(action) - 1)] = np.random.uniform((- 0.005), (- 0.015))
    else:
        action = [0, 0, 0, 0]
        grip_pos = env.sim.data.get_site_xpos('robot0:grip')
        target_rel_pos = (env.goal - grip_pos)
        if (np.linalg.norm(target_rel_pos, axis=(- 1)) < SLOW_DOWN_RADIUS):
            pos_act = (0.25 * get_linear_pos_act(np.zeros(3), target_rel_pos))
            pos_act += np.random.uniform(0.0, ACT_SLOW_NOISE_SCALE, 3)
        else:
            pos_act = get_linear_pos_act(np.zeros(3), target_rel_pos)
            pos_act += np.random.uniform(0.0, ACT_NOISE_SCALE, 3)
        for i in range(len(pos_act)):
            action[i] = pos_act[i]
        action[(len(action) - 1)] = np.random.uniform((- 0.005), (- 0.015))
    action = np.clip(action, (- 1.0), 1.0)
    return (action, {})
def get_action(self, obs, env, timestep):
    cur_stage = (- 1)
    if (not self.milestone_0_complete):
        if self.milestone_0_cond(obs):
            self.milestone_0_complete = True
            cur_stage = 1
        else:
            cur_stage = 0
    else:
        if (not self.milestone_1_complete):
            if self.milestone_1_cond(obs):
                self.milestone_1_complete = True
                self.first_time_all_complete = timestep
                print('solved')
        cur_stage = 1
    if (cur_stage == 0):
        grip_pos = env.sim.data.get_site_xpos('robot0:grip')
        action = [0, 0, 0, 0]
        pos_act = get_linear_pos_act(grip_pos, self.waypoint)
        pos_act += np.random.uniform(0.0, ACT_NOISE_SCALE, 3)
        for i in range(len(pos_act)):
            action[i] = pos_act[i]
        action[(len(action) - 1)] = np.random.uniform((- 0.005), (- 0.015))
    else:
        action = [0, 0, 0, 0]
        grip_pos = env.sim.data.get_site_xpos('robot0:grip')
        target_rel_pos = (env.goal - grip_pos)
        if (np.linalg.norm(target_rel_pos, axis=(- 1)) < SLOW_DOWN_RADIUS):
            pos_act = (0.25 * get_linear_pos_act(np.zeros(3), target_rel_pos))
            pos_act += np.random.uniform(0.0, ACT_SLOW_NOISE_SCALE, 3)
        else:
            pos_act = get_linear_pos_act(np.zeros(3), target_rel_pos)
            pos_act += np.random.uniform(0.0, ACT_NOISE_SCALE, 3)
        for i in range(len(pos_act)):
            action[i] = pos_act[i]
        action[(len(action) - 1)] = np.random.uniform((- 0.005), (- 0.015))
    action = np.clip(action, (- 1.0), 1.0)
    return (action, {})