def reset()

in mae_envs/wrappers/manipulation.py [0:0]

• 25 lines of code
• 10 McCabe index (conditional complexity)

    def reset(self):
        obs = self.env.reset()
        sim = self.unwrapped.sim

        if self.obj_in_game_metadata_keys is not None:
            self.actual_body_slice = np.concatenate([self.metadata[k] for k in self.obj_in_game_metadata_keys])
        else:
            self.actual_body_slice = np.ones((len(self.body_names))).astype(np.bool)

        actual_body_names = list(compress(self.body_names, self.actual_body_slice))
        self.n_obj = len(actual_body_names)

        # Cache ids
        self.obj_body_idxs = np.array([sim.model.body_name2id(body_name) for body_name in actual_body_names])
        self.obj_jnt_idxs = [np.where(sim.model.jnt_bodyid == body_idx)[0] for body_idx in self.obj_body_idxs]
        self.obj_geom_ids = [np.where(sim.model.geom_bodyid == body_idx)[0] for body_idx in self.obj_body_idxs]
        self.agent_body_idxs = np.array([sim.model.body_name2id(f"agent{i}:particle") for i in range(self.n_agents)])
        self.agent_body_idxs = self.agent_body_idxs[self.agent_idx_allowed_to_lock]
        self.agent_geom_ids = np.array([sim.model.geom_name2id(f'agent{i}:agent') for i in range(self.n_agents)])
        self.agent_geom_ids = self.agent_geom_ids[self.agent_idx_allowed_to_lock]

        self.unlock_objs()
        self.obj_locked = np.zeros((self.n_obj,), dtype=bool)
        self.which_locked = np.zeros((self.n_obj,), dtype=int)

        if self.agent_allowed_to_lock_keys is not None:
            self.agent_allowed_to_lock_mask = np.concatenate([obs[k] for k in self.agent_allowed_to_lock_keys])
        else:
            self.agent_allowed_to_lock_mask = np.ones((self.n_agents, self.n_obj))

        return self.observation(obs)