"""Original finite-state control lab: classical expert, BC and REINFORCE. Python + NumPy; no simulator assets, network downloads or GPU required. State: one-dimensional position in [-5,5]. Action: -1,0,+1. Goal: 0. This deliberately omits inertia, perception and contact physics. """ import argparse import json import time import numpy as np ACTIONS = np.array([-1, 0, 1]) TRAIN_STARTS = np.array([-5, -3, -1, 1, 3, 5]) TEST_STARTS = np.array([-4, -2, 2, 4]) HORIZON, GAMMA = 20, 1.0 def probabilities(logits, state): z = logits[state + 5] - logits[state + 5].max() e = np.exp(z) return e / e.sum() def gradient_checks(): logits, state, target = np.zeros((11, 3)), 2, 0 logits[state+5] = [0.2, -0.1, 0.4] probability = probabilities(logits, state) analytic = probability.copy() analytic[target] -= 1.0 numerical = np.zeros(3) epsilon = 1e-6 for col in range(3): plus, minus = logits.copy(), logits.copy() plus[state+5, col] += epsilon minus[state+5, col] -= epsilon numerical[col] = (-np.log(probabilities(plus, state)[target]) + np.log(probabilities(minus, state)[target])) / (2*epsilon) error = float(np.max(np.abs(analytic-numerical))) # The expectation of the log-policy score is zero for a fixed state; # this is why an action-independent baseline preserves its expectation. score_expectation = np.zeros(3) for action in range(3): score = -probability.copy() score[action] += 1 score_expectation += probability[action]*score assert error < 1e-8 assert np.max(np.abs(score_expectation)) < 1e-12 return {"softmax_cross_entropy_max_gradient_error": error, "expected_policy_score_max_abs": float(np.max(np.abs(score_expectation)))} def transition(state, action, rng, dropout=0.0): applied = 0 if rng.random() < dropout else action nxt = int(np.clip(state + applied, -5, 5)) reward = 1.0 if nxt == 0 else -0.04 return nxt, reward, nxt == 0 def expert_index(state): # Positive position -> action -1 -> index 0, and vice versa. return 2 if state < 0 else 0 if state > 0 else 1 def demonstrations(seed): rng = np.random.default_rng(seed) samples = [] for episode in range(120): state = int(rng.choice(TRAIN_STARTS)) for step in range(HORIZON): action_index = expert_index(state) samples.append((episode, state, action_index)) state, _, done = transition(state, int(ACTIONS[action_index]), rng) if done: break return samples def train_bc(seed): logits = np.zeros((11, 3)) samples = demonstrations(seed) train = [(s, a) for episode, s, a in samples if episode < 96] validation = [(s, a) for episode, s, a in samples if episode >= 96] for _ in range(300): gradient = np.zeros_like(logits) for state, action_index in train: gradient[state + 5] += probabilities(logits, state) gradient[state + 5, action_index] -= 1 logits -= 0.8 * gradient / len(train) loss = -np.mean([np.log(probabilities(logits, s)[a]) for s, a in validation]) return logits, {"synthetic_episodes": 120, "training_samples": len(train), "validation_samples": len(validation), "validation_cross_entropy": float(loss)} def train_reinforce(seed, batches=300): rng = np.random.default_rng(seed) logits, baseline = np.zeros((11, 3)), np.zeros(11) history, steps = [], 0 for batch in range(batches): gradient = np.zeros_like(logits) batch_returns = [] for _ in range(32): state = int(rng.choice(TRAIN_STARTS)) trajectory = [] for _ in range(HORIZON): probs = probabilities(logits, state) action_index = int(rng.choice(3, p=probs)) nxt, reward, done = transition(state, int(ACTIONS[action_index]), rng) trajectory.append((state, action_index, reward, probs)) state = nxt steps += 1 if done: break returns, total = [], 0.0 for _, _, reward, _ in reversed(trajectory): total = reward + GAMMA * total returns.append(total) returns.reverse() batch_returns.append(sum(row[2] for row in trajectory)) # Freeze the baseline over this episode; it must not depend on # the particular action whose log probability is differentiated. old_baseline = baseline.copy() for (state, action_index, _, probs), total in zip(trajectory, returns): score = -probs.copy() score[action_index] += 1 gradient[state + 5] += score * (total - old_baseline[state + 5]) baseline[state + 5] += 0.05 * (total - baseline[state + 5]) logits += 0.3 * gradient / 32 if batch % 50 == 0 or batch == batches - 1: history.append({"batch": batch, "mean_training_return": float(np.mean(batch_returns))}) return logits, {"episodes": batches * 32, "environment_steps": steps, "training_curve": history} def evaluate(logits, seed, expert_policy=False, random_policy=False, episodes=200): rng = np.random.default_rng(seed) successes, lengths, total_returns = 0, [], [] for _ in range(episodes): state, total = int(rng.choice(TEST_STARTS)), 0.0 for step in range(HORIZON): if expert_policy: action_index = expert_index(state) elif random_policy: action_index = int(rng.integers(3)) else: action_index = int(np.argmax(logits[state + 5])) state, reward, done = transition(state, int(ACTIONS[action_index]), rng, dropout=0.15) total += reward if done: successes += 1 break lengths.append(step + 1) total_returns.append(total) return {"episodes": episodes, "successes": successes, "success_rate": successes / episodes, "mean_steps": float(np.mean(lengths)), "mean_return": float(np.mean(total_returns)), "actuation_dropout": 0.15} def main(): parser = argparse.ArgumentParser() parser.add_argument("--seed", type=int, default=7) parser.add_argument("--batches", type=int, default=300) args = parser.parse_args() started = time.perf_counter() bc, bc_training = train_bc(args.seed) rl, rl_training = train_reinforce(args.seed, args.batches) eval_seed = args.seed + 10000 results = {"seed": args.seed, "numpy": np.__version__, "simulation": "finite-1D-integrator", "gradient_checks": gradient_checks(), "bc_training": bc_training, "rl_training": rl_training, "evaluation": {"classical": evaluate(None, eval_seed, expert_policy=True), "random": evaluate(None, eval_seed, random_policy=True), "bc": evaluate(bc, eval_seed), "reinforce": evaluate(rl, eval_seed)}, "elapsed_seconds": time.perf_counter() - started} assert all(np.isfinite(v["mean_return"]) for v in results["evaluation"].values()) print(json.dumps(results, indent=2)) if __name__ == "__main__": main()