def reset_then_step()

in polymetis/python/polysim/envs/bullet_locomotor.py [0:0]

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

    def reset_then_step(self, joint_des_state, torque):
        curr_state = self.get_current_full_state()

        joint_pos = joint_des_state[: self.n_dofs]
        base_pos = joint_des_state[self.n_dofs : self.n_dofs + 6]
        base_quat_orient = pybullet.getQuaternionFromEuler(base_pos[3:])
        joint_vel = joint_des_state[self.n_dofs + 6 : 2 * self.n_dofs + 6]
        base_velocity = joint_des_state[self.n_dofs * 2 + 6 :]

        self.robot.parts["daisy"].reset_pose(base_pos[:3], base_quat_orient)
        self.robot.parts["daisy"].reset_velocity(base_velocity)

        for idx, joint in enumerate(self.robot.ordered_joints):
            joint.reset_position(joint_pos[idx], joint_vel[idx])

        # add some checks for simulation inaccuracies

        foot_pos = self.get_current_foot_pos()[:, -1]
        if np.any(foot_pos < -0.01):
            # print("foot step under the ground ", foot_pos)
            return curr_state

        self._apply_control(torque)
        next_state = self.get_current_full_state()

        if (
            np.abs(next_state[20] - curr_state[20]) > 0.003
            or np.abs(next_state[21] - curr_state[21]) > 0.005
        ):
            next_state = curr_state
        if np.abs(next_state[22] - curr_state[22]) > 0.001:
            next_state[22] = curr_state[22]
        return next_state