upload self_train.py

220c07d verified about 1 month ago

16.3 kB

"""主动自我训练（用户数据 → 拆 train/test → 精英保留 + 模型回滚）
策略：
  1. 每轮同时采样模型生成 + 随机生成候选
  2. 在 train 集上评估 R²，跨轮保留精英候选（elite pool）
  3. 只对 R² > 0.1 的正向候选做伪标签 CE 微调（避免强化错误模式）
  4. REINFORCE 策略梯度用于模型采样的候选
  5. 保存每轮开始前的模型 state，如果本轮最佳 R² 下降则回滚
  6. 终止条件：R² > 0.995 或 连续 4 轮 test R² 无改善
"""

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

import config
from vocab import Vocab
from model import SymbolicRegressionModel
from data_generator import RandomExpressionGenerator

# 关键：stdout 行缓冲，保证 print() 立即刷新。
# 否则在后台/管道运行时（python 默认块缓冲），进度会攒在一旁，
# 直到进程结束才一次性输出 —— 这正是"训练过程不输出"的根因。
if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(line_buffering=True)

# 全局静音候选表达式在数值评估时的 numpy 警告（exp/sqrt/log/power 溢出、
# 定义域外、除零等均属常规现象，结果已用 isfinite/complex 过滤）。
# 全局设置可覆盖 lambdify 生成的已编译模块，避免每一轮都刷屏。
np.seterr(over="ignore", invalid="ignore", divide="ignore")
warnings.filterwarnings("ignore", module="lambdifygenerated")

COPIED_MODEL_STATE_KEY = "_saved_lstm_state"


def train_test_split_arrays(X, y, test_size=0.3, seed=None):
    """手写随机 train/test 拆分。"""
    rng = np.random.RandomState(seed)
    n = X.shape[0]
    idx = rng.permutation(n)
    n_te = max(1, int(test_size * n))
    te_idx, tr_idx = idx[:n_te], idx[n_te:]
    return X[tr_idx], X[te_idx], y[tr_idx], y[te_idx]


def eval_expr_on_data(expr_str: str, X: np.ndarray, y_true: np.ndarray) -> tuple:
    """计算给定表达式在数据上的预测值和 MSE。无效返回 NaN。"""
    import sympy as sp
    import warnings
    D = X.shape[1]
    try:
        expr = sp.sympify(expr_str)
        symbols = [sp.Symbol(f"x{i}") for i in range(D)]
        f = sp.lambdify(symbols, expr, modules="numpy")
        # 候选表达式常落在定义域外或数值溢出（exp/sqrt/log/power），
        # 静音这些 numpy RuntimeWarning，避免刷屏。结果 NaN/inf 会被下方过滤。
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            with np.errstate(over="ignore", invalid="ignore", divide="ignore"):
                y_pred = f(*[X[:, i] for i in range(D)])
        if np.iscomplexobj(y_pred) or not np.isfinite(y_pred).all():
            return np.full_like(y_true, np.nan), np.nan
        mse = float(np.nanmean((y_pred - y_true) ** 2))
        return y_pred, mse
    except Exception:
        return np.full_like(y_true, np.nan), np.nan


def _generate_random_candidates(num_candidates: int, D: int, seed: int = None) -> list[str]:
    """生成随机候选，混合深度。"""
    gen = RandomExpressionGenerator(seed=seed)
    rng = np.random.RandomState(seed)
    candidates = []
    depths = [2, 3, 4]
    while len(candidates) < num_candidates * 3:
        try:
            depth = rng.choice(depths)
            data = gen.generate(var_count=D, max_depth=int(depth))
            candidates.append(data["expr_str"])
        except Exception:
            continue
    return candidates[:num_candidates]


def _evaluate_candidates(expr_strings: list[str], X: np.ndarray, y: np.ndarray,
                         y_var: float) -> tuple:
    """批量评估候选，返回 rewards (R²) 和 mses。无效 reward=-1。"""
    n = len(expr_strings)
    rewards = np.full(n, -1.0)
    mses = np.full(n, np.nan)
    for i, expr_str in enumerate(expr_strings):
        _, mse = eval_expr_on_data(expr_str, X, y)
        mses[i] = mse
        if np.isfinite(mse):
            rewards[i] = 1.0 - mse / max(y_var, 1e-9)
    return rewards, mses


@torch.no_grad()
def _encode_input(X: np.ndarray, y: np.ndarray, device: str):
    """编码用户数据为模型输入 [1, N, D_MAX], [1, N]."""
    D = X.shape[1]
    N = X.shape[0]
    X_t = torch.zeros(1, N, config.D_MAX, dtype=torch.float32, device=device)
    X_t[0, :, :D] = torch.from_numpy(X).to(device)
    y_t = torch.from_numpy(y).float().unsqueeze(0).to(device)
    n_samples = torch.tensor([N], device=device)
    var_counts = torch.tensor([D], device=device)
    return X_t, y_t, n_samples, var_counts


@torch.enable_grad()
def _sample_model(model, X_t, y_t, n_samples, var_counts, device,
                  max_len=None, num_candidates=20, greedy=False, temperature=1.5):
    """从模型采样候选。greedy → argmax；否则多项式采样 + 噪声扰动。"""
    if max_len is None:
        max_len = config.T_MAX
    sample_size = num_candidates
    vocab = model.vocab

    z = model.encoder(X_t, y_t, n_samples, var_counts)

    if not greedy:
        noise_scale = 0.15
        z = z + torch.randn_like(z) * noise_scale
        temp = temperature
    else:
        temp = 0.1

    hidden = z.unsqueeze(0).expand(-1, sample_size, -1).contiguous()
    cell = z.unsqueeze(0).expand(-1, sample_size, -1).contiguous()
    if not greedy:
        hidden = hidden + torch.randn_like(hidden) * noise_scale
        cell = cell + torch.randn_like(cell) * noise_scale

    token = torch.full((sample_size, 1), fill_value=vocab.sos_id, device=device)
    sequences = [token]
    log_probs = []
    done = torch.zeros(sample_size, dtype=torch.bool, device=device)

    for _ in range(max_len):
        embedded = model.lstm.embedding(token)
        output, (hidden, cell) = model.lstm.lstm(embedded, (hidden, cell))
        logits = model.lstm.proj(output).squeeze(1)
        logits = logits / temp

        if greedy:
            token = logits.argmax(dim=-1, keepdim=True)
        else:
            log_probs_t = F.log_softmax(logits, dim=-1)
            probs_t = F.softmax(logits, dim=-1)
            dist = torch.distributions.Categorical(probs=probs_t)
            token = dist.sample().unsqueeze(-1)
            lp = torch.gather(log_probs_t, 1, token)
            lp = lp.masked_fill(done.unsqueeze(-1), 0.0)
            log_probs.append(lp)

        sequences.append(token)
        done = done | (token.squeeze(-1) == vocab.eos_id)
        if done.all():
            break

    tokens = torch.cat(sequences, dim=1)
    if not greedy:
        log_probs = torch.cat(log_probs, dim=1)
    tokens = tokens[:, 1:]

    expr_strings = []
    for i in range(sample_size):
        toks = tokens[i]
        if (toks == vocab.eos_id).any():
            eos = (toks == vocab.eos_id).nonzero(as_tuple=True)[0][0].item()
            toks = toks[:eos]
        expr_strings.append(vocab.decode(toks.tolist()))

    return tokens, log_probs if not greedy else None, expr_strings


def _save_model_snapshot(model):
    """保存 LSTM 参数快照到 model 的临时属性。"""
    sd = {k: v.data.clone() for k, v in model.lstm.state_dict().items()}
    setattr(model, COPIED_MODEL_STATE_KEY, sd)


def _restore_model_snapshot(model):
    """从快照恢复 LSTM 参数。"""
    sd = getattr(model, COPIED_MODEL_STATE_KEY, None)
    if sd is not None:
        for k, v in sd.items():
            model.lstm.state_dict()[k].data.copy_(v)


def _encode_expr(expr_str: str, vocab: Vocab, device: str):
    """把表达式字符串编码为 token 序列。"""
    token_ids = vocab.encode(expr_str)
    token_ids = [vocab.sos_id] + token_ids + [vocab.eos_id]
    token_ids = token_ids[:config.T_MAX]
    return torch.tensor(token_ids, dtype=torch.long, device=device).unsqueeze(0)


def self_train_on_user_data(model, X, y, device, split=0.7,
                            max_rounds=15, candidates=96, seed=None,
                            verbose=True):
    """针对用户数据 (X, y) 执行主动自我训练。
    返回 dict {best_expr, best_r2, converged, rounds, results, round_results}
    """
    vocab = model.vocab

    if X.ndim == 1:
        X = X.reshape(-1, 1)
    D = X.shape[1]

    Xtr, Xte, ytr, yte = train_test_split_arrays(
        X, y, test_size=1 - split, seed=seed
    )
    ytr_var = float(np.var(ytr))
    yte_var = float(np.var(yte))

    if verbose:
        print(f"  [数据] 训练集: {Xtr.shape[0]} 样本, 测试集: {Xte.shape[0]} 样本, 维度 D={D}")
        print(f"  [数据] y 方差: train={ytr_var:.4f}, test={yte_var:.4f}")

    Xtr_t, ytr_t, ntr, vctr = _encode_input(Xtr, ytr, device)

    optimizer = torch.optim.Adam(model.lstm.parameters(), lr=3e-4)

    # ---- 精英池：跨轮保留历史最佳候选 ----
    elite_pool = []
    PATIENCE = 4

    results = []
    best_r2 = -float("inf")
    best_expr = None
    no_improve = 0
    converged = False

    for rnd in range(1, max_rounds + 1):
        model.train()

        # --- 保存本轮开始前的模型 ---
        _save_model_snapshot(model)

        # --- 温度退火 ---
        temp = max(1.2, 3.0 - 0.18 * (rnd - 1))

        # --- 生成候选 ---
        n_model = int(candidates * 0.45)
        n_rand = candidates - n_model

        # 模型采样
        tokens, log_probs, model_exprs = _sample_model(
            model, Xtr_t, ytr_t, ntr, vctr, device,
            num_candidates=n_model, greedy=False, temperature=temp
        )

        # 随机生成（混合深度，提供多样性）
        rand_candidates = _generate_random_candidates(
            n_rand, D, seed=(seed or 0) + rnd * 131 + 7
        )

        # 加入精英池中的候选（去重后按 R² 排序取前10）
        elite_exprs = [e for e, _ in elite_pool[:10]]
        expr_strings = list(model_exprs) + list(rand_candidates) + elite_exprs

        # --- 评估候选 ---
        rewards, mses = _evaluate_candidates(expr_strings, Xtr, ytr, ytr_var)
        valid_mask = np.isfinite(mses)
        valid_idxs = np.where(valid_mask)[0]

        if valid_idxs.size == 0:
            if verbose:
                print(f"  [自我训练] round {rnd}/{max_rounds} | T={temp:.2f} | 所有候选无效，跳过此轮")
            results.append({"round": rnd, "train_r2": -1.0, "test_r2": -1.0,
                           "test_expr": "", "n_valid": 0})
            continue

        best_local_idx = valid_idxs[int(np.argmax(rewards[valid_idxs]))]
        best_local_expr = expr_strings[best_local_idx]
        best_local_r2 = float(rewards[best_local_idx])

        # --- 更新精英池 ---
        for i in valid_idxs:
            r2 = float(rewards[i])
            if r2 > 0.05:
                elite_pool.append((expr_strings[i], r2))
        elite_pool = sorted(set(elite_pool), key=lambda x: -x[1])[:20]

        # --- 训练步：伪标签 + REINFORCE ---
        # 只对 R² > 0.1 的正向候选做 CE 伪标签
        good_mask = rewards > 0.1
        good_idxs = np.where(good_mask)[0]

        ce_losses = []
        for idx in good_idxs:
            expr_str = expr_strings[idx]
            if expr_str == "":
                continue
            try:
                t_tensor = _encode_expr(expr_str, vocab, device)
                z = model.encoder(Xtr_t, ytr_t, ntr, vctr)
                hidden = z.unsqueeze(0)
                cell = z.unsqueeze(0)
                logits, _, _ = model.lstm(t_tensor[:, :-1], hidden, cell)
                targets = t_tensor[:, 1:]
                ce = F.cross_entropy(logits.reshape(-1, vocab.vocab_size), targets.reshape(-1))
                ce_losses.append(ce * 0.3)
            except Exception:
                continue

        if ce_losses and verbose:
            pos_count = len(ce_losses)
            best_r2_among_pos = max(rewards[good_idxs])
            print(f"    └─ 伪标签: {pos_count} 个正向候选 (R²>{0.1:.1f}, 最高={best_r2_among_pos:.3f})")

        # REINFORCE：仅对模型采样候选（前 n_model 个）中有效且导通的
        reinforce_loss = None
        model_valid_idx = valid_idxs[valid_idxs < n_model]
        non_collapsed = []
        for idx in model_valid_idx:
            if idx < n_model and log_probs is not None:
                lp_row = log_probs[idx]
                if lp_row.numel() > 0 and lp_row.sum() < -0.01:
                    non_collapsed.append(idx)

        if len(non_collapsed) >= 3:
            m_rewards = rewards[non_collapsed]
            m_rewards = np.where(np.isfinite(m_rewards), m_rewards, -1.0)
            m_advantage_t = torch.tensor(m_rewards - m_rewards.mean(),
                                         dtype=torch.float32, device=device)
            m_logp = log_probs[non_collapsed]
            m_logp_sum = m_logp.sum(dim=-1)
            reinforce_loss = -(m_advantage_t * m_logp_sum).mean()

        if ce_losses or reinforce_loss is not None:
            all_losses = ce_losses[:]
            if reinforce_loss is not None:
                all_losses.append(reinforce_loss * 0.5)
            loss = torch.stack(all_losses).mean()
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.lstm.parameters(), 1.0)
            optimizer.step()

        # --- test 评估 ---
        te_expr = best_local_expr
        _, te_mse = eval_expr_on_data(te_expr, Xte, yte)
        if np.isfinite(te_mse):
            te_r2 = 1.0 - te_mse / max(yte_var, 1e-9)
        else:
            te_r2 = -float("inf")

        # --- 模型回滚：如果本轮最佳 R² 比上一轮差，恢复模型 ---
        prev_r2 = results[-1]["best_local_r2"] if results else -float("inf")
        if best_local_r2 < prev_r2 - 0.1 and len(results) >= 1:
            _restore_model_snapshot(model)
            if verbose:
                print(f"    └─ 模型回滚（best R² {best_local_r2:.3f} < 上一轮 {prev_r2:.3f}）")
            # 重新用恢复后的模型采样一次作为补充
            tokens2, lp2, model_exprs2 = _sample_model(
                model, Xtr_t, ytr_t, ntr, vctr, device,
                num_candidates=n_model, greedy=False, temperature=temp
            )
            expr_strings2 = list(model_exprs2) + elite_exprs[:5]
            rewards2, mses2 = _evaluate_candidates(expr_strings2, Xtr, ytr, ytr_var)
            vm2 = np.isfinite(mses2)
            vi2 = np.where(vm2)[0]
            if vi2.size > 0:
                best_i2 = vi2[int(np.argmax(rewards2[vi2]))]
                best_local_expr = expr_strings2[best_i2]
                best_local_r2 = float(rewards2[best_i2])

                _, te_mse2 = eval_expr_on_data(best_local_expr, Xte, yte)
                if np.isfinite(te_mse2):
                    te_r2 = 1.0 - te_mse2 / max(yte_var, 1e-9)
                else:
                    te_r2 = -float("inf")

        if te_r2 > best_r2:
            best_r2 = te_r2
            best_expr = te_expr
            no_improve = 0
        else:
            no_improve += 1

        nv = int(valid_mask.sum())
        results.append({
            "round": rnd, "train_r2": float(np.nanmean(rewards)),
            "test_r2": float(te_r2), "test_expr": te_expr,
            "best_local_expr": best_local_expr,
            "best_local_r2": best_local_r2, "n_valid": nv,
            "n_good_ce": len(ce_losses),
        })

        if verbose:
            avg_r2 = float(np.nanmean(rewards))
            print(f"  [自我训练] round {rnd}/{max_rounds} | T={temp:.2f} | "
                  f"valid={nv}/{candidates} | "
                  f"avg R²={avg_r2:.4f} | "
                  f"best R²={best_local_r2:.4f} | "
                  f"test R²={te_r2:.4f} | {te_expr}")

        # --- 终止判定 ---
        if te_r2 > 0.995:
            if verbose:
                print(f"    → R² 达标 ({te_r2:.4f})，停止自我训练 ✓")
            converged = True
            break
        if no_improve >= PATIENCE:
            if verbose:
                print(f"    → test R² 连续 {PATIENCE} 轮无改善，停止自我训练（最佳: {best_r2:.4f}）")
            converged = True
            break

    return {
        "best_expr": best_expr,
        "best_r2": best_r2,
        "converged": converged,
        "rounds": len(results),
        "round_results": results,
    }

	"""主动自我训练（用户数据 → 拆 train/test → 精英保留 + 模型回滚）

	策略：
	1. 每轮同时采样模型生成 + 随机生成候选
	2. 在 train 集上评估 R²，跨轮保留精英候选（elite pool）
	3. 只对 R² > 0.1 的正向候选做伪标签 CE 微调（避免强化错误模式）
	4. REINFORCE 策略梯度用于模型采样的候选
	5. 保存每轮开始前的模型 state，如果本轮最佳 R² 下降则回滚
	6. 终止条件：R² > 0.995 或连续 4 轮 test R² 无改善
	"""

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

	import config
	from vocab import Vocab
	from model import SymbolicRegressionModel
	from data_generator import RandomExpressionGenerator

	# 关键：stdout 行缓冲，保证 print() 立即刷新。
	# 否则在后台/管道运行时（python 默认块缓冲），进度会攒在一旁，
	# 直到进程结束才一次性输出 —— 这正是"训练过程不输出"的根因。
	if hasattr(sys.stdout, "reconfigure"):
	sys.stdout.reconfigure(line_buffering=True)

	# 全局静音候选表达式在数值评估时的 numpy 警告（exp/sqrt/log/power 溢出、
	# 定义域外、除零等均属常规现象，结果已用 isfinite/complex 过滤）。
	# 全局设置可覆盖 lambdify 生成的已编译模块，避免每一轮都刷屏。
	np.seterr(over="ignore", invalid="ignore", divide="ignore")
	warnings.filterwarnings("ignore", module="lambdifygenerated")

	COPIED_MODEL_STATE_KEY = "_saved_lstm_state"


	def train_test_split_arrays(X, y, test_size=0.3, seed=None):
	"""手写随机 train/test 拆分。"""
	rng = np.random.RandomState(seed)
	n = X.shape[0]
	idx = rng.permutation(n)
	n_te = max(1, int(test_size * n))
	te_idx, tr_idx = idx[:n_te], idx[n_te:]
	return X[tr_idx], X[te_idx], y[tr_idx], y[te_idx]


	def eval_expr_on_data(expr_str: str, X: np.ndarray, y_true: np.ndarray) -> tuple:
	"""计算给定表达式在数据上的预测值和 MSE。无效返回 NaN。"""
	import sympy as sp
	import warnings
	D = X.shape[1]
	try:
	expr = sp.sympify(expr_str)
	symbols = [sp.Symbol(f"x{i}") for i in range(D)]
	f = sp.lambdify(symbols, expr, modules="numpy")
	# 候选表达式常落在定义域外或数值溢出（exp/sqrt/log/power），
	# 静音这些 numpy RuntimeWarning，避免刷屏。结果 NaN/inf 会被下方过滤。
	with warnings.catch_warnings():
	warnings.simplefilter("ignore")
	with np.errstate(over="ignore", invalid="ignore", divide="ignore"):
	y_pred = f(*[X[:, i] for i in range(D)])
	if np.iscomplexobj(y_pred) or not np.isfinite(y_pred).all():
	return np.full_like(y_true, np.nan), np.nan
	mse = float(np.nanmean((y_pred - y_true) ** 2))
	return y_pred, mse
	except Exception:
	return np.full_like(y_true, np.nan), np.nan


	def _generate_random_candidates(num_candidates: int, D: int, seed: int = None) -> list[str]:
	"""生成随机候选，混合深度。"""
	gen = RandomExpressionGenerator(seed=seed)
	rng = np.random.RandomState(seed)
	candidates = []
	depths = [2, 3, 4]
	while len(candidates) < num_candidates * 3:
	try:
	depth = rng.choice(depths)
	data = gen.generate(var_count=D, max_depth=int(depth))
	candidates.append(data["expr_str"])
	except Exception:
	continue
	return candidates[:num_candidates]


	def _evaluate_candidates(expr_strings: list[str], X: np.ndarray, y: np.ndarray,
	y_var: float) -> tuple:
	"""批量评估候选，返回 rewards (R²) 和 mses。无效 reward=-1。"""
	n = len(expr_strings)
	rewards = np.full(n, -1.0)
	mses = np.full(n, np.nan)
	for i, expr_str in enumerate(expr_strings):
	_, mse = eval_expr_on_data(expr_str, X, y)
	mses[i] = mse
	if np.isfinite(mse):
	rewards[i] = 1.0 - mse / max(y_var, 1e-9)
	return rewards, mses


	@torch.no_grad()
	def _encode_input(X: np.ndarray, y: np.ndarray, device: str):
	"""编码用户数据为模型输入 [1, N, D_MAX], [1, N]."""
	D = X.shape[1]
	N = X.shape[0]
	X_t = torch.zeros(1, N, config.D_MAX, dtype=torch.float32, device=device)
	X_t[0, :, :D] = torch.from_numpy(X).to(device)
	y_t = torch.from_numpy(y).float().unsqueeze(0).to(device)
	n_samples = torch.tensor([N], device=device)
	var_counts = torch.tensor([D], device=device)
	return X_t, y_t, n_samples, var_counts


	@torch.enable_grad()
	def _sample_model(model, X_t, y_t, n_samples, var_counts, device,
	max_len=None, num_candidates=20, greedy=False, temperature=1.5):
	"""从模型采样候选。greedy → argmax；否则多项式采样 + 噪声扰动。"""
	if max_len is None:
	max_len = config.T_MAX
	sample_size = num_candidates
	vocab = model.vocab

	z = model.encoder(X_t, y_t, n_samples, var_counts)

	if not greedy:
	noise_scale = 0.15
	z = z + torch.randn_like(z) * noise_scale
	temp = temperature
	else:
	temp = 0.1

	hidden = z.unsqueeze(0).expand(-1, sample_size, -1).contiguous()
	cell = z.unsqueeze(0).expand(-1, sample_size, -1).contiguous()
	if not greedy:
	hidden = hidden + torch.randn_like(hidden) * noise_scale
	cell = cell + torch.randn_like(cell) * noise_scale

	token = torch.full((sample_size, 1), fill_value=vocab.sos_id, device=device)
	sequences = [token]
	log_probs = []
	done = torch.zeros(sample_size, dtype=torch.bool, device=device)

	for _ in range(max_len):
	embedded = model.lstm.embedding(token)
	output, (hidden, cell) = model.lstm.lstm(embedded, (hidden, cell))
	logits = model.lstm.proj(output).squeeze(1)
	logits = logits / temp

	if greedy:
	token = logits.argmax(dim=-1, keepdim=True)
	else:
	log_probs_t = F.log_softmax(logits, dim=-1)
	probs_t = F.softmax(logits, dim=-1)
	dist = torch.distributions.Categorical(probs=probs_t)
	token = dist.sample().unsqueeze(-1)
	lp = torch.gather(log_probs_t, 1, token)
	lp = lp.masked_fill(done.unsqueeze(-1), 0.0)
	log_probs.append(lp)

	sequences.append(token)
	done = done \| (token.squeeze(-1) == vocab.eos_id)
	if done.all():
	break

	tokens = torch.cat(sequences, dim=1)
	if not greedy:
	log_probs = torch.cat(log_probs, dim=1)
	tokens = tokens[:, 1:]

	expr_strings = []
	for i in range(sample_size):
	toks = tokens[i]
	if (toks == vocab.eos_id).any():
	eos = (toks == vocab.eos_id).nonzero(as_tuple=True)[0][0].item()
	toks = toks[:eos]
	expr_strings.append(vocab.decode(toks.tolist()))

	return tokens, log_probs if not greedy else None, expr_strings


	def _save_model_snapshot(model):
	"""保存 LSTM 参数快照到 model 的临时属性。"""
	sd = {k: v.data.clone() for k, v in model.lstm.state_dict().items()}
	setattr(model, COPIED_MODEL_STATE_KEY, sd)


	def _restore_model_snapshot(model):
	"""从快照恢复 LSTM 参数。"""
	sd = getattr(model, COPIED_MODEL_STATE_KEY, None)
	if sd is not None:
	for k, v in sd.items():
	model.lstm.state_dict()[k].data.copy_(v)


	def _encode_expr(expr_str: str, vocab: Vocab, device: str):
	"""把表达式字符串编码为 token 序列。"""
	token_ids = vocab.encode(expr_str)
	token_ids = [vocab.sos_id] + token_ids + [vocab.eos_id]
	token_ids = token_ids[:config.T_MAX]
	return torch.tensor(token_ids, dtype=torch.long, device=device).unsqueeze(0)


	def self_train_on_user_data(model, X, y, device, split=0.7,
	max_rounds=15, candidates=96, seed=None,
	verbose=True):
	"""针对用户数据 (X, y) 执行主动自我训练。

	返回 dict {best_expr, best_r2, converged, rounds, results, round_results}
	"""
	vocab = model.vocab

	if X.ndim == 1:
	X = X.reshape(-1, 1)
	D = X.shape[1]

	Xtr, Xte, ytr, yte = train_test_split_arrays(
	X, y, test_size=1 - split, seed=seed
	)
	ytr_var = float(np.var(ytr))
	yte_var = float(np.var(yte))

	if verbose:
	print(f" [数据] 训练集: {Xtr.shape[0]} 样本, 测试集: {Xte.shape[0]} 样本, 维度 D={D}")
	print(f" [数据] y 方差: train={ytr_var:.4f}, test={yte_var:.4f}")

	Xtr_t, ytr_t, ntr, vctr = _encode_input(Xtr, ytr, device)

	optimizer = torch.optim.Adam(model.lstm.parameters(), lr=3e-4)

	# ---- 精英池：跨轮保留历史最佳候选 ----
	elite_pool = []
	PATIENCE = 4

	results = []
	best_r2 = -float("inf")
	best_expr = None
	no_improve = 0
	converged = False

	for rnd in range(1, max_rounds + 1):
	model.train()

	# --- 保存本轮开始前的模型 ---
	_save_model_snapshot(model)

	# --- 温度退火 ---
	temp = max(1.2, 3.0 - 0.18 * (rnd - 1))

	# --- 生成候选 ---
	n_model = int(candidates * 0.45)
	n_rand = candidates - n_model

	# 模型采样
	tokens, log_probs, model_exprs = _sample_model(
	model, Xtr_t, ytr_t, ntr, vctr, device,
	num_candidates=n_model, greedy=False, temperature=temp
	)

	# 随机生成（混合深度，提供多样性）
	rand_candidates = _generate_random_candidates(
	n_rand, D, seed=(seed or 0) + rnd * 131 + 7
	)

	# 加入精英池中的候选（去重后按 R² 排序取前10）
	elite_exprs = [e for e, _ in elite_pool[:10]]
	expr_strings = list(model_exprs) + list(rand_candidates) + elite_exprs

	# --- 评估候选 ---
	rewards, mses = _evaluate_candidates(expr_strings, Xtr, ytr, ytr_var)
	valid_mask = np.isfinite(mses)
	valid_idxs = np.where(valid_mask)[0]

	if valid_idxs.size == 0:
	if verbose:
	print(f" [自我训练] round {rnd}/{max_rounds} \| T={temp:.2f} \| 所有候选无效，跳过此轮")
	results.append({"round": rnd, "train_r2": -1.0, "test_r2": -1.0,
	"test_expr": "", "n_valid": 0})
	continue

	best_local_idx = valid_idxs[int(np.argmax(rewards[valid_idxs]))]
	best_local_expr = expr_strings[best_local_idx]
	best_local_r2 = float(rewards[best_local_idx])

	# --- 更新精英池 ---
	for i in valid_idxs:
	r2 = float(rewards[i])
	if r2 > 0.05:
	elite_pool.append((expr_strings[i], r2))
	elite_pool = sorted(set(elite_pool), key=lambda x: -x[1])[:20]

	# --- 训练步：伪标签 + REINFORCE ---
	# 只对 R² > 0.1 的正向候选做 CE 伪标签
	good_mask = rewards > 0.1
	good_idxs = np.where(good_mask)[0]

	ce_losses = []
	for idx in good_idxs:
	expr_str = expr_strings[idx]
	if expr_str == "":
	continue
	try:
	t_tensor = _encode_expr(expr_str, vocab, device)
	z = model.encoder(Xtr_t, ytr_t, ntr, vctr)
	hidden = z.unsqueeze(0)
	cell = z.unsqueeze(0)
	logits, _, _ = model.lstm(t_tensor[:, :-1], hidden, cell)
	targets = t_tensor[:, 1:]
	ce = F.cross_entropy(logits.reshape(-1, vocab.vocab_size), targets.reshape(-1))
	ce_losses.append(ce * 0.3)
	except Exception:
	continue

	if ce_losses and verbose:
	pos_count = len(ce_losses)
	best_r2_among_pos = max(rewards[good_idxs])
	print(f" └─ 伪标签: {pos_count} 个正向候选 (R²>{0.1:.1f}, 最高={best_r2_among_pos:.3f})")

	# REINFORCE：仅对模型采样候选（前 n_model 个）中有效且导通的
	reinforce_loss = None
	model_valid_idx = valid_idxs[valid_idxs < n_model]
	non_collapsed = []
	for idx in model_valid_idx:
	if idx < n_model and log_probs is not None:
	lp_row = log_probs[idx]
	if lp_row.numel() > 0 and lp_row.sum() < -0.01:
	non_collapsed.append(idx)

	if len(non_collapsed) >= 3:
	m_rewards = rewards[non_collapsed]
	m_rewards = np.where(np.isfinite(m_rewards), m_rewards, -1.0)
	m_advantage_t = torch.tensor(m_rewards - m_rewards.mean(),
	dtype=torch.float32, device=device)
	m_logp = log_probs[non_collapsed]
	m_logp_sum = m_logp.sum(dim=-1)
	reinforce_loss = -(m_advantage_t * m_logp_sum).mean()

	if ce_losses or reinforce_loss is not None:
	all_losses = ce_losses[:]
	if reinforce_loss is not None:
	all_losses.append(reinforce_loss * 0.5)
	loss = torch.stack(all_losses).mean()
	optimizer.zero_grad()
	loss.backward()
	torch.nn.utils.clip_grad_norm_(model.lstm.parameters(), 1.0)
	optimizer.step()

	# --- test 评估 ---
	te_expr = best_local_expr
	_, te_mse = eval_expr_on_data(te_expr, Xte, yte)
	if np.isfinite(te_mse):
	te_r2 = 1.0 - te_mse / max(yte_var, 1e-9)
	else:
	te_r2 = -float("inf")

	# --- 模型回滚：如果本轮最佳 R² 比上一轮差，恢复模型 ---
	prev_r2 = results[-1]["best_local_r2"] if results else -float("inf")
	if best_local_r2 < prev_r2 - 0.1 and len(results) >= 1:
	_restore_model_snapshot(model)
	if verbose:
	print(f" └─ 模型回滚（best R² {best_local_r2:.3f} < 上一轮 {prev_r2:.3f}）")
	# 重新用恢复后的模型采样一次作为补充
	tokens2, lp2, model_exprs2 = _sample_model(
	model, Xtr_t, ytr_t, ntr, vctr, device,
	num_candidates=n_model, greedy=False, temperature=temp
	)
	expr_strings2 = list(model_exprs2) + elite_exprs[:5]
	rewards2, mses2 = _evaluate_candidates(expr_strings2, Xtr, ytr, ytr_var)
	vm2 = np.isfinite(mses2)
	vi2 = np.where(vm2)[0]
	if vi2.size > 0:
	best_i2 = vi2[int(np.argmax(rewards2[vi2]))]
	best_local_expr = expr_strings2[best_i2]
	best_local_r2 = float(rewards2[best_i2])

	_, te_mse2 = eval_expr_on_data(best_local_expr, Xte, yte)
	if np.isfinite(te_mse2):
	te_r2 = 1.0 - te_mse2 / max(yte_var, 1e-9)
	else:
	te_r2 = -float("inf")

	if te_r2 > best_r2:
	best_r2 = te_r2
	best_expr = te_expr
	no_improve = 0
	else:
	no_improve += 1

	nv = int(valid_mask.sum())
	results.append({
	"round": rnd, "train_r2": float(np.nanmean(rewards)),
	"test_r2": float(te_r2), "test_expr": te_expr,
	"best_local_expr": best_local_expr,
	"best_local_r2": best_local_r2, "n_valid": nv,
	"n_good_ce": len(ce_losses),
	})

	if verbose:
	avg_r2 = float(np.nanmean(rewards))
	print(f" [自我训练] round {rnd}/{max_rounds} \| T={temp:.2f} \| "
	f"valid={nv}/{candidates} \| "
	f"avg R²={avg_r2:.4f} \| "
	f"best R²={best_local_r2:.4f} \| "
	f"test R²={te_r2:.4f} \| {te_expr}")

	# --- 终止判定 ---
	if te_r2 > 0.995:
	if verbose:
	print(f" → R² 达标 ({te_r2:.4f})，停止自我训练 ✓")
	converged = True
	break
	if no_improve >= PATIENCE:
	if verbose:
	print(f" → test R² 连续 {PATIENCE} 轮无改善，停止自我训练（最佳: {best_r2:.4f}）")
	converged = True
	break

	return {
	"best_expr": best_expr,
	"best_r2": best_r2,
	"converged": converged,
	"rounds": len(results),
	"round_results": results,
	}