tinybaselines 0.1.0a1__py3-none-any.whl

tinybaselines/__init__.py

@@ -0,0 +1,5 @@
+from .ppo import PPO
+from .policies import GaussianActorCritic
+from .rollout import RolloutBatch, RolloutState, rollout
+from .training import TrainingStep, learn
+from .evaluation import evaluate

tinybaselines/evaluation.py

@@ -0,0 +1,6 @@
+from .rollout import rollout
+
+
+def evaluate(env, policy, num_episodes, deterministic=False):
+    """Return total rewards from num_episodes fresh episodes."""
+    return rollout(env, policy, num_episodes=num_episodes, deterministic=deterministic)

tinybaselines/policies.py

@@ -0,0 +1,75 @@
+import json
+
+import numpy as np
+from tinygrad import Tensor, TinyJit, nn
+
+
+class MLP:
+    def __init__(self, input_dim, output_dim, hidden_dim, layer_count):
+        self.layers = [nn.Linear(input_dim, hidden_dim)]
+        self.layers.extend(nn.Linear(hidden_dim, hidden_dim) for _ in range(layer_count - 1))
+        self.layers.append(nn.Linear(hidden_dim, output_dim))
+
+    def __call__(self, x):
+        for layer in self.layers[:-1]:
+            x = layer(x).relu()
+        return self.layers[-1](x)
+
+
+class GaussianActorCritic:
+    def __init__(self, state_dim, action_dim, hidden_dim, actor_layers, critic_layers, observation_scale=None):
+        self.config = dict(state_dim=state_dim, action_dim=action_dim, hidden_dim=hidden_dim,
+                           actor_layers=actor_layers, critic_layers=critic_layers)
+        self.actor = MLP(state_dim, action_dim, hidden_dim, actor_layers)
+        self.critic = MLP(state_dim, 1, hidden_dim, critic_layers)
+        self.actor.layers[-1].weight.assign(self.actor.layers[-1].weight * 0.01).realize()
+        self.actor.layers[-1].bias.assign(Tensor.zeros_like(self.actor.layers[-1].bias)).realize()
+        self.log_std = Tensor.full((action_dim,), -0.7).realize()
+        scale = np.ones(state_dim, np.float32) if observation_scale is None else np.asarray(observation_scale, np.float32)
+        self.observation_scale = Tensor(scale.copy()).is_param_(False).realize()
+        self.jit_shape = None
+        self.actor_jit = TinyJit(lambda state: self.actor(state).realize())
+        self.refresh_inference()
+
+    def __call__(self, state):
+        action_mean = self.actor(state)
+        log_std = self.log_std.clip(-5, 2)
+        value = self.critic(state)
+        return action_mean, log_std, value
+
+    def preprocess(self, observations):
+        return Tensor(observations.astype(np.float32), device=self.observation_scale.device) * self.observation_scale
+
+    def act(self, state, deterministic=False):
+        state_tensor = self.preprocess(state)
+        if self.jit_shape != state_tensor.shape:
+            self.actor_jit.reset()
+            self.jit_shape = state_tensor.shape
+        mean = self.actor_jit(state_tensor).numpy()
+        if deterministic:
+            action = np.clip(mean, -1, 1)
+            return action, mean
+        sample = np.random.normal(mean, self.action_std).astype(np.float32)
+        action = np.clip(sample, -1, 1)
+        return action, sample
+
+    def state_dict(self):
+        return nn.state.get_state_dict(dict(actor=self.actor, critic=self.critic, log_std=self.log_std,
+                                            observation_scale=self.observation_scale))
+
+    def save(self, path):
+        nn.state.safe_save(self.state_dict(), path, metadata={'config': json.dumps(self.config)})
+
+    @classmethod
+    def load(cls, path):
+        _, _, header = nn.state.safe_load_metadata(path)
+        model = cls(**json.loads(header['__metadata__']['config']))
+        nn.state.load_state_dict(model.state_dict(), nn.state.safe_load(path), verbose=False)
+        model.refresh_inference()
+        return model
+
+    def refresh_inference(self):
+        self.actor_jit.reset()
+        self.action_std = np.exp(self.log_std.clip(-5, 2).numpy())
+        if not np.isfinite(self.action_std).all():
+            raise FloatingPointError('non-finite policy standard deviation')

tinybaselines/ppo.py

@@ -0,0 +1,155 @@
+import json
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import numpy as np
+from tinygrad import Context, Tensor, TinyJit, nn
+
+from .policies import GaussianActorCritic
+
+LOG_2PI = float(np.log(2.0 * np.pi))
+NORMAL_ENTROPY_CONSTANT = float(0.5 * np.log(2.0 * np.pi * np.e))
+
+
+def normal_log_prob(action_mean, log_std, actions):
+    normalized = (actions - action_mean) / log_std.exp()
+    return (-0.5 * (normalized.square() + 2.0 * log_std + LOG_2PI)).sum(axis=-1)
+
+
+def clip_grad_norm(parameters, max_norm):
+    gradients = [parameter.grad for parameter in parameters if parameter.grad is not None]
+    if not gradients:
+        return
+    total_norm = sum((gradient.square().sum() for gradient in gradients)).sqrt()
+    scale = (max_norm / (total_norm + 1e-6)).clip(max_=1.0)
+    for parameter in parameters:
+        if parameter.grad is not None:
+            parameter.grad = parameter.grad * scale
+
+
+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.device = actor_critic.log_std.device
+        self.actor_parameters = nn.state.get_parameters(actor_critic.actor) + [actor_critic.log_std]
+        self.critic_parameters = nn.state.get_parameters(actor_critic.critic)
+        self.pi_optimizer = nn.optim.Adam(self.actor_parameters, lr=pi_lr)
+        self.vf_optimizer = nn.optim.Adam(self.critic_parameters, lr=vf_lr)
+        self.optimizer = nn.optim.OptimizerGroup(self.pi_optimizer, self.vf_optimizer)
+        self.gamma, self.lamda, self.K_epochs = gamma, lamda, K_epochs
+        self.eps_clip, self.batch_size = eps_clip, batch_size
+        self.vf_coef, self.entropy_coef = vf_coef, entropy_coef
+        self._train_steps = {}
+
+    @property
+    def config(self):
+        return dict(pi_lr=self.pi_optimizer.lr.item(), vf_lr=self.vf_optimizer.lr.item(),
+                    gamma=self.gamma, lamda=self.lamda, K_epochs=self.K_epochs, eps_clip=self.eps_clip,
+                    batch_size=self.batch_size, vf_coef=self.vf_coef, entropy_coef=self.entropy_coef)
+
+    def state_dict(self):
+        state = dict(policy=self.actor_critic.state_dict())
+        for name in ('pi_optimizer', 'vf_optimizer'):
+            optimizer = getattr(self, name)
+            state[name] = {key: getattr(optimizer, key) for key in ('lr', 'b1_t', 'b2_t', 'm', 'v')}
+        return nn.state.get_state_dict(state)
+
+    def save(self, path):
+        path = Path(path)
+        metadata = dict(format='tinybaselines.ppo.v1', policy_config=json.dumps(self.actor_critic.config),
+                        ppo_config=json.dumps(self.config))
+        with TemporaryDirectory(dir=path.parent) as directory:
+            temporary = Path(directory) / path.name
+            nn.state.safe_save(self.state_dict(), temporary, metadata=metadata)
+            temporary.replace(path)
+
+    @classmethod
+    def load(cls, path):
+        _, _, header = nn.state.safe_load_metadata(path)
+        metadata = header.get('__metadata__', {})
+        if metadata.get('format') != 'tinybaselines.ppo.v1':
+            raise ValueError('unsupported PPO checkpoint format')
+        policy = GaussianActorCritic(**json.loads(metadata['policy_config']))
+        ppo = cls(policy, **json.loads(metadata['ppo_config']))
+        nn.state.load_state_dict(ppo.state_dict(), nn.state.safe_load(path), verbose=False)
+        policy.refresh_inference()
+        return ppo
+
+    def compute_advantages(self, rewards, state_values, terminated, truncated, next_values):
+        T, N = rewards.shape
+        advantages = np.zeros_like(rewards, dtype=np.float32)
+        gae = np.zeros(N, dtype=np.float32)
+        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() + 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, log_std, state_values = self.actor_critic(batch_states)
+        action_logprobs = normal_log_prob(action_means, log_std, batch_actions)
+        ratios = (action_logprobs - batch_logprobs).exp()
+        unclipped = ratios * batch_advantages
+        clipped = ratios.clip(1 - self.eps_clip, 1 + self.eps_clip) * batch_advantages
+        actor_loss = -unclipped.minimum(clipped).mean()
+        critic_loss = (state_values.squeeze(-1) - batch_returns).square().mean()
+        entropy = (log_std + NORMAL_ENTROPY_CONSTANT).sum()
+        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):
+        size = batch_states.shape[0]
+        if size not in self._train_steps:
+            self._train_steps[size] = TinyJit(self._train_step)
+        loss = self._train_steps[size](batch_states, batch_actions, batch_logprobs, batch_advantages, batch_returns)
+        self.actor_critic.refresh_inference()
+        return loss
+
+    @Context(TRAINING=1)
+    def _train_step(self, batch_states, batch_actions, batch_logprobs, batch_advantages, batch_returns):
+        self.optimizer.zero_grad()
+        loss = self.compute_loss(
+            batch_states, batch_actions, batch_logprobs, batch_advantages, batch_returns
+        )
+        loss.backward()
+        clip_grad_norm(self.actor_parameters, 0.5)
+        clip_grad_norm(self.critic_parameters, 0.5)
+        return loss.realize(*self.optimizer.schedule_step())
+
+    def prepare(self, batch):
+        T, N = batch.rewards.shape
+        obs_flat = self.actor_critic.preprocess(batch.observations.reshape(T * N, -1))
+        raw_act_np = batch.raw_actions.reshape(T * N, -1).astype(np.float32)
+        raw_act_flat = Tensor(raw_act_np, device=self.device)
+        old_mean, old_log_std, old_value = self.actor_critic(obs_flat)
+        old_logprobs_np = normal_log_prob(old_mean, old_log_std, raw_act_flat).numpy()
+        old_values_np = old_value.numpy().reshape(T, N)
+        next_obs_t = self.actor_critic.preprocess(batch.next_observations.reshape(T * N, -1))
+        next_values_np = self.actor_critic.critic(next_obs_t).numpy().reshape(T, N)
+        advantages_np, returns_np = self.compute_advantages(
+            batch.rewards, old_values_np, batch.terminated, batch.truncated, next_values_np
+        )
+
+        old_logprobs = Tensor(old_logprobs_np, device=self.device)
+        advantages = Tensor(advantages_np, device=self.device)
+        returns = Tensor(returns_np, device=self.device)
+        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.shape[0]
+
+        for _ in range(self.K_epochs):
+            permutation = np.random.permutation(num_samples).astype(np.int32)
+            for start in range(0, num_samples, self.batch_size):
+                idx = Tensor(permutation[start:start + self.batch_size], device=self.device)
+                batch_states = obs_flat[idx].contiguous().realize()
+                batch_actions = raw_act_flat[idx].contiguous().realize()
+                batch_logprobs = old_logprobs[idx].contiguous().realize()
+                batch_advantages = advantages[idx].contiguous().realize()
+                batch_returns = returns[idx].contiguous().realize()
+                self.train_step(
+                    batch_states, batch_actions, batch_logprobs, batch_advantages, batch_returns
+                )

tinybaselines/rollout.py

@@ -0,0 +1,80 @@
+from dataclasses import dataclass
+from typing import NamedTuple
+
+import numpy as np
+
+
+@dataclass
+class RolloutState:
+    observations: np.ndarray
+    episode_returns: np.ndarray
+
+
+class RolloutBatch(NamedTuple):
+    observations: np.ndarray
+    raw_actions: np.ndarray
+    rewards: np.ndarray
+    terminated: np.ndarray
+    truncated: np.ndarray
+    next_observations: np.ndarray
+
+
+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
+
+
+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, *states.shape), dtype=np.float32)
+        next_obs = np.empty_like(obs)
+        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:
+            if t == 0:
+                raw_act = np.empty((T, *raw_actions.shape), dtype=np.float32)
+            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
+            final_obs = infos.get('final_obs')
+            for idx in np.flatnonzero(d):
+                if final_obs is not None and final_obs[idx] is not None:
+                    next_obs[t, idx] = final_obs[idx]
+                elif truncated[idx] and not terminated[idx]:
+                    raise ValueError('truncated transitions require final_obs')
+
+        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

tinybaselines/training.py

@@ -0,0 +1,32 @@
+from time import perf_counter
+from typing import NamedTuple
+
+import numpy as np
+
+from .rollout import RolloutState, rollout
+
+
+class TrainingStep(NamedTuple):
+    timesteps: int
+    episode_returns: list[float]
+    rollout_seconds: float
+    update_seconds: float
+
+
+def learn(env, ppo, total_timesteps, rollout_steps):
+    """Train in full rollouts, yielding progress after each update. The timestep budget is a minimum."""
+    if total_timesteps < 0:
+        raise ValueError('total_timesteps must be nonnegative')
+    if rollout_steps <= 0 or rollout_steps % env.num_envs:
+        raise ValueError('rollout_steps must be a positive multiple of env.num_envs')
+
+    state = None
+    for timesteps in range(0, total_timesteps, rollout_steps):
+        if state is None:
+            state = RolloutState(env.reset()[0], np.zeros(env.num_envs))
+        started = perf_counter()
+        batch, returns = rollout(env, ppo.actor_critic, num_steps=rollout_steps, state=state)
+        collected = perf_counter()
+        ppo.update(batch)
+        updated = perf_counter()
+        yield TrainingStep(timesteps + batch.rewards.size, returns, collected - started, updated - collected)

tinybaselines-0.1.0a1.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,11 @@
+tinybaselines/__init__.py,sha256=YF5iiW_X6vwTmNRdiaQHMNgE7VMa2idpulh6IfqRzds,195
+tinybaselines/evaluation.py,sha256=3ti_T2gtJgpulgX1C0HEVxXITXfxPGkACsJnKfUOE_o,246
+tinybaselines/policies.py,sha256=u1fOphgpKufGYBquVdpE3cWhotBUUVcq79-XfQcL3DM,3358
+tinybaselines/ppo.py,sha256=PmChTu-iHEHIhCr1NxL60qK3AXs2jItQOPMj0zhcGY8,7877
+tinybaselines/rollout.py,sha256=9niFh9epX-W4P5rlzPc3b-UYrmHHqM7MOK-CnE4L9O8,2864
+tinybaselines/training.py,sha256=5VfZq4S7xlBs6PtdIgDP51Ei1uNGdBqi6f-SNESLYls,1171
+tinybaselines-0.1.0a1.dist-info/licenses/LICENSE,sha256=voIoq73Yv9HUCAMenLnIuZ-t4TJa_CfIO7TKUlwhWlc,1069
+tinybaselines-0.1.0a1.dist-info/METADATA,sha256=O_qFoN-6Jeja_wmcxoslMIdglQAK3bq1ZqLwQZeq9q0,3728
+tinybaselines-0.1.0a1.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+tinybaselines-0.1.0a1.dist-info/top_level.txt,sha256=P-g_X6Qa5L4-vsMCPSgUkolHGQMiDfxhra0_KSeEhpg,14
+tinybaselines-0.1.0a1.dist-info/RECORD,,

tinybaselines-0.1.0a1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

tinybaselines-0.1.0a1.dist-info/licenses/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.dist-info/top_level.txt

@@ -0,0 +1 @@
+tinybaselines