def initialize_episode()

in gym_aloha/tasks/sim_end_effector.py [0:0]

• 12 lines of code
• 2 McCabe index (conditional complexity)

    def initialize_episode(self, physics):
        """Sets the state of the environment at the start of each episode."""
        self.initialize_robots(physics)
        # randomize peg and socket position
        peg_pose, socket_pose = sample_insertion_pose()

        def id2index(j_id):
            return 16 + (j_id - 16) * 7  # first 16 is robot qpos, 7 is pose dim # hacky

        peg_start_id = physics.model.name2id("red_peg_joint", "joint")
        peg_start_idx = id2index(peg_start_id)
        np.copyto(physics.data.qpos[peg_start_idx : peg_start_idx + 7], peg_pose)
        # print(f"randomized cube position to {cube_position}")

        socket_start_id = physics.model.name2id("blue_socket_joint", "joint")
        socket_start_idx = id2index(socket_start_id)
        np.copyto(physics.data.qpos[socket_start_idx : socket_start_idx + 7], socket_pose)
        # print(f"randomized cube position to {cube_position}")

        super().initialize_episode(physics)