def reset()

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

• 21 lines of code
• 9 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 body ids
        self.obj_body_idxs = np.array([sim.model.body_name2id(body_name) for body_name in actual_body_names])
        self.agent_body_idxs = np.array([sim.model.body_name2id(f"agent{i}:particle") for i in range(self.n_agents)])

        # Cache geom ids
        self.obj_geom_ids = np.array([sim.model.geom_name2id(body_name) for body_name in actual_body_names])
        self.agent_geom_ids = np.array([sim.model.geom_name2id(f'agent{i}:agent') for i in range(self.n_agents)])

        # Cache constraint ids
        self.agent_eq_ids = np.array(
            [i for i, obj1 in enumerate(sim.model.eq_obj1id)
             if sim.model.body_names[obj1] == f"agent{i}:particle"])
        assert len(self.agent_eq_ids) == self.n_agents

        # turn off equality constraints
        sim.model.eq_active[self.agent_eq_ids] = 0
        self.obj_grabbed = np.zeros((self.n_agents, self.n_obj), dtype=bool)
        self.last_obj_grabbed = np.zeros((self.n_agents, self.n_obj), dtype=bool)

        return self.observation(obs)