ccrl

English | 中文说明

ccrl is a Python package implementing the method described in the paper:

Haoze Wu, Shisheng Zhong, Minghang Zhao, Xuyun Fu, Yongjian Zhang, Song Fu, "Continual contrastive reinforcement learning: Towards stronger agent for environment-aware fault diagnosis of aero-engines through long-term optimization under highly imbalance scenarios", Advanced Engineering Informatics, DOI: 10.1016/j.aei.2025.103297.

• Repository: https://github.com/haozewu/ccrl
• Paper DOI: 10.1016/j.aei.2025.103297

The package is published as ccrl on PyPI. This repository contains its source code and documentation.

The CCRL pipeline implemented here includes:

• LSTM autoencoder pretraining
• contrastive representation learning
• D3QN-based type identification
• imbalanced reward design
• repeated imbalanced cross-validation

Installation

pip install ccrl

For local development:

pip install .

Quick Start

from ccrl import CCRLConfig, run_ccrl_diag

config = CCRLConfig()
config.data.pretrain_label = "normal"

result = run_ccrl_diag(
    data_path="data/fault_dataset.pkl",
    config=config,
    repeats=1,
    seed=2024,
    log_dir="logs",
)

print(result.mean_f1, result.std_f1)

CLI

ccrl --data data/fault_dataset.pkl --repeats 1 --seed 2024 --log-dir logs

Input Data Format

The input file must be a pickle containing a dict[str, samples].

• each key is a class label, such as normal, fault_a, or bearing_outer
• each value must be convertible to a 3D array with shape (num_samples, seq_len, feature_dim)
• all classes must share the same seq_len and feature_dim

Minimal example:

{
    "normal": [
        [[0.1, 1.2], [0.2, 1.1], [0.3, 1.0]],
        [[0.0, 1.0], [0.1, 0.9], [0.2, 0.8]],
    ],
    "fault_a": [
        [[1.2, 0.1], [1.1, 0.2], [1.0, 0.3]],
        [[0.9, 0.0], [0.8, 0.1], [0.7, 0.2]],
    ],
    "fault_b": [
        [[0.5, 2.0], [0.6, 1.9], [0.7, 1.8]],
        [[0.4, 2.1], [0.5, 2.0], [0.6, 1.9]],
    ],
}

This example has:

• num_classes = 3
• seq_len = 3
• feature_dim = 2

Save it as pickle:

import pickle

data = {
    "normal": [
        [[0.1, 1.2], [0.2, 1.1], [0.3, 1.0]],
        [[0.0, 1.0], [0.1, 0.9], [0.2, 0.8]],
    ],
    "fault_a": [
        [[1.2, 0.1], [1.1, 0.2], [1.0, 0.3]],
        [[0.9, 0.0], [0.8, 0.1], [0.7, 0.2]],
    ],
    "fault_b": [
        [[0.5, 2.0], [0.6, 1.9], [0.7, 1.8]],
        [[0.4, 2.1], [0.5, 2.0], [0.6, 1.9]],
    ],
}

with open("data/fault_dataset.pkl", "wb") as f:
    pickle.dump(data, f)

Configuration Notes

from ccrl import CCRLConfig

config = CCRLConfig()
config.data.class_order = ["normal", "fault_a", "fault_b"]
config.data.pretrain_label = "normal"
config.data.test_samples_per_class = 2

• class_order controls label encoding order and reporting order
• pretrain_label chooses which class is used for AE pretraining
• test_samples_per_class controls how many samples per class are drawn into the test split

Notes

This repository is the source distribution of the ccrl package, intended for research, reproduction, and further development. It is not an official release from the paper authors unless explicitly stated otherwise. Because this codebase has been extracted, refactored, and generalized for reuse, it may contain engineering adaptations, simplifications, or implementation mistakes. If anything in this repository conflicts with the paper, the original paper should be treated as the authoritative reference.

License

This project is released under the Apache License 2.0. See LICENSE.