def modify_sim_step()

in mae_envs/modules/objects.py [0:0]

• 11 lines of code
• 7 McCabe index (conditional complexity)

    def modify_sim_step(self, env, sim):
        # Cache qpos, qvel idxs
        self.box_geom_idxs = np.array([sim.model.geom_name2id(f'moveable_box{i}')
                                       for i in range(self.curr_n_boxes)])
        self.box_qpos_idxs = np.array([qpos_idxs_from_joint_prefix(sim, f'moveable_box{i}:')
                                       for i in range(self.curr_n_boxes)])
        self.box_qvel_idxs = np.array([qvel_idxs_from_joint_prefix(sim, f'moveable_box{i}:')
                                       for i in range(self.curr_n_boxes)])
        if self.mark_box_corners:
            self.box_corner_idxs = np.array([sim.model.site_name2id(f'moveable_box{i}_corner{j}')
                                             for i in range(self.curr_n_boxes)
                                             for j in range(4)])