Machine learning en producción: Optuna, FT-Transformer y CQL

Lo que separa un prototipo de un sistema

Un modelo que funciona en un notebook y uno que funciona en producción se parecen poco. El segundo tiene que responder en milisegundos, aguantar clases muy desbalanceadas, decidir sin poder experimentar con clientes reales y seguir sirviendo cuando los datos cambian. Una búsqueda en rejilla sobre un gradient boosting sin ajustar no llega ahí.

Este artículo cubre tres técnicas que atacan tres de esos problemas:

• Optimización bayesiana con poda (TPE en Optuna): encontrar buenos hiperparámetros cuando cada evaluación cuesta minutos y el presupuesto de cómputo es finito.
• Transformers para datos tabulares (FT-Transformer): cuándo tiene sentido reemplazar los árboles por atención, y cuándo no.
• Aprendizaje por refuerzo offline (Conservative Q-Learning): aprender una política de decisión solo con el histórico, sin que el modelo se invente acciones milagrosas.

Cierro con despliegue: exportación a ONNX y monitoreo de drift.

Todo el código lo ejecuté, y las salidas están copiadas de la terminal. Versiones: Python 3.13, Optuna 5.0, CatBoost 1.2.10, PyTorch 2.14 y ONNX Runtime 1.30, en CPU.

1. Optimización bayesiana con TPE

La búsqueda en rejilla prueba todas las combinaciones y la aleatoria prueba puntos al azar. Las dos ignoran lo que ya aprendieron: el trial número 40 se elige igual de a ciegas que el primero. Cuando entrenar un modelo cuesta minutos y el espacio tiene seis o siete dimensiones, eso sale caro.

La optimización bayesiana usa los resultados anteriores para elegir el siguiente punto. Mantiene un modelo sustituto de la función objetivo $f(x)$ —la métrica de validación en función de los hiperparámetros— y elige el siguiente $x$ donde ese modelo promete más. El sustituto clásico es un proceso gaussiano, que modela $p(y \mid x)$ directamente, pero su costo crece con el cubo del número de observaciones: $O(N^3)$.

Cómo decide TPE

El Tree-structured Parzen Estimator (Bergstra et al., 2011) le da la vuelta al problema. En lugar de modelar $p(y \mid x)$, modela $p(x \mid y)$. Ordena las observaciones por su pérdida $y$, elige un umbral $y^*$ tal que una fracción $\gamma$ quede por debajo, y ajusta dos densidades:

$\ell(x)$ es la densidad de los hiperparámetros que dieron los mejores resultados, y $g(x)$ la del resto. Con eso, la mejora esperada (Expected Improvement) sobre el umbral queda así:

Un matiz que cambia resultados: por defecto, TPE modela cada hiperparámetro por separado y no ve interacciones, como que un learning_rate bajo necesite más árboles. Optuna tiene TPESampler(multivariate=True) para modelarlas juntas.

La poda es otra cosa

TPE decide qué probar. La poda decide cuándo dejar de probar algo que va mal. Son mecanismos independientes, y en Optuna se configuran por separado: el sampler (TPE, CMA-ES, un proceso gaussiano) y el pruner (mediana, Hyperband, successive halving).

El MedianPruner compara el valor intermedio de un trial en un paso con la mediana de los trials anteriores en ese mismo paso. Si queda por debajo, lo corta. Para que funcione, el trial tiene que poder informar valores intermedios comparables: por ejemplo, la métrica después de cada fold de la validación cruzada.

2. Transformers para datos tabulares: FT-Transformer

En datos tabulares, los árboles con gradient boosting (XGBoost, LightGBM, CatBoost) siguen siendo la opción por defecto, y con razón. La comparación más citada, de Grinsztajn y colaboradores (2022), mostró que siguen ganando a las redes neuronales en la mayoría de datasets medianos. El propio artículo del FT-Transformer (Gorishniy et al., 2021) concluye que no hay un ganador universal.

Entonces, ¿para qué un Transformer? Para los casos donde la red neuronal aporta algo que un árbol no tiene: cuando el modelo tabular tiene que entrenarse junto con otras partes neuronales (texto, imágenes), cuando necesitas las representaciones vectoriales de las filas para otra tarea, o cuando quieres seguir entrenando con gradiente a medida que llegan datos.

La arquitectura

La idea es convertir cada columna en un token. Para una columna numérica $x_j$, el tokenizador aprende un vector de pesos y uno de sesgo propios de esa columna:

Las columnas categóricas pasan por una tabla de embeddings, también con un sesgo por columna. A la secuencia de tokens se le añade un token especial [CLS], y todo pasa por bloques Transformer con normalización previa (pre-norm):

La predicción sale del estado final del token [CLS]. Si vienes de cómo funciona la atención, aquí la atención compara columnas entre sí en lugar de palabras: cada columna decide qué otras columnas mirar para esta fila en particular.

3. Aprendizaje por refuerzo offline y Conservative Q-Learning

En finanzas, salud, recomendación o precios, aprender una política probando en vivo (el refuerzo online, como PPO o DQN) es caro o peligroso: nadie quiere que un modelo aprenda a fijar precios experimentando con clientes reales. El refuerzo offline aprende solo con un histórico fijo de decisiones pasadas:

El problema: acciones que nunca viste

El Q-learning estima $Q(s, a)$, el valor de tomar la acción $a$ en el estado $s$, y en cada actualización usa $\max_{a'} Q(s', a')$. Ese máximo recorre todas las acciones, también las que el histórico nunca probó. Para esas, la red no tiene datos: su valor es una extrapolación. Y el máximo elige justamente la extrapolación más optimista, que se vuelve a usar como objetivo en la siguiente actualización. El error no se corrige, se acumula.

Online, el agente probaría esa acción, vería que era mala y corregiría. Offline, no hay forma de probarla.

La solución de CQL

Conservative Q-Learning (Kumar et al., 2020) añade un término que empuja hacia abajo los valores $Q$ de todas las acciones y hacia arriba los de las acciones que sí están en el histórico:

$\hat{\pi}_\beta$ es la política que generó el histórico, y $\hat{\mathcal{B}}^{\pi}\hat{Q}^{k}$ el objetivo de Bellman calculado con la red de la iteración anterior. El segundo término es el Q-learning de siempre. El primero es el nuevo: el log-sum-exp es un máximo suave sobre todas las acciones, y restarle el valor de las acciones del histórico penaliza a las que la red valora alto sin evidencia.

El artículo original demuestra que, con $\alpha$ suficientemente grande y bajo sus supuestos, el valor esperado de la política según la $Q$ aprendida es una cota inferior de su valor real. Conviene leerlo con cuidado: con redes neuronales y datos finitos esa garantía no se hereda tal cual. Lo que sí se hereda es el sesgo deliberado hacia abajo, que en la práctica es lo que evita las decisiones temerarias. Y tiene un costo: con $\alpha$ demasiado alto, la política se limita a copiar el histórico.

Las librerías

Ejercicio 1: TPE con poda en CatBoost, con fraude desbalanceado

Un clasificador de fraude sobre 20.000 transacciones sintéticas, con un 3 % de positivos. El 20 % se aparta como test antes de empezar y no se toca durante la búsqueda: es la única cifra que no quedó optimizada a mano.

import time

import numpy as np
import optuna
import pandas as pd
from catboost import CatBoostClassifier
from sklearn.datasets import make_classification
from sklearn.metrics import average_precision_score, roc_auc_score
from sklearn.model_selection import StratifiedKFold, train_test_split

N_FOLDS = 3


def dataset_fraude(n=20_000, n_features=25, seed=42):
    X, y = make_classification(
        n_samples=n, n_features=n_features, n_informative=18, n_redundant=5,
        weights=[0.97, 0.03], flip_y=0.01, random_state=seed,
    )
    return pd.DataFrame(X, columns=[f"f{i}" for i in range(n_features)]), pd.Series(y)


def objetivo(X, y):
    folds = list(StratifiedKFold(N_FOLDS, shuffle=True, random_state=42).split(X, y))

    def evaluar(trial: optuna.Trial) -> float:
        params = {
            "iterations": 1000,
            "learning_rate": trial.suggest_float("learning_rate", 1e-2, 0.3, log=True),
            "depth": trial.suggest_int("depth", 4, 8),
            "l2_leaf_reg": trial.suggest_float("l2_leaf_reg", 0.1, 10.0, log=True),
            "bootstrap_type": "Bernoulli",
            "subsample": trial.suggest_float("subsample", 0.5, 1.0),
            "random_strength": trial.suggest_float("random_strength", 1e-2, 10.0, log=True),
            # Reponderar las clases es una decisión, no un dogma: que la busque también.
            "auto_class_weights": trial.suggest_categorical("auto_class_weights", ["None", "SqrtBalanced", "Balanced"]),
            "eval_metric": "PRAUC",
            "early_stopping_rounds": 50,
            "random_seed": 42,
            "verbose": False,
        }
        if params["auto_class_weights"] == "None":
            del params["auto_class_weights"]

        puntajes, iteraciones = [], []
        for k, (tr, va) in enumerate(folds):
            modelo = CatBoostClassifier(**params)
            modelo.fit(X.iloc[tr], y.iloc[tr], eval_set=(X.iloc[va], y.iloc[va]))
            puntajes.append(average_precision_score(y.iloc[va], modelo.predict_proba(X.iloc[va])[:, 1]))
            iteraciones.append(modelo.get_best_iteration() + 1)
            # Un valor por fold, con el fold como paso: la media acumulada es
            # comparable entre trials en el mismo punto de su evaluación.
            trial.report(float(np.mean(puntajes)), step=k)
            if trial.should_prune():
                raise optuna.TrialPruned()
        # El modelo final no tendrá conjunto de validación: guarda cuántos árboles necesitó.
        trial.set_user_attr("iteraciones", int(np.mean(iteraciones)))
        return float(np.mean(puntajes))

    return evaluar


if __name__ == "__main__":
    X, y = dataset_fraude()
    X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.2, stratify=y, random_state=42)

    base = CatBoostClassifier(random_seed=42, verbose=False).fit(X_tr, y_tr)
    p_base = base.predict_proba(X_te)[:, 1]

    optuna.logging.set_verbosity(optuna.logging.WARNING)
    study = optuna.create_study(
        direction="maximize",
        sampler=optuna.samplers.TPESampler(seed=42),
        pruner=optuna.pruners.MedianPruner(n_startup_trials=5, n_warmup_steps=0),
    )
    t0 = time.time()
    study.optimize(objetivo(X_tr, y_tr), n_trials=30)
    minutos = (time.time() - t0) / 60

    estados = [t.state for t in study.trials]
    completos = estados.count(optuna.trial.TrialState.COMPLETE)
    podados = estados.count(optuna.trial.TrialState.PRUNED)
    folds_usados = sum(len(t.intermediate_values) for t in study.trials)

    best = {**study.best_params}
    if best["auto_class_weights"] == "None":
        del best["auto_class_weights"]
    final = CatBoostClassifier(
        iterations=study.best_trial.user_attrs["iteraciones"], bootstrap_type="Bernoulli",
        random_seed=42, verbose=False, **best,
    ).fit(X_tr, y_tr)
    p_final = final.predict_proba(X_te)[:, 1]

    print(f"Trials: {completos} completos, {podados} podados · {minutos:.1f} min")
    print(f"Folds entrenados: {folds_usados} de {len(study.trials) * N_FOLDS} posibles")
    print(f"PR-AUC en validación cruzada (mejor trial): {study.best_value:.4f}")
    print(f"Test · CatBoost por defecto : PR-AUC {average_precision_score(y_te, p_base):.4f} · ROC-AUC {roc_auc_score(y_te, p_base):.4f}")
    print(f"Test · CatBoost optimizado  : PR-AUC {average_precision_score(y_te, p_final):.4f} · ROC-AUC {roc_auc_score(y_te, p_final):.4f}")
    print(f"Prevalencia en test: {y_te.mean():.2%} (PR-AUC de un modelo al azar)")

Trials: 13 completos, 17 podados · 14.3 min
Folds entrenados: 56 de 90 posibles
PR-AUC en validación cruzada (mejor trial): 0.7866
Test · CatBoost por defecto : PR-AUC 0.7598 · ROC-AUC 0.9087
Test · CatBoost optimizado  : PR-AUC 0.7896 · ROC-AUC 0.9072
Prevalencia en test: 3.50% (PR-AUC de un modelo al azar)

Qué muestra

La métrica importa más que el optimizador. La búsqueda subió la PR-AUC en test de 0,760 a 0,790, y la ROC-AUC ni se enteró: bajó de 0,909 a 0,907. Con un 3,5 % de positivos, la ROC-AUC está dominada por los miles de negativos fáciles de ordenar y apenas distingue un modelo de otro. La PR-AUC mide lo que le importa a quien revisa alertas de fraude: de lo que el modelo marca, cuánto es fraude de verdad. Su línea base no es 0,5 sino la prevalencia: un modelo al azar saca 0,035.

La poda ahorró un 38 % del cómputo. De 90 folds posibles se entrenaron 56. Los primeros cinco trials se completan siempre (n_startup_trials=5) para tener una mediana con la que comparar; a partir de ahí, un trial cuya media tras el primer o segundo fold queda por debajo de la mediana se corta.

El número de árboles se hereda. Cada fold para con early stopping, y el modelo final, que se entrena sin conjunto de validación, usa la media de esas iteraciones. Entrenarlo con las 1.000 iteraciones del espacio de búsqueda sería otro modelo, no el que se validó.

for step in range(1, 501, 50):
    model.fit(train_pool, eval_set=val_pool, early_stopping_rounds=30)
    trial.report(score, step)
    if trial.should_prune():
        raise optuna.TrialPruned()

Ejercicio 2: FT-Transformer contra CatBoost, en datos reales

El dataset es Adult, un clásico del aprendizaje supervisado tabular: 45.222 personas del censo de Estados Unidos, con 6 columnas numéricas y 8 categóricas, y la pregunta de si ganan más de 50.000 dólares al año. La implementación sigue el artículo original: pesos y sesgos propios por columna, sesgo también en las categóricas, bloques pre-norm sin la primera normalización en el primer bloque, FFN con activación ReGLU e inicialización uniforme en $\pm 1/\sqrt{d}$.

import copy
import math
import time

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from catboost import CatBoostClassifier
from sklearn.datasets import fetch_openml
from sklearn.metrics import accuracy_score, roc_auc_score
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import OrdinalEncoder, QuantileTransformer


# ——— 1. Tokenizador: cada columna se convierte en un vector de dimensión d ———
class FeatureTokenizer(nn.Module):
    def __init__(self, n_num: int, cardinalidades: list[int], d: int):
        super().__init__()
        s = 1 / math.sqrt(d)
        # e_j = x_j · w_j + b_j: un w y un b propios por columna numérica.
        self.w_num = nn.Parameter(torch.empty(n_num, d).uniform_(-s, s))
        self.b_num = nn.Parameter(torch.empty(n_num, d).uniform_(-s, s))
        # Categóricas: una tabla de embeddings compartida, con un desplazamiento por columna.
        self.register_buffer("offsets", torch.tensor([0, *np.cumsum(cardinalidades)[:-1]]))
        self.emb_cat = nn.Embedding(sum(cardinalidades), d)
        nn.init.uniform_(self.emb_cat.weight, -s, s)
        self.b_cat = nn.Parameter(torch.empty(len(cardinalidades), d).uniform_(-s, s))
        self.cls = nn.Parameter(torch.empty(1, 1, d).uniform_(-s, s))

    def forward(self, x_num, x_cat):
        num = x_num.unsqueeze(-1) * self.w_num + self.b_num          # [B, n_num, d]
        cat = self.emb_cat(x_cat + self.offsets) + self.b_cat        # [B, n_cat, d]
        cls = self.cls.expand(len(x_num), -1, -1)                    # [B, 1, d]
        # Sin codificación posicional: la identidad de cada columna ya está en su w_j.
        return torch.cat([cls, num, cat], dim=1)


# ——— 2. Bloque pre-norm con FFN ReGLU, como en el paper ———
class ReGLU(nn.Module):
    def forward(self, x):
        a, b = x.chunk(2, dim=-1)
        return a * F.relu(b)


class Bloque(nn.Module):
    def __init__(self, d, heads, primer_bloque: bool, d_ffn: int):
        super().__init__()
        # El paper quita la primera LayerNorm del primer bloque: en sus experimentos rendía mejor así.
        self.norm_attn = nn.Identity() if primer_bloque else nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, heads, dropout=0.2, batch_first=True)
        self.norm_ffn = nn.LayerNorm(d)
        self.ffn = nn.Sequential(nn.Linear(d, 2 * d_ffn), ReGLU(), nn.Dropout(0.1), nn.Linear(d_ffn, d))

    def forward(self, z):
        h = self.norm_attn(z)
        z = z + self.attn(h, h, h, need_weights=False)[0]
        return z + self.ffn(self.norm_ffn(z))


class FTTransformer(nn.Module):
    def __init__(self, n_num, cardinalidades, d=96, bloques=3, heads=8):
        super().__init__()
        self.tokenizer = FeatureTokenizer(n_num, cardinalidades, d)
        d_ffn = int(d * 4 / 3)
        self.bloques = nn.ModuleList(Bloque(d, heads, i == 0, d_ffn) for i in range(bloques))
        self.head = nn.Sequential(nn.LayerNorm(d), nn.ReLU(), nn.Linear(d, 1))

    def forward(self, x_num, x_cat):
        z = self.tokenizer(x_num, x_cat)
        for bloque in self.bloques:
            z = bloque(z)
        return self.head(z[:, 0]).squeeze(-1)  # la predicción sale del token [CLS]


# ——— 3. Datos: Adult (48.842 personas, ¿gana más de 50.000 dólares al año?) ———
def cargar_adult():
    df = fetch_openml("adult", version=2, as_frame=True).frame.dropna()
    y = (df.pop("class") == ">50K").astype(int).to_numpy()
    num_cols = df.select_dtypes("number").columns.tolist()
    cat_cols = [c for c in df.columns if c not in num_cols]
    return df, y, num_cols, cat_cols


if __name__ == "__main__":
    torch.manual_seed(0)
    df, y, num_cols, cat_cols = cargar_adult()
    idx_tr, idx_te = train_test_split(np.arange(len(df)), test_size=0.2, stratify=y, random_state=0)
    idx_tr, idx_va = train_test_split(idx_tr, test_size=0.15, stratify=y[idx_tr], random_state=0)

    # Numéricas a cuantiles normales (como en el paper); categóricas a enteros. Todo ajustado solo en train.
    qt = QuantileTransformer(output_distribution="normal", random_state=0).fit(df.iloc[idx_tr][num_cols])
    oe = OrdinalEncoder(handle_unknown="use_encoded_value", unknown_value=-1).fit(df.iloc[idx_tr][cat_cols])
    Xn = torch.tensor(qt.transform(df[num_cols]), dtype=torch.float32)
    Xc = torch.tensor(oe.transform(df[cat_cols]) + 1, dtype=torch.long)  # 0 = categoría no vista
    card = [len(c) + 1 for c in oe.categories_]
    Y = torch.tensor(y, dtype=torch.float32)

    modelo = FTTransformer(len(num_cols), card)
    opt = torch.optim.AdamW(modelo.parameters(), lr=1e-4, weight_decay=1e-5)
    print(f"FT-Transformer: {sum(p.numel() for p in modelo.parameters()):,} parámetros")

    def predecir(idx):
        modelo.eval()
        with torch.no_grad():
            return torch.sigmoid(modelo(Xn[idx], Xc[idx])).numpy()

    mejor, mejor_estado, paciencia, t0 = 0.0, None, 0, time.time()
    for epoca in range(1, 101):
        modelo.train()
        for lote in torch.randperm(len(idx_tr)).split(256):
            i = idx_tr[lote.numpy()]
            loss = F.binary_cross_entropy_with_logits(modelo(Xn[i], Xc[i]), Y[i])
            opt.zero_grad()
            loss.backward()
            opt.step()
        auc = roc_auc_score(y[idx_va], predecir(idx_va))
        if auc > mejor:
            mejor, mejor_estado, paciencia = auc, copy.deepcopy(modelo.state_dict()), 0
        else:
            paciencia += 1
        print(f"Época {epoca:2d} | ROC-AUC validación {auc:.4f}")
        if paciencia == 5:  # parada temprana
            break
    modelo.load_state_dict(mejor_estado)
    p_ft = predecir(idx_te)
    t_ft = time.time() - t0

    t0 = time.time()
    cb = CatBoostClassifier(iterations=2000, early_stopping_rounds=100, random_seed=0, verbose=False, cat_features=cat_cols)
    cb.fit(df.iloc[idx_tr], y[idx_tr], eval_set=(df.iloc[idx_va], y[idx_va]))
    p_cb = cb.predict_proba(df.iloc[idx_te])[:, 1]
    t_cb = time.time() - t0

    for nombre, p, t in [("FT-Transformer", p_ft, t_ft), ("CatBoost (por defecto)", p_cb, t_cb)]:
        print(f"Test · {nombre:<24} ROC-AUC {roc_auc_score(y[idx_te], p):.4f} · "
              f"accuracy {accuracy_score(y[idx_te], p > 0.5):.4f} · entrenamiento {t:.0f} s")
    torch.save({"estado": modelo.state_dict(), "n_num": len(num_cols), "card": card}, "ft_adult.pt")

FT-Transformer: 236,833 parámetros
Época  1 | ROC-AUC validación 0.8962
Época  2 | ROC-AUC validación 0.8990
...
Época 21 | ROC-AUC validación 0.9102
...
Época 26 | ROC-AUC validación 0.9091
Test · FT-Transformer           ROC-AUC 0.9115 · accuracy 0.8515 · entrenamiento 670 s
Test · CatBoost (por defecto)   ROC-AUC 0.9248 · accuracy 0.8664 · entrenamiento 83 s

Qué muestra

CatBoost sin ajustar le gana al FT-Transformer. Más ROC-AUC (0,925 frente a 0,912), más accuracy (0,866 frente a 0,852) y ocho veces menos tiempo de entrenamiento en CPU. No es un fallo de la implementación: es lo que dice la literatura para un dataset de este tamaño, y es la razón para no reemplazar un gradient boosting que funciona. El FT-Transformer que ves aquí tampoco está ajustado; con una búsqueda como la del ejercicio 1 la distancia se acorta, pero en Adult no es habitual que se invierta.

Lo que sí te da el Transformer es lo que un árbol no tiene: el estado del token [CLS] es un vector de 96 dimensiones por fila, que puedes reutilizar como representación en otro modelo, y la red entera se puede entrenar junto con un codificador de texto o imágenes. Si tu problema es solo "predecir esta columna con estas otras", usa árboles.

Tres detalles que la implementación ingenua suele perder. nn.TransformerEncoderLayer usa post-norm por defecto, así que la fórmula de arriba solo se cumple con norm_first=True o con un bloque propio como este. No acepta ReGLU como activación, porque ReGLU parte en dos la salida de la primera capa lineal y cambia sus dimensiones: torch.nn.functional no tiene una reglu, y comprobar si existe para elegirla deja la red en ReLU sin avisar. Y el [CLS] no necesita posición: ninguna la necesita.

Ejercicio 3: CQL contra Q-learning offline, en un simulador de precios

Para ver si la penalización conservadora sirve, hace falta algo que en producción no existe: saber cuánto vale de verdad la política aprendida. Por eso este ejercicio usa un simulador de precios con la verdad conocida. Entrenar con estados y recompensas aleatorias no enseña nada: no hay nada que aprender, y no hay forma de ver si la política es buena.

El escenario: una tienda decide cada día qué descuento ofrecer a cada cliente, entre 0, 10, 20 y 40 %. La demanda sube con el descuento, pero con un 40 % se vende a pérdida. Y un descuento hoy adelanta compras: baja la demanda de mañana. El histórico lo generó un equipo de precios que usó el 0, el 10 y el 20 %, y nunca probó el 40 %.

El descuento entra a la red como un número, igual que un precio en la práctica. Eso es lo que permite extrapolar, y lo que hace real el problema.

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

DESCUENTOS = np.array([0.00, 0.10, 0.20, 0.40])
PRECIO, COSTO = 100.0, 62.0
STATE_DIM, N_ACTIONS, GAMMA = 4, len(DESCUENTOS), 0.9
W_DEMANDA = np.array([0.8, -0.5, 0.3, 0.6])


# ——— 1. El mundo (en producción no lo conoces; aquí sí, para poder evaluar) ———
def prob_compra(s, a):
    """La demanda depende del contexto y sube con el descuento."""
    return 1 / (1 + np.exp(-(s @ W_DEMANDA - 2.5 + 15.0 * DESCUENTOS[a])))


def recompensa_esperada(s, a):
    margen = PRECIO * (1 - DESCUENTOS[a]) - COSTO  # con 40 % de descuento, se vende a pérdida
    return prob_compra(s, a) * margen


def paso(s, a, rng):
    compra = rng.random(len(s)) < prob_compra(s, a)
    r = compra * (PRECIO * (1 - DESCUENTOS[a]) - COSTO)
    # Un descuento hoy adelanta compras: baja la demanda de mañana (dimensión 3).
    s2 = 0.8 * s + rng.normal(0, 0.3, s.shape)
    s2[:, 3] -= 0.5 * DESCUENTOS[a]
    return r, s2


# ——— 2. El histórico: lo que decidió el equipo de precios durante un año ———
def politica_historica(n, rng):
    # Nunca se probó el 40 %: es la acción fuera de distribución.
    return rng.choice(N_ACTIONS, size=n, p=[0.40, 0.35, 0.25, 0.00])


def dataset_offline(n=50_000, seed=0):
    rng = np.random.default_rng(seed)
    s = rng.normal(0, 1, (n, STATE_DIM))
    a = politica_historica(n, rng)
    r, s2 = paso(s, a, rng)
    t = lambda x, dt=torch.float32: torch.as_tensor(x, dtype=dt)
    return t(s), t(a, torch.long), t(r).unsqueeze(1), t(s2)


# ——— 3. Red Q y agente ———
class QNetwork(nn.Module):
    """Q(s, descuento): el descuento entra como número, igual que un precio en la práctica."""

    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(STATE_DIM + 1, 128), nn.ReLU(),
            nn.Linear(128, 128), nn.ReLU(),
            nn.Linear(128, 1),
        )
        self.register_buffer("descuentos", torch.tensor(DESCUENTOS, dtype=torch.float32))

    def forward(self, s):
        """Devuelve Q para las cuatro acciones: [B, N_ACTIONS]."""
        s_rep = s.unsqueeze(1).expand(-1, N_ACTIONS, -1)
        d_rep = self.descuentos.view(1, -1, 1).expand(len(s), -1, -1)
        return self.net(torch.cat([s_rep, d_rep], dim=-1)).squeeze(-1)


def entrenar(datos, alpha: float, pasos=15_000, batch=256, seed=0):
    torch.manual_seed(seed)
    s, a, r, s2 = datos
    q, q_target = QNetwork(), QNetwork()
    q_target.load_state_dict(q.state_dict())
    opt = torch.optim.Adam(q.parameters(), lr=3e-4)
    for i in range(pasos):
        idx = torch.randint(0, len(s), (batch,))
        with torch.no_grad():
            target = r[idx] + GAMMA * q_target(s2[idx]).max(dim=1, keepdim=True).values
        q_all = q(s[idx])
        q_data = q_all.gather(1, a[idx].unsqueeze(1))
        td = F.mse_loss(q_data, target)
        # CQL(H): empuja hacia abajo todas las acciones (log-sum-exp) y hacia
        # arriba la que de verdad está en el dataset.
        conservador = (torch.logsumexp(q_all, dim=1, keepdim=True) - q_data).mean()
        loss = td + alpha * conservador
        opt.zero_grad()
        loss.backward()
        opt.step()
        if i % 500 == 0:
            q_target.load_state_dict(q.state_dict())
    return q


# ——— 4. Evaluación contra la verdad del simulador ———
def valor_real(politica, n=20_000, horizonte=40, seed=1):
    rng = np.random.default_rng(seed)
    s = rng.normal(0, 1, (n, STATE_DIM))
    total, desc = np.zeros(n), 1.0
    for _ in range(horizonte):
        a = politica(s)
        r, s = paso(s, a, rng)
        total += desc * r
        desc *= GAMMA
    return total.mean()


def greedy(q):
    def pol(s):
        with torch.no_grad():
            return q(torch.as_tensor(s, dtype=torch.float32)).argmax(1).numpy()
    return pol


if __name__ == "__main__":
    datos = dataset_offline()
    s0 = np.random.default_rng(2).normal(0, 1, (20_000, STATE_DIM))
    rng_hist = np.random.default_rng(3)
    hist = lambda s: politica_historica(len(s), rng_hist)
    print(f"Política histórica          valor real {valor_real(hist):7.2f}")
    r_media = [recompensa_esperada(s0, np.full(len(s0), k)).mean() for k in range(N_ACTIONS)]
    print("Recompensa inmediata real por acción:", " ".join(f"{d:.0%}={r:.2f}" for d, r in zip(DESCUENTOS, r_media)))
    for alpha in [0.0, 1.0, 5.0]:
        q = entrenar(datos, alpha)
        with torch.no_grad():
            q_s0 = q(torch.as_tensor(s0, dtype=torch.float32))
        acciones = q_s0.argmax(1).numpy()
        estimado = q_s0.max(1).values.mean().item()
        reparto = np.bincount(acciones, minlength=N_ACTIONS) / len(acciones)
        print(
            f"alpha={alpha:<4} Q estimada {estimado:7.2f} | valor real {valor_real(greedy(q)):7.2f} | "
            f"elige 0/10/20/40 %: {' '.join(f'{p:4.0%}' for p in reparto)} | Q(40 %) {q_s0[:, 3].mean().item():7.2f}"
        )

Política histórica          valor real   65.76
Recompensa inmediata real por acción: 0%=4.34 10%=8.72 20%=10.76 40%=-1.90
alpha=0.0  Q estimada  162.07 | valor real  -18.65 | elige 0/10/20/40 %:   0%   0%   0% 100% | Q(40 %)  162.07
alpha=1.0  Q estimada   92.21 | valor real  101.50 | elige 0/10/20/40 %:   0%  19%  80%   0% | Q(40 %)   83.93
alpha=5.0  Q estimada   75.16 | valor real  100.20 | elige 0/10/20/40 %:   2%  32%  66%   0% | Q(40 %)   69.09

Qué muestra

Sin penalización, el agente elige lo único que nunca se probó. Con alpha=0, que es Q-learning offline sin más, la política pone el 40 % de descuento al 100 % de los clientes. Cree que eso vale 162. En realidad vale −18,65: pierde dinero, y queda muy por debajo del equipo de precios al que quería mejorar (65,76).

La causa es la de la teoría. En el histórico, pasar del 10 al 20 % de descuento siempre dio más. La red aprendió esa pendiente y la prolongó hasta el 40 %, donde no tenía ni un dato. El máximo de Bellman eligió esa extrapolación, la usó como objetivo y la infló en cada actualización. Nada en los datos podía corregirla, porque nadie probó nunca el 40 %.

Con CQL, la política mejora el histórico sin inventar. Con alpha=1 vale 101,50, un 54 % más que el equipo de precios, y usa sobre todo el 20 %, la mejor acción que sí aparece en los datos. Nunca elige el 40 %.

La estimación queda por debajo de la realidad. Con alpha=1 la red cree que su política vale 92,21 y vale 101,50; con alpha=5, 75,16 frente a 100,20. Es la cota inferior de la teoría, observada aquí en este problema, no garantizada en general. Y con alpha=5 se ve el costo del exceso de prudencia: la política se parece más al histórico, usa más el 10 %, y rinde un poco menos.

Despliegue: ONNX y latencia real

Entrenar en PyTorch no obliga a servir en PyTorch. Exportar el modelo a ONNX desacopla el motor de entrenamiento del de ejecución: ONNX Runtime ejecuta el grafo en C++, sin el intérprete de Python en cada operación. El script importa el modelo del ejercicio 2, guardado como ej2_ft.py.

import time

import numpy as np
import onnxruntime as ort
import torch

from ej2_ft import FTTransformer, cargar_adult

# ——— 1. Exportar a ONNX ———
df, y, num_cols, cat_cols = cargar_adult()
ckpt = torch.load("ft_adult.pt")
modelo = FTTransformer(ckpt["n_num"], ckpt["card"])
modelo.load_state_dict(ckpt["estado"])
modelo.eval()

x_num = torch.randn(1, len(num_cols))
x_cat = torch.zeros(1, len(cat_cols), dtype=torch.long)
lote = torch.export.Dim("lote", min=1, max=4096)
torch.onnx.export(
    modelo, (x_num, x_cat), "ft_transformer.onnx",
    input_names=["x_num", "x_cat"], output_names=["logit"],
    dynamic_shapes={"x_num": {0: lote}, "x_cat": {0: lote}},
)

sesion = ort.InferenceSession("ft_transformer.onnx", providers=["CPUExecutionProvider"])
entrada = {"x_num": x_num.numpy(), "x_cat": x_cat.numpy()}
with torch.no_grad():
    ref = modelo(x_num, x_cat).numpy()
print(f"Diferencia máxima PyTorch vs ONNX: {np.abs(sesion.run(None, entrada)[0] - ref).max():.2e}")


def p50_p99(fn, n=2000):
    for _ in range(100):
        fn()
    t = []
    for _ in range(n):
        t0 = time.perf_counter()
        fn()
        t.append((time.perf_counter() - t0) * 1000)
    return np.percentile(t, 50), np.percentile(t, 99)


def torch_eager():
    with torch.no_grad():
        modelo(x_num, x_cat)


torch.set_num_threads(1)
so = ort.SessionOptions()
so.intra_op_num_threads = 1
sesion1 = ort.InferenceSession("ft_transformer.onnx", so, providers=["CPUExecutionProvider"])
for nombre, fn in [("PyTorch (eager)", torch_eager), ("ONNX Runtime", lambda: sesion1.run(None, entrada))]:
    p50, p99 = p50_p99(fn)
    print(f"Latencia de una predicción · {nombre:<16} p50 {p50:.3f} ms · p99 {p99:.3f} ms")

Diferencia máxima PyTorch vs ONNX: 2.38e-07
Latencia de una predicción · PyTorch (eager)  p50 1.044 ms · p99 1.973 ms
Latencia de una predicción · ONNX Runtime     p50 0.286 ms · p99 0.598 ms

La salida es la misma hasta la séptima cifra decimal, y la latencia de una predicción baja 3,6 veces, con un solo hilo en ambos casos. Mide siempre en tu hardware y con tu tamaño de lote: la ganancia depende del modelo, y en lotes grandes se reduce, porque ahí el costo ya no está en el intérprete sino en las multiplicaciones de matrices.

Dos detalles de versión. En PyTorch 2.14 torch.onnx.export usa por defecto el exportador basado en torch.export, y las dimensiones variables se declaran con dynamic_shapes y torch.export.Dim; el viejo dynamic_axes pertenece al exportador anterior. Y CatBoost también exporta a ONNX, pero no con variables categóricas: con el modelo de Adult, save_model(format="onnx") falla con ONNX-ML format export does yet not support categorical features. Si sirves CatBoost con categóricas, usa su propio runtime en C++ o codifica las categóricas antes de entrenar.

Monitoreo: drift con PSI y Kolmogorov-Smirnov

Un modelo en producción se degrada cuando cambian los datos. Hay dos tipos de cambio: el data drift, cuando cambia la distribución de las entradas, y el concept drift, cuando cambia la relación entre las entradas y lo que predices. Las métricas de abajo solo ven el primero. El segundo solo se ve midiendo el rendimiento con etiquetas nuevas.

El índice de estabilidad de la población (PSI) compara la distribución de referencia $P$ con la de producción $Q$, repartidas en los mismos $k$ intervalos:

Es una divergencia de Kullback-Leibler simetrizada, y tiene una trampa práctica: si un intervalo queda vacío, el logaritmo explota. Por eso el código suma un $\varepsilon$ y corta los intervalos por cuantiles de la referencia.

El test de Kolmogorov-Smirnov mide la máxima distancia vertical $D$ entre las dos funciones de distribución acumulada, y da un valor $p$ para la hipótesis de que vienen de la misma población.

import numpy as np
from scipy.stats import ks_2samp

from ej2_ft import cargar_adult

df, y, num_cols, cat_cols = cargar_adult()


def psi(referencia, actual, bins=10, eps=1e-6):
    """Population Stability Index con cortes por cuantiles de la referencia."""
    cortes = np.unique(np.quantile(referencia, np.linspace(0, 1, bins + 1)))
    cortes[0], cortes[-1] = -np.inf, np.inf
    p = np.histogram(referencia, cortes)[0] / len(referencia) + eps
    q = np.histogram(actual, cortes)[0] / len(actual) + eps
    return float(np.sum((p - q) * np.log(p / q)))


rng = np.random.default_rng(0)
edad = df["age"].to_numpy(dtype=float)
referencia = rng.choice(edad, 20_000, replace=False)
escenarios = {
    "misma población": rng.choice(edad, 5_000),
    "clientes 3 años mayores": rng.choice(edad, 5_000) + 3,
    "clientes 8 años mayores": rng.choice(edad, 5_000) + 8,
}
for nombre, actual in escenarios.items():
    ks = ks_2samp(referencia, actual)
    print(f"{nombre:<24} PSI {psi(referencia, actual):.3f} · K-S D={ks.statistic:.3f} p={ks.pvalue:.2g}")

misma población          PSI 0.001 · K-S D=0.007 p=0.99
clientes 3 años mayores  PSI 0.103 · K-S D=0.084 p=2.9e-25
clientes 8 años mayores  PSI 1.296 · K-S D=0.217 p=2.2e-165

La fila del medio es la lección. Con 20.000 y 5.000 observaciones, K-S da un valor $p$ de $10^{-25}$ para un cambio de tres años en la edad: con muestras grandes, K-S declara significativo casi cualquier cambio, relevante o no. El PSI mide cuánto se movió la distribución, y eso es lo que sirve para decidir. Las convenciones de la industria dicen que por debajo de 0,1 no hay cambio relevante, entre 0,1 y 0,25 hay que vigilar, y por encima hay un cambio importante. Son convenciones, no leyes: un PSI alto dice que hay que investigar, no que haya que reentrenar. Lo que decide si reentrenar es si el rendimiento del modelo cayó.

Elegir la herramienta, no la complejidad

Ninguna de estas tres técnicas es una mejora automática. En los tres ejercicios, lo que decidió el resultado fue una elección que no tenía que ver con la técnica en sí: optimizar la métrica correcta, reconocer que un gradient boosting sin ajustar ya era mejor que el Transformer, y no creerle a un modelo su propia estimación. La complejidad de un modelo se justifica cuando resuelve un problema que el modelo simple no resuelve, y eso se mide, no se supone: en latencia, en costo de entrenamiento y en estabilidad, antes de ponerlo en producción.

Para llevarte

• TPE elige el siguiente punto maximizando $\ell(x)/g(x)$, con un costo lineal en el número de observaciones. La poda es un mecanismo aparte, y solo funciona si cada trial informa valores intermedios comparables.
• Con clases muy desbalanceadas, optimiza PR-AUC: la ROC-AUC casi no distingue entre modelos, y la línea base de la PR-AUC es la prevalencia, no 0,5.
• El FT-Transformer convierte cada columna en un token y no necesita posiciones. En datos tabulares puros, un gradient boosting suele ganarle; úsalo cuando necesites lo que solo una red da.
• En refuerzo offline, el máximo de Bellman premia las acciones que nunca viste. CQL las penaliza y deja una estimación por debajo de la realidad, que es el error que conviene cometer.
• Exporta a ONNX y mide la latencia en tu hardware. Vigila el drift con PSI por magnitud, no con el valor $p$ de K-S, y decide si reentrenar por el rendimiento, no por el drift.