def __init__()

in mae_envs/envs/box_locking.py [0:0]

• 32 lines of code
• 3 McCabe index (conditional complexity)

    def __init__(self, env, n_objs, task='all', fixed_order=False,
                 obj_lock_obs_key='obj_lock', obj_pos_obs_key='box_pos',
                 act_lock_key='action_glue', agent_pos_key='agent_pos',
                 lock_reward=5.0, unlock_penalty=10.0, shaped_reward_scale=1.0,
                 success_reward=1, return_threshold=0.1):
        super().__init__(env)
        self.n_objs = n_objs
        self.task = task or 'all'
        assert task in ['all', 'order', 'all-return', 'order-return'], (
            f'task {task} is currently not supported')
        self.need_return = 'return' in self.task
        self.return_threshold = return_threshold
        if self.need_return:
            self.task = self.task.replace('-return', '')
        self.n_agents = self.unwrapped.n_agents
        assert self.n_agents == 1, 'The locking tasks only support 1 agent'
        self.agent_key = agent_pos_key
        self.obj_order = list(range(self.n_objs))
        self.fixed_order = fixed_order
        self.lock_key = obj_lock_obs_key
        self.pos_key = obj_pos_obs_key
        self.act_key = act_lock_key
        self.lock_reward = lock_reward
        self.unlock_penalty = unlock_penalty
        self.shaped_reward_scale = shaped_reward_scale
        self.success_reward = success_reward
        self.objs_locked = np.zeros((n_objs, ), dtype=np.int8)
        self.spawn_pos = None
        self.spawn_pos_dist = None
        self.next_obj = None
        self.next_obj_dist = 0
        self.unlocked_objs = []