Imitation learning offers a powerful approach for enabling robotic arms to perform complex tasks in unstructured environments. However, many state-of-the-art methods are hindered by redundant input data, leading to inefficient training and limited trajectory prediction accuracy in complex tasks. To address these issues, this paper proposes KPT-O, a method for optimizing demonstrated trajectories by extracting keypoints. The method filters keypoints to streamline the learning data and optimizes their distribution to enhance prediction accuracy. To validate its performance, KPT-O was trained within a state-of-the-art framework and compared against leading methods on the HelloWorld and RoboTasks datasets. The results demonstrate that KPT-O not only significantly reduces training time but also achieves superior trajectory prediction accuracy. Furthermore, evaluations on a physical robot platform confirm the method′s effectiveness in real-world robotic arm tasks involving changes in both position and orientation.