tinybaselines 0.1.0a1__tar.gz

tinybaselines-0.1.0a1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Utkarsh Gill
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

tinybaselines-0.1.0a1/MANIFEST.in

@@ -0,0 +1 @@
+recursive-include examples *.py *.txt *.json *.jsonl

tinybaselines-0.1.0a1/PKG-INFO

@@ -0,0 +1,94 @@
+Metadata-Version: 2.4
+Name: tinybaselines
+Version: 0.1.0a1
+Summary: PPO in tinygrad.
+License-Expression: MIT
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: tinygrad==0.14.0
+Requires-Dist: numpy<3,>=1.26
+Dynamic: license-file
+
+# tinybaselines
+
+PPO in tinygrad. Continuous actions clipped to `[-1, 1]`.
+
+Python 3.11 and clang.
+
+```sh
+python3.11 -m venv .venv
+source .venv/bin/activate
+pip install -e . -r examples/lunar_lander/requirements.txt
+```
+
+Train, save, and load.
+
+```python
+from statistics import mean
+from tinybaselines import GaussianActorCritic, PPO, evaluate, learn
+from examples.lunar_lander.config import OBS_SCALE
+from examples.lunar_lander.train import make_env
+
+env = make_env(4, asynchronous=False)
+policy = GaussianActorCritic(state_dim=8, action_dim=2, hidden_dim=128,
+                            actor_layers=4, critic_layers=4, observation_scale=OBS_SCALE)
+ppo = PPO(policy, pi_lr=1e-3, vf_lr=1e-3, gamma=0.99, lamda=0.95,
+          K_epochs=20, eps_clip=0.2, batch_size=10_000, vf_coef=1.0, entropy_coef=0.001)
+
+for step in learn(env, ppo, total_timesteps=2_000_000, rollout_steps=20_000):
+    print(step.timesteps, step.rollout_seconds, step.update_seconds)
+
+ppo.save('ppo.safetensors')
+ppo = PPO.load('ppo.safetensors')
+policy = ppo.actor_critic
+observations, _ = env.reset()
+actions, raw_actions = policy.act(observations, deterministic=True)
+env.close()
+
+eval_env = make_env(4, asynchronous=False)
+returns = evaluate(eval_env, policy, num_episodes=100)
+eval_env.close()
+print(mean(returns))
+```
+
+`total_timesteps` and `rollout_steps` count transitions across all environments.
+Use a positive multiple of `env.num_envs` for `rollout_steps`. Training stops
+after the first full rollout that reaches the budget. Each step contains completed
+`episode_returns` and collection/update times. The caller owns the environment.
+Use a separate environment for evaluation during training.
+`evaluate` resets its environment and returns complete episode returns. It samples
+actions by default; `deterministic=True` uses clipped policy means.
+`actions` go to the environment; `raw_actions` go into PPO.
+The collector uses Gymnasium vector environments with same-step autoreset.
+For time-limit bootstrapping, the collector reads pre-reset observations from `infos['final_obs']`.
+
+PPO checkpoints contain the policy, observation scale, Adam state, and PPO settings.
+Continue training by passing the loaded PPO to `learn` with a new transition budget.
+Each `learn` call starts fresh episodes and a new transition counter.
+`policy.save(path)` exports a policy for inference.
+
+The full LunarLander trainer, evaluation, and PyTorch reference live in
+`examples/lunar_lander`.
+
+```sh
+python -m examples.lunar_lander.train
+```
+
+Pendulum uses three observations and one action. Its example rescales policy actions
+from `[-1, 1]` to torque in `[-2, 2]`.
+
+```sh
+pip install -e . -r examples/pendulum/requirements.txt
+DEV=CPU BEAM=0 python -m examples.pendulum.train --output runs/pendulum
+```
+
+The output directory contains the configuration, episode returns, update times,
+and a PPO checkpoint. Initial and final evaluations each use 100 fresh stochastic episodes.
+
+Three independent four-environment CPU runs reached mean returns of -141.4, -162.2,
+and -150.8 after 200,704 transitions each. Each final evaluation used the reloaded checkpoint.
+Our acceptance cutoff is -200 over 100 fresh stochastic episodes per run; all three passed.
+Gymnasium defines no solved threshold for Pendulum. We chose this cutoff after the first run.
+Run records: [1](examples/pendulum/results/cpu.jsonl),
+[2](examples/pendulum/results/cpu-2.jsonl), [3](examples/pendulum/results/cpu-3.jsonl).

tinybaselines-0.1.0a1/README.md

@@ -0,0 +1,82 @@
+# tinybaselines
+
+PPO in tinygrad. Continuous actions clipped to `[-1, 1]`.
+
+Python 3.11 and clang.
+
+```sh
+python3.11 -m venv .venv
+source .venv/bin/activate
+pip install -e . -r examples/lunar_lander/requirements.txt
+```
+
+Train, save, and load.
+
+```python
+from statistics import mean
+from tinybaselines import GaussianActorCritic, PPO, evaluate, learn
+from examples.lunar_lander.config import OBS_SCALE
+from examples.lunar_lander.train import make_env
+
+env = make_env(4, asynchronous=False)
+policy = GaussianActorCritic(state_dim=8, action_dim=2, hidden_dim=128,
+                            actor_layers=4, critic_layers=4, observation_scale=OBS_SCALE)
+ppo = PPO(policy, pi_lr=1e-3, vf_lr=1e-3, gamma=0.99, lamda=0.95,
+          K_epochs=20, eps_clip=0.2, batch_size=10_000, vf_coef=1.0, entropy_coef=0.001)
+
+for step in learn(env, ppo, total_timesteps=2_000_000, rollout_steps=20_000):
+    print(step.timesteps, step.rollout_seconds, step.update_seconds)
+
+ppo.save('ppo.safetensors')
+ppo = PPO.load('ppo.safetensors')
+policy = ppo.actor_critic
+observations, _ = env.reset()
+actions, raw_actions = policy.act(observations, deterministic=True)
+env.close()
+
+eval_env = make_env(4, asynchronous=False)
+returns = evaluate(eval_env, policy, num_episodes=100)
+eval_env.close()
+print(mean(returns))
+```
+
+`total_timesteps` and `rollout_steps` count transitions across all environments.
+Use a positive multiple of `env.num_envs` for `rollout_steps`. Training stops
+after the first full rollout that reaches the budget. Each step contains completed
+`episode_returns` and collection/update times. The caller owns the environment.
+Use a separate environment for evaluation during training.
+`evaluate` resets its environment and returns complete episode returns. It samples
+actions by default; `deterministic=True` uses clipped policy means.
+`actions` go to the environment; `raw_actions` go into PPO.
+The collector uses Gymnasium vector environments with same-step autoreset.
+For time-limit bootstrapping, the collector reads pre-reset observations from `infos['final_obs']`.
+
+PPO checkpoints contain the policy, observation scale, Adam state, and PPO settings.
+Continue training by passing the loaded PPO to `learn` with a new transition budget.
+Each `learn` call starts fresh episodes and a new transition counter.
+`policy.save(path)` exports a policy for inference.
+
+The full LunarLander trainer, evaluation, and PyTorch reference live in
+`examples/lunar_lander`.
+
+```sh
+python -m examples.lunar_lander.train
+```
+
+Pendulum uses three observations and one action. Its example rescales policy actions
+from `[-1, 1]` to torque in `[-2, 2]`.
+
+```sh
+pip install -e . -r examples/pendulum/requirements.txt
+DEV=CPU BEAM=0 python -m examples.pendulum.train --output runs/pendulum
+```
+
+The output directory contains the configuration, episode returns, update times,
+and a PPO checkpoint. Initial and final evaluations each use 100 fresh stochastic episodes.
+
+Three independent four-environment CPU runs reached mean returns of -141.4, -162.2,
+and -150.8 after 200,704 transitions each. Each final evaluation used the reloaded checkpoint.
+Our acceptance cutoff is -200 over 100 fresh stochastic episodes per run; all three passed.
+Gymnasium defines no solved threshold for Pendulum. We chose this cutoff after the first run.
+Run records: [1](examples/pendulum/results/cpu.jsonl),
+[2](examples/pendulum/results/cpu-2.jsonl), [3](examples/pendulum/results/cpu-3.jsonl).

tinybaselines-0.1.0a1/examples/lunar_lander/config.py

@@ -0,0 +1,5 @@
+import numpy as np
+
+ENV_NAME = 'LunarLanderContinuous-v3'
+# https://gymnasium.farama.org/environments/box2d/lunar_lander/
+OBS_SCALE = np.array([10, 6.666, 5, 7.5, 1, 2.5, 1, 1], dtype=np.float32)

tinybaselines-0.1.0a1/examples/lunar_lander/reference.py

@@ -0,0 +1,335 @@
+import gymnasium as gym
+import numpy as np
+import time
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import os
+import warnings
+from tqdm import trange
+from .config import ENV_NAME, OBS_SCALE
+from tinybaselines.rollout import RolloutBatch, RolloutState
+
+warnings.filterwarnings('ignore', message='pkg_resources is deprecated')
+
+state_dim, action_dim = 8, 2
+num_envs = int(os.getenv('NUM_ENVS', 4))
+
+max_epochs, steps_per_epoch = 100, 20_000
+log_interval, eval_interval = 5, 5
+eval_episodes = 100
+batch_size, K_epochs = 10_000, 20
+hidden_dim = 128
+actor_layers, critic_layers = 4, 4
+pi_lr, vf_lr = 1e-3, 1e-3
+gamma, gae_lambda, eps_clip = 0.99, 0.95, 0.2
+vf_coef, entropy_coef = 1.0, 0.001
+solved_threshold = 200
+
+PLOT = bool(int(os.getenv('PLOT', '0')))
+RENDER = bool(int(os.getenv('RENDER', '0')))
+RENDER_EPISODES = int(os.getenv('RENDER_EPISODES', '3'))
+METAL = bool(int(os.getenv('METAL', '0')))
+
+device = torch.device('mps' if METAL and torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu')
+
+if PLOT:
+    import matplotlib.pyplot as plt
+
+def make_env(n, render=False, asynchronous=True):
+    render_mode = 'human' if render else None
+    env_class = gym.vector.AsyncVectorEnv if asynchronous else gym.vector.SyncVectorEnv
+    return env_class(
+        [lambda: gym.make(ENV_NAME, render_mode=render_mode) for _ in range(n)],
+        autoreset_mode=gym.vector.AutoresetMode.SAME_STEP,
+    )
+
+def update_plot(ax, returns, threshold):
+    ax.clear()
+    ax.plot(returns, alpha=0.3, label='Episode Returns')
+    if len(returns) >= 100:
+        ma = np.convolve(returns, np.ones(100)/100, mode='valid')
+        ax.plot(range(99, len(returns)), ma, label='100-ep MA', linewidth=2)
+    ax.axhline(threshold, color='red', linestyle='--', alpha=0.5, label=f'Solved ({threshold})')
+    ax.legend()
+    ax.set_xlabel('Episode')
+    ax.set_ylabel('Return')
+    plt.pause(0.01)
+
+
+def track_episode_returns(done_mask, ep_returns, ep_rets, remaining=None):
+    for idx in np.where(done_mask)[0]:
+        if remaining is None or remaining[idx] > 0:
+            ep_returns.append(ep_rets[idx])
+            if remaining is not None:
+                remaining[idx] -= 1
+        ep_rets[idx] = 0.0
+
+class ActorCritic(nn.Module):
+    def __init__(self, state_dim, action_dim, hidden_dim, actor_layers, critic_layers):
+        super(ActorCritic, self).__init__()
+        self.config = dict(state_dim=state_dim, action_dim=action_dim, hidden_dim=hidden_dim,
+                           actor_layers=actor_layers, critic_layers=critic_layers)
+        self.register_buffer('observation_scale', torch.from_numpy(OBS_SCALE.copy()))
+        
+        actor = [nn.Linear(state_dim, hidden_dim), nn.ReLU()]
+        for _ in range(actor_layers - 1):
+            actor.extend([nn.Linear(hidden_dim, hidden_dim), nn.ReLU()])
+        actor.append(nn.Linear(hidden_dim, action_dim))
+        self.actor = nn.Sequential(*actor)
+        # small mu init → near-zero thrust initially (slow descent > random crashes)
+        self.actor[-1].weight.data.mul_(0.01)
+        self.actor[-1].bias.data.zero_()
+        self.log_std = nn.Parameter(torch.full((action_dim,), -0.7))
+        
+        critic = [nn.Linear(state_dim, hidden_dim), nn.ReLU()]
+        for _ in range(critic_layers - 1):
+            critic.extend([nn.Linear(hidden_dim, hidden_dim), nn.ReLU()])
+        critic.append(nn.Linear(hidden_dim, 1))
+        self.critic = nn.Sequential(*critic)
+
+    def forward(self, state):
+        action_mean = self.actor(state)
+        action_std = self.log_std.clamp(-5, 2).exp().expand_as(action_mean)
+        value = self.critic(state)
+        return action_mean, action_std, value
+    
+    def preprocess(self, observations):
+        return torch.as_tensor(observations, device=self.observation_scale.device, dtype=torch.float32) * self.observation_scale
+
+    @torch.inference_mode()
+    def act(self, state, deterministic=False):
+        state_tensor = self.preprocess(state)
+        action_mean = self.actor(state_tensor)
+        action_std = self.log_std.clamp(-5, 2).exp().expand_as(action_mean)
+        if deterministic:
+            action = action_mean.clamp(-1, 1)
+            return action.cpu().numpy(), action_mean.cpu().numpy()
+        sample = torch.distributions.Normal(action_mean, action_std).sample()
+        action = sample.clamp(-1, 1)
+        return action.cpu().numpy(), sample.cpu().numpy()
+
+    def save(self, path):
+        torch.save({'config': self.config, 'state_dict': self.state_dict()}, path)
+
+    @classmethod
+    def load(cls, path):
+        checkpoint = torch.load(path, map_location=device, weights_only=True)
+        model = cls(**checkpoint['config']).to(device)
+        model.load_state_dict(checkpoint['state_dict'])
+        return model
+
+
+class PPO:
+    def __init__(self, actor_critic, pi_lr, vf_lr, gamma, lamda, K_epochs, eps_clip, batch_size, vf_coef, entropy_coef):
+        self.actor_critic = actor_critic
+        self.pi_optimizer = optim.Adam(list(actor_critic.actor.parameters()) + [actor_critic.log_std], lr=pi_lr)
+        self.vf_optimizer = optim.Adam(actor_critic.critic.parameters(), lr=vf_lr)
+        self.gamma, self.lamda, self.K_epochs = gamma, lamda, K_epochs
+        self.eps_clip, self.batch_size, self.vf_coef, self.entropy_coef = eps_clip, batch_size, vf_coef, entropy_coef
+
+    def compute_advantages(self, rewards, state_values, terminated, truncated, next_values):
+        T, N = rewards.shape
+        advantages, gae = torch.zeros_like(rewards), torch.zeros(N, device=rewards.device)
+        episode_ends = terminated | truncated
+        for t in reversed(range(T)):
+            delta = rewards[t] + self.gamma * next_values[t] * ~terminated[t] - state_values[t]
+            gae = delta + self.gamma * self.lamda * ~episode_ends[t] * gae
+            advantages[t] = gae
+        returns = advantages + state_values
+        advantages = (advantages - advantages.mean()) / (advantages.std(correction=0) + 1e-8)
+        return advantages.reshape(-1), returns.reshape(-1)
+    
+    def compute_loss(self, batch_states, batch_actions, batch_logprobs, batch_advantages, batch_returns):
+        action_means, action_stds, state_values = self.actor_critic(batch_states)
+        dist = torch.distributions.Normal(action_means, action_stds)
+        action_logprobs = dist.log_prob(batch_actions).sum(-1)
+        ratios = torch.exp(action_logprobs - batch_logprobs)
+        actor_loss = -torch.min(ratios * batch_advantages, torch.clamp(ratios, 1-self.eps_clip, 1+self.eps_clip) * batch_advantages).mean()
+        critic_loss = F.mse_loss(state_values.squeeze(-1), batch_returns)
+        entropy = dist.entropy().sum(-1).mean()
+        
+        return actor_loss + self.vf_coef * critic_loss - self.entropy_coef * entropy
+    
+    def train_step(self, batch_states, batch_actions, batch_logprobs, batch_advantages, batch_returns):
+        self.pi_optimizer.zero_grad()
+        self.vf_optimizer.zero_grad()
+        loss = self.compute_loss(batch_states, batch_actions, batch_logprobs, batch_advantages, batch_returns)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(list(self.actor_critic.actor.parameters()) + [self.actor_critic.log_std], max_norm=0.5)
+        torch.nn.utils.clip_grad_norm_(self.actor_critic.critic.parameters(), max_norm=0.5)
+        self.pi_optimizer.step()
+        self.vf_optimizer.step()
+        return loss.detach()
+
+    def prepare(self, batch):
+        T, N = batch.rewards.shape
+        obs_flat = self.actor_critic.preprocess(batch.observations.reshape(T * N, -1))
+        raw_act_flat = torch.from_numpy(batch.raw_actions.reshape(T * N, -1)).to(device=device, dtype=torch.float32)
+        rew_t = torch.from_numpy(batch.rewards).to(device=device, dtype=torch.float32)
+        terminated_t = torch.from_numpy(batch.terminated).to(device=device)
+        truncated_t = torch.from_numpy(batch.truncated).to(device=device)
+        
+        with torch.no_grad():
+            mean, std, val = self.actor_critic(obs_flat)
+            dist = torch.distributions.Normal(mean, std)
+            old_logprobs = dist.log_prob(raw_act_flat).sum(-1)
+            old_values = val.squeeze(-1).view(T, N)
+            next_obs_t = self.actor_critic.preprocess(batch.next_observations.reshape(T * N, -1))
+            next_values = self.actor_critic.critic(next_obs_t).squeeze(-1).view(T, N)
+            advantages, returns = self.compute_advantages(rew_t, old_values, terminated_t, truncated_t, next_values)
+        
+        return obs_flat, raw_act_flat, old_logprobs, advantages, returns
+
+    def update(self, batch):
+        obs_flat, raw_act_flat, old_logprobs, advantages, returns = self.prepare(batch)
+        num_samples = obs_flat.size(0)
+        for _ in range(self.K_epochs):
+            perm = torch.randperm(num_samples, device=device)
+            for start in range(0, num_samples, self.batch_size):
+                idx = perm[start:start + self.batch_size]
+                self.train_step(obs_flat[idx], raw_act_flat[idx], old_logprobs[idx], advantages[idx], returns[idx])
+
+def rollout(env, actor_critic, num_steps=None, num_episodes=None, deterministic=False, state=None):
+    assert (num_steps is None) != (num_episodes is None), "Specify exactly one: num_steps or num_episodes"
+    if num_episodes is not None and num_episodes <= 0:
+        raise ValueError('num_episodes must be positive')
+    
+    N = env.num_envs
+    if state is None:
+        state = RolloutState(env.reset()[0], np.zeros(N))
+    states = state.observations
+    ep_returns, ep_rets = [], state.episode_returns
+    
+    collect = num_steps is not None
+    remaining = None if collect else (num_episodes + np.arange(N)) // N
+    if collect:
+        T = num_steps // N
+        obs = np.empty((T, N, state_dim), dtype=np.float32)
+        next_obs = np.empty_like(obs)
+        raw_act = np.empty((T, N, action_dim), dtype=np.float32)
+        rew = np.empty((T, N), dtype=np.float32)
+        terminations = np.empty((T, N), dtype=bool)
+        truncations = np.empty((T, N), dtype=bool)
+    
+    t = 0
+    while True:
+        actions, raw_actions = actor_critic.act(states, deterministic=deterministic)
+        if collect:
+            obs[t], raw_act[t] = states, raw_actions
+        
+        states, rewards, terminated, truncated, infos = env.step(actions)
+        d = np.logical_or(terminated, truncated)
+        if collect:
+            rew[t], terminations[t], truncations[t] = rewards, terminated, truncated
+            next_obs[t] = states
+            for idx in np.flatnonzero(d):
+                next_obs[t, idx] = infos['final_obs'][idx]
+        
+        ep_rets += rewards
+        track_episode_returns(d, ep_returns, ep_rets, remaining)
+        t += 1
+        
+        if (collect and t >= T) or (not collect and not remaining.any()):
+            break
+    
+    state.observations = states
+    return (RolloutBatch(obs, raw_act, rew, terminations, truncations, next_obs), ep_returns) if collect else ep_returns
+
+class TrainingContext:
+    def __init__(self):
+        self.ac = ActorCritic(state_dim, action_dim, hidden_dim, actor_layers, critic_layers).to(device)
+        self.ppo = PPO(self.ac, pi_lr, vf_lr, gamma, gae_lambda, K_epochs, eps_clip, batch_size, vf_coef, entropy_coef)
+        
+        self.env = make_env(num_envs)
+        self.rollout_state = RolloutState(self.env.reset()[0], np.zeros(num_envs))
+        self.eval_env = make_env(min(num_envs, eval_episodes), asynchronous=False)
+        self.all_episode_returns = []
+        self.last_eval_stochastic = float('-inf')
+        self.pbar = trange(max_epochs, desc="Training", unit='epoch')
+        self.rollout_times = []
+        self.update_times = []
+        
+        if PLOT:
+            plt.ion()
+            _, self.ax = plt.subplots()
+        else:
+            self.ax = None
+    
+    def cleanup(self):
+        self.env.close()
+        self.eval_env.close()
+        self.pbar.close()
+        if PLOT:
+            plt.ioff()
+            plt.show()
+
+def train_one_epoch(epoch, ctx):
+    t0 = time.perf_counter()
+    batch, ep_rets = rollout(
+        ctx.env, ctx.ac, num_steps=steps_per_epoch, state=ctx.rollout_state
+    )
+    t1 = time.perf_counter()
+    ctx.rollout_times.append(t1 - t0)
+    
+    t0 = time.perf_counter()
+    ctx.ppo.update(batch)
+    t1 = time.perf_counter()
+    ctx.update_times.append(t1 - t0)
+    
+    ctx.all_episode_returns.extend(ep_rets)
+    ctx.pbar.update(1)
+    
+    if epoch % eval_interval == 0:
+        ctx.last_eval_stochastic = evaluate_policy(
+            ctx.ac, env=ctx.eval_env, deterministic=False
+        )
+        if RENDER:
+            evaluate_policy(ctx.ac, render=True, num_episodes=RENDER_EPISODES, deterministic=False)
+    
+    if epoch % log_interval == 0:
+        s = ctx.ac.log_std.clamp(-5, 2).exp().detach().cpu().numpy()
+        rollout_ms = np.mean(ctx.rollout_times[-log_interval:]) * 1000
+        update_ms = np.mean(ctx.update_times[-log_interval:]) * 1000
+        total_ms = rollout_ms + update_ms
+        ctx.pbar.write(f"Epoch {epoch:3d}  n_ep={len(ep_rets):3d}  ret={np.mean(ep_rets):7.1f}±{np.std(ep_rets):5.1f}  eval_stoch={ctx.last_eval_stochastic:6.1f}  σ=[{s[0]:.2f} {s[1]:.2f}]  ⏱ {total_ms:.0f}ms (rollout:{rollout_ms:.0f}ms update:{update_ms:.0f}ms)")
+    
+    if PLOT and epoch % (log_interval * 2) == 0:
+        update_plot(ctx.ax, ctx.all_episode_returns, solved_threshold)
+    
+    if ctx.last_eval_stochastic >= solved_threshold:
+        ctx.pbar.write(f"\n{'='*60}\nSOLVED at epoch {epoch}! eval_stoch={ctx.last_eval_stochastic:.1f} ≥ {solved_threshold}\n{'='*60}")
+        if RENDER:
+            evaluate_policy(ctx.ac, render=True, num_episodes=RENDER_EPISODES, deterministic=False)
+        return True
+    
+    return False
+
+def train():
+    ctx = TrainingContext()
+    for epoch in range(max_epochs):
+        if train_one_epoch(epoch, ctx):
+            break
+    ctx.cleanup()
+
+def evaluate_policy(actor_critic, num_episodes=eval_episodes, render=False, env=None, deterministic=True):
+    close_env = env is None
+    if env is None:
+        env = make_env(1 if render else num_episodes, render)
+    ep_rets = rollout(env, actor_critic, num_episodes=num_episodes, deterministic=deterministic)
+    if close_env:
+        env.close()
+    return float(np.mean(ep_rets)) if ep_rets else 0.0
+
+if __name__ == '__main__':
+    print(f"Using {device} device")
+    if METAL:
+        print(f"MPS available: {torch.backends.mps.is_available()}")
+        print(f"MPS built: {torch.backends.mps.is_built()}")
+        if device.type == 'mps':
+            print("✅ MPS device active")
+        else:
+            print("⚠️  METAL=1 but device is not MPS")
+    print(f"PyTorch threads: {torch.get_num_threads()}")
+    train()

tinybaselines-0.1.0a1/examples/lunar_lander/requirements-reference.txt

@@ -0,0 +1,2 @@
+-r requirements.txt
+torch==2.8.0

tinybaselines-0.1.0a1/examples/lunar_lander/requirements.txt

@@ -0,0 +1,4 @@
+-r ../../requirements.txt
+gymnasium[box2d]==1.1.1
+tqdm>=4.65
+matplotlib>=3.5

tinybaselines-0.1.0a1/examples/lunar_lander/results/2026-10-26/phase-initial-tinygrad-thread4.jsonl

@@ -0,0 +1 @@
+{"seconds": 0.614734499999031, "event": "config", "framework": "tinygrad", "environment": "LunarLanderContinuous-v3", "source_commit": "8f0b31a96b03e23f43b35dcf12e36302e0036e40", "device": "CPU", "dtype": "dtypes.float", "platform": "macOS-27.0.1-arm64-arm-64bit", "python": "3.11.14", "versions": {"gymnasium": "1.1.1", "numpy": "2.4.6", "tinygrad": "0.13.0"}, "backend_source": {"url": "https://github.com/tinygrad/tinygrad.git", "vcs_info": {"vcs": "git", "commit_id": "5b60a09ab0a7f20c3426e505251c48cac020d1fa", "requested_revision": "5b60a09ab0a7f20c3426e505251c48cac020d1fa"}}, "settings": {"num_envs": 4, "steps_per_epoch": 20000, "max_epochs": 3, "batch_size": 10000, "K_epochs": 20, "pi_lr": 0.001, "vf_lr": 0.001, "gamma": 0.99, "gae_lambda": 0.95, "eps_clip": 0.2, "vf_coef": 1.0, "entropy_coef": 0.001, "eval_interval": 5, "eval_episodes": 100, "solved_threshold": 200}, "architecture": {"state_dim": 8, "action_dim": 2, "hidden_dim": 128, "actor_layers": 4, "critic_layers": 4}, "beam": "1", "torch_threads": null}

tinybaselines-0.1.0a1/examples/lunar_lander/results/2026-10-26/phase-initial-torch.jsonl

@@ -0,0 +1,6 @@
+{"seconds": 0.48239595799532253, "event": "config", "framework": "torch", "environment": "LunarLanderContinuous-v3", "source_commit": "8f0b31a96b03e23f43b35dcf12e36302e0036e40", "device": "cpu", "dtype": "torch.float32", "platform": "macOS-27.0.1-arm64-arm-64bit", "python": "3.11.14", "versions": {"gymnasium": "1.1.1", "numpy": "2.4.6", "torch": "2.8.0"}, "backend_source": {}, "settings": {"num_envs": 4, "steps_per_epoch": 20000, "max_epochs": 3, "batch_size": 10000, "K_epochs": 20, "pi_lr": 0.001, "vf_lr": 0.001, "gamma": 0.99, "gae_lambda": 0.95, "eps_clip": 0.2, "vf_coef": 1.0, "entropy_coef": 0.001, "eval_interval": 5, "eval_episodes": 100, "solved_threshold": 200}, "architecture": {"state_dim": 8, "action_dim": 2, "hidden_dim": 128, "actor_layers": 4, "critic_layers": 4}, "beam": "1", "torch_threads": 4}
+{"seconds": 2.1622258749994216, "event": "epoch", "steps": 20000, "seconds_total": 1.679733083001338, "rollout_seconds": 0.8259355000045616, "update_seconds": 0.43762720800441457, "other_seconds": 0.41617037499236176}
+{"seconds": 2.1622854999950505, "event": "evaluation", "steps": 20000, "score": -144.3682780688725}
+{"seconds": 3.3910218749952037, "event": "epoch", "steps": 40000, "seconds_total": 1.2287172080032178, "rollout_seconds": 0.8063115420009126, "update_seconds": 0.4221108749989071, "other_seconds": 0.0002947910033981316}
+{"seconds": 4.608377207994636, "event": "epoch", "steps": 60000, "seconds_total": 1.2172969999955967, "rollout_seconds": 0.7961072079997393, "update_seconds": 0.4209024999945541, "other_seconds": 0.00028729200130328536}
+{"seconds": 5.081694374996005, "event": "final", "steps": 60000, "score": -104.08706254110945, "seconds_to_stop": 4.608431249995192, "stopped_on_threshold": false, "policy": "policy.pt"}

tinybaselines-0.1.0a1/examples/lunar_lander/results/2026-10-26/phase-protocol.json

@@ -0,0 +1,53 @@
+{
+  "date": "2026-10-26",
+  "hardware": "Apple M4",
+  "purpose": "Locate wall time in the real training path. These are independent stochastic runs, not identical-input backend comparisons.",
+  "source_commit": "483d1da6b213611ff1dccd9520149edd4777671d",
+  "selected_runs": {
+    "torch": "phase-torch-thread4.jsonl",
+    "tinygrad": "phase-tinygrad-thread10.jsonl"
+  },
+  "run_order": [
+    "phase-initial-torch.jsonl",
+    "phase-initial-tinygrad-thread4.jsonl",
+    "phase-tinygrad-thread4-error.jsonl",
+    "phase-torch-thread4.jsonl",
+    "phase-tinygrad-thread10.jsonl"
+  ],
+  "environment": {
+    "METAL": "0",
+    "DEV": "CPU",
+    "NUM_ENVS": "4",
+    "NUM_CPU_THREADS": "10",
+    "BEAM": "1",
+    "JIT": "1",
+    "XDG_CACHE_HOME": "$PWD/.venv/cache"
+  },
+  "requested_threads": {
+    "torch": 4,
+    "tinygrad": 10
+  },
+  "compiler_cache_directory_reused": true,
+  "epochs_per_completed_run": 3,
+  "warmup": "No epochs excluded from the records. Article table shows zero-based epoch 2; first epoch includes initial JIT calls.",
+  "timings": {
+    "seconds_total": "Wall time around train_one_epoch. Excludes backend import, environment construction, save/load and the final evaluation.",
+    "rollout_seconds": "Collection of 20000 transitions, including policy.act, environment step/IPC, and rollout buffer work.",
+    "update_seconds": "PPO.update, including preparation/GAE, shuffling/indexing, and 40 optimizer steps.",
+    "other_seconds": "seconds_total minus rollout_seconds and update_seconds. Includes scheduled evaluation, progress display and logging."
+  },
+  "evaluation": "100 fresh stochastic episodes at epoch 0 and after save/load. Three-epoch probes are not convergence runs.",
+  "initial_runs": "First Torch run and the first four-thread tinygrad failure used 8f0b31a. The initial failure has only a config record because error logging was not yet implemented.",
+  "failed_run": {
+    "file": "phase-tinygrad-thread4-error.jsonl",
+    "NUM_CPU_THREADS": "4",
+    "error_type": "KernelOptError",
+    "message": "too many threads",
+    "location": "Replaying cached beam-search options during PPO.prepare, before any epoch completed.",
+    "cache_manually_cleared": false
+  },
+  "reproduce": [
+    "METAL=0 DEV=CPU NUM_ENVS=4 NUM_CPU_THREADS=10 BEAM=1 JIT=1 XDG_CACHE_HOME=\"$PWD/.venv/cache\" .venv/bin/python -c 'from pathlib import Path; from examples.lunar_lander import reference; from examples.lunar_lander.run_learning import run; reference.max_epochs=3; run(\"torch\", Path(\"runs/torch-phases\"))'",
+    "METAL=0 DEV=CPU NUM_ENVS=4 NUM_CPU_THREADS=10 BEAM=1 JIT=1 XDG_CACHE_HOME=\"$PWD/.venv/cache\" .venv/bin/python -c 'from pathlib import Path; from examples.lunar_lander import train; from examples.lunar_lander.run_learning import run; train.max_epochs=3; run(\"tinygrad\", Path(\"runs/tinygrad-phases\"))'"
+  ]
+}

tinybaselines-0.1.0a1/examples/lunar_lander/results/2026-10-26/phase-tinygrad-thread10.jsonl

@@ -0,0 +1,6 @@
+{"seconds": 0.707170124995173, "event": "config", "framework": "tinygrad", "environment": "LunarLanderContinuous-v3", "source_commit": "483d1da6b213611ff1dccd9520149edd4777671d", "device": "CPU", "dtype": "dtypes.float", "platform": "macOS-27.0.1-arm64-arm-64bit", "python": "3.11.14", "versions": {"gymnasium": "1.1.1", "numpy": "2.4.6", "tinygrad": "0.13.0"}, "backend_source": {"url": "https://github.com/tinygrad/tinygrad.git", "vcs_info": {"vcs": "git", "commit_id": "5b60a09ab0a7f20c3426e505251c48cac020d1fa", "requested_revision": "5b60a09ab0a7f20c3426e505251c48cac020d1fa"}}, "settings": {"num_envs": 4, "steps_per_epoch": 20000, "max_epochs": 3, "batch_size": 10000, "K_epochs": 20, "pi_lr": 0.001, "vf_lr": 0.001, "gamma": 0.99, "gae_lambda": 0.95, "eps_clip": 0.2, "vf_coef": 1.0, "entropy_coef": 0.001, "eval_interval": 5, "eval_episodes": 100, "solved_threshold": 200}, "architecture": {"state_dim": 8, "action_dim": 2, "hidden_dim": 128, "actor_layers": 4, "critic_layers": 4}, "beam": "1", "torch_threads": null, "tinygrad_cpu_threads": 10}
+{"seconds": 27.337733790998755, "event": "epoch", "steps": 20000, "seconds_total": 26.63045024999883, "rollout_seconds": 10.67004020800232, "update_seconds": 10.632532167001045, "other_seconds": 5.327877874995465}
+{"seconds": 27.33781787499902, "event": "evaluation", "steps": 20000, "score": -148.26613410664652}
+{"seconds": 41.51010841599782, "event": "epoch", "steps": 40000, "seconds_total": 14.172251832998882, "rollout_seconds": 9.591514917003224, "update_seconds": 4.580453624999791, "other_seconds": 0.0002832909958669916}
+{"seconds": 55.548277791000146, "event": "epoch", "steps": 60000, "seconds_total": 14.038110916000733, "rollout_seconds": 9.761195333005162, "update_seconds": 4.276613417001499, "other_seconds": 0.00030216599407140166}
+{"seconds": 62.28786079099518, "event": "final", "steps": 60000, "score": -90.02494137380951, "seconds_to_stop": 55.548339540997404, "stopped_on_threshold": false, "policy": "policy.safetensors"}