train_gan.py

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
import numpy as np
import argparse
import tensorflow as tf
tf.get_logger().setLevel('ERROR')
tf.config.optimizer.set_jit(True)
os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError:
        pass
from tensorflow import keras
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import AdamW
from tensorflow.keras.callbacks import (
    EarlyStopping, ReduceLROnPlateau, ModelCheckpoint, CSVLogger
)
from conditional_trajectory import load_defensive_encoder
from gan_trajectory import (
    build_trajectory_discriminator, build_gan_model, AdversarialLoss
)
from metrics_utils import seq2seq_masked_rmse_loss, seq2seq_masked_rmse_metric, seq2seq_pure_rmse_metric, create_physics_informed_loss_for_delta
from data_utils import load_and_split_data, generate_decoder_input
import joblib
import gc
def precompute_intent_to_memmap(encoder, X_P, X_S, save_path, batch_size=2048):
    num_samples = len(X_P)
    output_dim = 256
    if os.path.exists(save_path):
        try:
            mmap = np.memmap(save_path, dtype='float32', mode='r', shape=(num_samples, output_dim))
            if mmap.shape == (num_samples, output_dim):
                return mmap
        except Exception:
            try:
                os.remove(save_path)
            except Exception:
                pass
    fp = np.memmap(save_path, dtype='float32', mode='w+', shape=(num_samples, output_dim))
    X_S_input = X_S[:, :, :18] if X_S.shape[2] > 18 else X_S
    chunk_size = 50000
    num_chunks = (num_samples + chunk_size - 1) // chunk_size
    from tqdm import tqdm
    for i in tqdm(range(num_chunks), desc="precompute defense"):
        start_idx = i * chunk_size
        end_idx = min((i + 1) * chunk_size, num_samples)
        batch_X_P = X_P[start_idx:end_idx]
        batch_X_S = X_S_input[start_idx:end_idx]
        preds = encoder.predict([batch_X_P, batch_X_S], batch_size=batch_size, verbose=0)
        fp[start_idx:end_idx] = preds
        if i % 10 == 0:
            fp.flush()
        del preds, batch_X_P, batch_X_S
        gc.collect()
    fp.flush()
    del fp
    gc.collect()
    return np.memmap(save_path, dtype='float32', mode='r', shape=(num_samples, output_dim))
def create_dataset_from_arrays_gan(X_Players, X_Static, defensive_intent, player_indices, Y_Decoder_Input, Y_Target, adversarial_y_true, batch_size, shuffle=True, prefetch_buffer=2, chunk_size=100000, cycle_length=4):
    num_samples = len(X_Players)
    num_chunks = (num_samples + chunk_size - 1) // chunk_size
    if num_samples < chunk_size * 2:
        dataset = tf.data.Dataset.from_tensor_slices({
            'input_players': X_Players,
            'input_static': X_Static,
            'input_defensive_intent': defensive_intent,
            'input_player_index': player_indices,
            'input_decoder': Y_Decoder_Input,
            'output_trajectory': Y_Target,
            'output_adversarial': adversarial_y_true
        })
    else:
        def create_chunk_dataset(chunk_idx):
            def get_chunk_data(chunk_id):
                chunk_id_np = int(chunk_id.numpy())
                start_np = chunk_id_np * chunk_size
                end_np = min(start_np + chunk_size, num_samples)
                return (
                    X_Players[start_np:end_np],
                    X_Static[start_np:end_np],
                    defensive_intent[start_np:end_np],
                    player_indices[start_np:end_np],
                    Y_Decoder_Input[start_np:end_np],
                    Y_Target[start_np:end_np],
                    adversarial_y_true[start_np:end_np]
                )
            chunk_data_tuple = tf.py_function(
                get_chunk_data,
                [chunk_idx],
                (tf.float32, tf.float32, tf.float32, tf.float32, tf.float32, tf.float32, tf.float32)
            )
            chunk_data = {
                'input_players': chunk_data_tuple[0],
                'input_static': chunk_data_tuple[1],
                'input_defensive_intent': chunk_data_tuple[2],
                'input_player_index': chunk_data_tuple[3],
                'input_decoder': chunk_data_tuple[4],
                'output_trajectory': chunk_data_tuple[5],
                'output_adversarial': chunk_data_tuple[6]
            }
            chunk_data['input_players'].set_shape([None, X_Players.shape[1], X_Players.shape[2]])
            chunk_data['input_static'].set_shape([None, X_Static.shape[1], X_Static.shape[2]])
            chunk_data['input_defensive_intent'].set_shape([None, defensive_intent.shape[1]])
            chunk_data['input_player_index'].set_shape([None, player_indices.shape[1]])
            chunk_data['input_decoder'].set_shape([None, Y_Decoder_Input.shape[1], Y_Decoder_Input.shape[2]])
            chunk_data['output_trajectory'].set_shape([None, Y_Target.shape[1], Y_Target.shape[2]])
            chunk_data['output_adversarial'].set_shape([None, adversarial_y_true.shape[1]])
            return tf.data.Dataset.from_tensor_slices(chunk_data)
        chunk_indices = tf.data.Dataset.range(num_chunks)
        dataset = chunk_indices.interleave(
            create_chunk_dataset,
            cycle_length=cycle_length,
            block_length=1,
            num_parallel_calls=tf.data.AUTOTUNE
        )
    if shuffle:
        buffer_size = min(10000, num_samples)
        dataset = dataset.shuffle(buffer_size=buffer_size, reshuffle_each_iteration=True)
    dataset = dataset.batch(batch_size, drop_remainder=False)
    dataset = dataset.prefetch(prefetch_buffer)
    def map_to_tuple(x):
        inputs = (
            x['input_players'],
            x['input_static'],
            x['input_defensive_intent'],
            x['input_player_index'],
            x['input_decoder']
        )
        outputs = (
            x['output_trajectory'],
            x['output_adversarial']
        )
        return inputs, outputs
    dataset = dataset.map(map_to_tuple, num_parallel_calls=tf.data.AUTOTUNE)
    return dataset
class GANTrainingCallback(keras.callbacks.Callback):
    def __init__(self, X_train, Y_train, X_val, Y_val, Y_Target_90_val, 
                 discriminator, batch_size, adversarial_weight=0.1):
        super().__init__()
        self.X_train = X_train
        self.Y_train = Y_train
        self.X_val = X_val
        self.Y_val = Y_val
        self.Y_Target_90_val = Y_Target_90_val
        self.discriminator = discriminator
        self.batch_size = batch_size
        self.adversarial_weight = adversarial_weight
        self.best_value = float('inf')
    def on_epoch_end(self, epoch, logs=None):
        if logs is None:
            return
        gan_model = self.model
        gen_inputs = gan_model.inputs
        gen_output = gan_model.outputs[0]
        generator = Model(inputs=gen_inputs, outputs=gen_output, name="Generator_Extracted")
        train_size = len(self.X_train[0]) if isinstance(self.X_train, list) else len(self.X_train)
        sample_size = min(10000, train_size)
        sample_indices = np.random.choice(train_size, sample_size, replace=False)
        if isinstance(self.X_train, list):
            X_train_sample = [x[sample_indices] for x in self.X_train]
        else:
            X_train_sample = self.X_train[sample_indices]
        Y_train_sample = self.Y_train[sample_indices]
        fake_trajectories = generator.predict(X_train_sample, batch_size=self.batch_size, verbose=0)
        real_labels = np.ones((len(Y_train_sample), 1))
        fake_labels = np.zeros((len(fake_trajectories), 1))
        disc_X = np.concatenate([Y_train_sample, fake_trajectories], axis=0)
        disc_y = np.concatenate([real_labels, fake_labels], axis=0)
        indices = np.random.permutation(len(disc_X))
        disc_X = disc_X[indices]
        disc_y = disc_y[indices]
        disc_batch_size = min(self.batch_size, 512)
        num_batches = int(np.ceil(len(disc_X) / disc_batch_size))
        for i in range(num_batches):
            start_idx = i * disc_batch_size
            end_idx = min((i + 1) * disc_batch_size, len(disc_X))
            batch_X = disc_X[start_idx:end_idx]
            batch_y = disc_y[start_idx:end_idx]
            self.discriminator.train_on_batch(batch_X, batch_y)
        del fake_trajectories, disc_X, disc_y, X_train_sample, Y_train_sample
        gc.collect()
        val_predictions = generator.predict(self.X_val, batch_size=self.batch_size, verbose=0)
        val_predictions_90 = np.zeros((len(val_predictions), 90), dtype=np.float32)
        val_predictions_90[:, 0:30] = val_predictions[:, :, 0]
        val_predictions_90[:, 30:60] = val_predictions[:, :, 1]
        val_rmse = seq2seq_masked_rmse_metric(
            tf.constant(self.Y_Target_90_val, dtype=tf.float32),
            tf.constant(val_predictions_90, dtype=tf.float32)
        ).numpy()
        val_pure_rmse = seq2seq_pure_rmse_metric(
            tf.constant(self.Y_Target_90_val, dtype=tf.float32),
            tf.constant(val_predictions_90, dtype=tf.float32)
        ).numpy()
        if val_rmse < self.best_value:
            self.best_value = val_rmse
        logs['val_seq2seq_masked_rmse_metric'] = float(val_rmse)
        logs['val_seq2seq_pure_rmse_metric'] = float(val_pure_rmse)
        del val_predictions, val_predictions_90
        gc.collect()
def build_gan_model_from_hps(hps, generator, input_shape_players, input_shape_static, y_scaler=None):
    trajectory_weight = hps.get('trajectory_weight', 1.0)
    adversarial_weight = hps.get('adversarial_weight', 0.1)
    discriminator_lr = hps.get('discriminator_lr', 1e-4)
    vel_weight = hps.get('vel_weight', 0.1)
    accel_weight = hps.get('accel_weight', 0.05)
    if y_scaler is not None:
        trajectory_loss_fn = create_physics_informed_loss_for_delta(y_scaler, vel_weight, accel_weight)
    else:
        trajectory_loss_fn = seq2seq_masked_rmse_loss
    discriminator = build_trajectory_discriminator(
        input_shape=(30, 2),
        hidden_units=128,
        num_layers=2,
        dropout_rate=0.3
    )
    discriminator_optimizer = AdamW(learning_rate=discriminator_lr, weight_decay=1e-4)
    discriminator.compile(
        optimizer=discriminator_optimizer,
        loss='binary_crossentropy',
        metrics=['accuracy']
    )
    gan = build_gan_model(generator, discriminator)
    generator_optimizer = AdamW(learning_rate=1e-5, weight_decay=1e-4)
    gan.compile(
        optimizer=generator_optimizer,
        loss=[trajectory_loss_fn, AdversarialLoss()],
        loss_weights=[trajectory_weight, adversarial_weight],
        metrics=[[], []]
    )
    return gan, discriminator
def train_gan_trajectory(
    X_Players_train, X_Static_train, Y_Target_train, Y_Decoder_Input_train,
    X_Players_val, X_Static_val, Y_Target_val, Y_Decoder_Input_val,
    defensive_intent_train, defensive_intent_val,
    player_indices_train_onehot, player_indices_val_onehot,
    generator_path=None,
    epochs=200, batch_size=32,
    generator_lr=1e-5, discriminator_lr=1e-4,
    trajectory_weight=1.0, adversarial_weight=0.1,
    vel_weight=0.1, accel_weight=0.05,
    exp_name='exp15_stage3',
    model_dir='models',
    log_dir='logs',
    y_scaler=None
):
    print(f"\n{'='*60}")
    print("Stage3: GAN training")
    print(f"{'='*60}\n")
    print(f"Train: {len(X_Players_train)}")
    print(f"Val: {len(X_Players_val)}")
    input_shape_players = (X_Players_train.shape[1], X_Players_train.shape[2])
    input_shape_static = (X_Static_train.shape[1], X_Static_train.shape[2])
    print(f"Input dim: Players={input_shape_players}, Static={input_shape_static}")
    strategy = tf.distribute.get_strategy()
    with strategy.scope():
        print("Build generator")
        if generator_path and os.path.exists(generator_path):
            print(f"Load generator: {generator_path}")
            generator = keras.models.load_model(generator_path, compile=False, safe_mode=False)
        else:
            raise ValueError("generator_path required")
        print("Build discriminator")
        discriminator = build_trajectory_discriminator(
            input_shape=(30, 2),
            hidden_units=128,
            num_layers=2,
            dropout_rate=0.3
        )
    hps = {
        'trajectory_weight': trajectory_weight,
        'adversarial_weight': adversarial_weight,
        'discriminator_lr': discriminator_lr,
        'vel_weight': vel_weight,
        'accel_weight': accel_weight
    }
    gan, discriminator = build_gan_model_from_hps(hps, generator, input_shape_players, input_shape_static, y_scaler)
    print(f"Generator params: {generator.count_params():,}")
    print(f"Discriminator params: {discriminator.count_params():,}")
    print(f"GAN params: {gan.count_params():,}")
    Y_Target_train_reshaped = Y_Target_train.reshape(-1, 30, 2)
    Y_Target_val_reshaped = Y_Target_val.reshape(-1, 30, 2)
    Y_Target_90_val = np.zeros((len(Y_Target_val_reshaped), 90), dtype=np.float32)
    Y_Target_90_val[:, 0:30] = Y_Target_val_reshaped[:, :, 0]
    Y_Target_90_val[:, 30:60] = Y_Target_val_reshaped[:, :, 1]
    coords_finite_val = np.isfinite(Y_Target_val_reshaped)
    frame_valid_val = np.any(coords_finite_val, axis=2)
    Y_Target_90_val[:, 60:90] = frame_valid_val.astype(np.float32)
    adversarial_y_true_train = np.ones((len(Y_Target_train_reshaped), 1))
    adversarial_y_true_val = np.ones((len(Y_Target_val_reshaped), 1))
    train_dataset = create_dataset_from_arrays_gan(
        X_Players_train, X_Static_train, defensive_intent_train, player_indices_train_onehot,
        Y_Decoder_Input_train, Y_Target_train_reshaped, adversarial_y_true_train,
        batch_size, shuffle=True
    )
    val_dataset = create_dataset_from_arrays_gan(
        X_Players_val, X_Static_val, defensive_intent_val, player_indices_val_onehot,
        Y_Decoder_Input_val, Y_Target_val_reshaped, adversarial_y_true_val,
        batch_size, shuffle=False
    )
    X_train_list = [X_Players_train, X_Static_train, defensive_intent_train, player_indices_train_onehot, Y_Decoder_Input_train]
    X_val_list = [X_Players_val, X_Static_val, defensive_intent_val, player_indices_val_onehot, Y_Decoder_Input_val]
    os.makedirs(model_dir, exist_ok=True)
    os.makedirs(log_dir, exist_ok=True)
    model_filename = os.path.join(model_dir, f"gan_trajectory_{exp_name}_best.keras")
    log_filename = os.path.join(log_dir, f"gan_trajectory_{exp_name}.csv")
    initial_epoch = 0
    best_val_loss_from_log = None
    if os.path.exists(log_filename):
        try:
            import pandas as pd
            existing_df = pd.read_csv(log_filename)
            if len(existing_df) > 0:
                initial_epoch = int(existing_df['epoch'].max()) + 1
                if 'val_loss' in existing_df.columns:
                    best_val_loss_from_log = existing_df['val_loss'].min()
                if os.path.exists(model_filename):
                    gan.load_weights(model_filename)
        except Exception:
            pass
    gan_training_callback = GANTrainingCallback(
        X_train=X_train_list,
        Y_train=Y_Target_train_reshaped,
        X_val=X_val_list,
        Y_val=Y_Target_val_reshaped,
        Y_Target_90_val=Y_Target_90_val,
        discriminator=discriminator,
        batch_size=batch_size,
        adversarial_weight=adversarial_weight
    )
    callbacks = [
        gan_training_callback,
        EarlyStopping(
            monitor='val_seq2seq_masked_rmse_metric',
            patience=20,
            restore_best_weights=True,
            verbose=1
        ),
        ReduceLROnPlateau(
            monitor='val_seq2seq_masked_rmse_metric',
            factor=0.7,
            patience=10,
            min_lr=1e-8,
            cooldown=5,
            verbose=1
        ),
        ModelCheckpoint(
            model_filename,
            monitor='val_seq2seq_masked_rmse_metric',
            save_best_only=True,
            save_weights_only=False,
            verbose=1
        ),
        CSVLogger(log_filename, append=True),
        tf.keras.callbacks.TerminateOnNaN()
    ]
    if best_val_loss_from_log is not None:
        callbacks[2].best = float(best_val_loss_from_log)
    print("Training GAN")
    if initial_epoch > 0:
        print(f"Resume from epoch {initial_epoch} (total {epochs})")
    gan.fit(
        train_dataset,
        validation_data=val_dataset,
        epochs=epochs,
        initial_epoch=initial_epoch,
        callbacks=callbacks,
        verbose=1
    )
    return model_filename
def main():
    parser = argparse.ArgumentParser(description='Stage3: GAN trajectory training')
    parser.add_argument('--data_root', type=str, required=True,
                        help='raw CSV path')
    parser.add_argument('--processed_path', type=str, required=True,
                        help='processed path')
    parser.add_argument('--generator_path', type=str, required=True,
                        help='stage2 generator path')
    parser.add_argument('--defensive_encoder_path', type=str, default=None,
                        help='defense encoder path')
    parser.add_argument('--train_week_threshold', type=int, default=15,
                        help='train week threshold')
    parser.add_argument('--epochs', type=int, default=200,
                        help='epochs')
    parser.add_argument('--batch_size', type=int, default=32,
                        help='batch size')
    parser.add_argument('--generator_lr', type=float, default=1e-5,
                        help='generator lr')
    parser.add_argument('--discriminator_lr', type=float, default=1e-4,
                        help='discriminator lr')
    parser.add_argument('--trajectory_weight', type=float, default=1.0,
                        help='trajectory weight')
    parser.add_argument('--adversarial_weight', type=float, default=0.1,
                        help='adversarial weight')
    parser.add_argument('--vel_weight', type=float, default=0.1,
                        help='velocity weight')
    parser.add_argument('--accel_weight', type=float, default=0.05,
                        help='accel weight')
    parser.add_argument('--exp_name', type=str, default='exp15_stage3',
                        help='exp name')
    parser.add_argument('--model_dir', type=str, default='models',
                        help='model dir')
    parser.add_argument('--log_dir', type=str, default='logs',
                        help='log dir')
    parser.add_argument('--use_memmap', action='store_true', default=False,
                        help='use memmap')
    args = parser.parse_args()
    np.random.seed(42)
    tf.random.set_seed(42)
    print("Load data")
    try:
        X_P_train, X_S_train, Y_train, X_P_val, X_S_val, Y_val, player_indices_train, player_indices_val = load_and_split_data(
            args.data_root, 
            args.processed_path, 
            train_week_threshold=args.train_week_threshold,
            use_memmap=args.use_memmap,
            sumer_coverage_path=None,
            extend_y_target=False
        )
    except Exception as e:
        import traceback
        traceback.print_exc()
        return
    print("\n--- 正在将 player_indices 转换为 One-Hot 编码... ---")
    def to_one_hot(indices, num_classes=22):
        n_samples = len(indices)
        one_hot = np.zeros((n_samples, num_classes), dtype=np.float32)
        one_hot[np.arange(n_samples), indices] = 1.0
        return one_hot
    player_indices_train_onehot = to_one_hot(player_indices_train)
    player_indices_val_onehot = to_one_hot(player_indices_val)
    Y_Decoder_Input_train = generate_decoder_input(Y_train)
    Y_Decoder_Input_val = generate_decoder_input(Y_val)
    print(f"\n--- 正在加载Y_Target scaler... ---")
    y_scaler_path = os.path.join(args.processed_path, "y_target_scaler_exp15.joblib")
    if os.path.exists(y_scaler_path):
        y_scaler = joblib.load(y_scaler_path)
    else:
        y_scaler = None
    print("Init defense encoder")
    if args.defensive_encoder_path is None:
        generator_dir = os.path.dirname(args.generator_path)
        defensive_encoder_path_candidates = [
            os.path.join(generator_dir, "Defensive_Encoder_exp15_stage1_stage1_best.keras"),
            os.path.join(generator_dir, "Defensive_Encoder_best.keras"),
            os.path.join(generator_dir, "Defensive_Encoder_exp15_stage1.keras"),
        ]
        defensive_encoder_path = None
        for candidate_path in defensive_encoder_path_candidates:
            if candidate_path and os.path.exists(candidate_path):
                defensive_encoder_path = candidate_path
                break
        if defensive_encoder_path is None:
            raise FileNotFoundError(f"未找到防守意图编码器！尝试过的路径：\n" + "\n".join([f"  - {p}" for p in defensive_encoder_path_candidates if p]))
    else:
        defensive_encoder_path = args.defensive_encoder_path
    temp_encoder = load_defensive_encoder(defensive_encoder_path)
    cache_dir = os.path.join(args.processed_path, "_cache_intent_vectors")
    os.makedirs(cache_dir, exist_ok=True)
    train_intent_path = os.path.join(cache_dir, f"train_intent_w{args.train_week_threshold}.npy")
    val_intent_path = os.path.join(cache_dir, f"val_intent_w{args.train_week_threshold}.npy")
    defensive_intent_train = precompute_intent_to_memmap(
        temp_encoder, X_P_train, X_S_train, train_intent_path, batch_size=128
    )
    defensive_intent_val = precompute_intent_to_memmap(
        temp_encoder, X_P_val, X_S_val, val_intent_path, batch_size=128
    )
    del temp_encoder
    tf.keras.backend.clear_session()
    gc.collect()
    print("Start training")
    print("="*60)
    train_gan_trajectory(
        X_P_train, X_S_train, Y_train, Y_Decoder_Input_train,
        X_P_val, X_S_val, Y_val, Y_Decoder_Input_val,
        defensive_intent_train, defensive_intent_val,
        player_indices_train_onehot, player_indices_val_onehot,
        args.generator_path,
        args.epochs, args.batch_size,
        args.generator_lr, args.discriminator_lr,
        args.trajectory_weight, args.adversarial_weight,
        args.vel_weight, args.accel_weight,
        args.exp_name, args.model_dir, args.log_dir,
        y_scaler
    )
if __name__ == '__main__':
    main()