An adaptive fault diagnosis framework under class-imbalanced conditions based on contrastive augmented deep reinforcement learning

In practical scenarios, it is difficult to acquire fault data from rotating machinery, resulting in class-imbalanced problems in the fault diagnosis field. Training a fault diagnosis model directly on an imbalanced dataset may lead to overfitting for minority classes and skewness toward majority classes. Thus, we propose a fault diagnosis framework based on contrastive augmented deep reinforcement learning (CADRL) with two-stage training. In the pretraining stage, we obtain sample pairs based on the batch construction strategy to calculate the contrastive loss. Then, this loss is used to train a feature extraction model to reduce intraclass distances and increase interclass distances. During the fine-tuning stage, an adaptive reward function that updates with the sample label distribution is adopted in the fault diagnosis model. This function can balance the attention given by the model to different fault modes and improve fault diagnosis performance on imbalanced data without prior knowledge. Case studies conducted on two public datasets demonstrate that the pretraining stage can provide a well-trained feature extraction model, which can be merged into the proposed fault diagnosis model to achieve better fault diagnosis performance than that of other advanced models.