tinybaselines 0.1.0a1__tar.gz → 0.1.0a2__tar.gz

{tinybaselines-0.1.0a1/tinybaselines.egg-info → tinybaselines-0.1.0a2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tinybaselines
-Version: 0.1.0a1
+Version: 0.1.0a2
 Summary: PPO in tinygrad.
 License-Expression: MIT
 Requires-Python: >=3.11
@@ -14,42 +14,51 @@ Dynamic: license-file
 
 PPO in tinygrad. Continuous actions clipped to `[-1, 1]`.
 
-Python 3.11 and clang.
+## LunarLander from PyPI
+
+Use Python 3.11 and clang. The library install does not include Gymnasium.
 
 ```sh
 python3.11 -m venv .venv
 source .venv/bin/activate
-pip install -e . -r examples/lunar_lander/requirements.txt
+python -m pip install tinybaselines==0.1.0a2
 ```
 
-Train, save, and load.
+Install Gymnasium and Box2D to run this example:
+
+```sh
+python -m pip install 'gymnasium[box2d]==1.4.0'
+```
+
+Save this as `train_lander.py`. It runs one PPO update and evaluates the reloaded checkpoint.
 
 ```python
+from contextlib import closing
 from statistics import mean
+
+import gymnasium as gym
+
 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)
+                             actor_layers=4, critic_layers=4,
+                             observation_scale=[10, 6.666, 5, 7.5, 1, 2.5, 1, 1])
 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()
+with closing(gym.vector.SyncVectorEnv(
+    [lambda: gym.make('LunarLander-v3', continuous=True)] * 4,
+    autoreset_mode=gym.vector.AutoresetMode.SAME_STEP,
+)) as env:
+    for step in learn(env, ppo, total_timesteps=20_000, rollout_steps=20_000):
+        print(step.timesteps, step.rollout_seconds, step.update_seconds)
+    ppo.save('ppo.safetensors')
+    policy = PPO.load('ppo.safetensors').actor_critic
+    print(mean(evaluate(env, policy, num_episodes=10)))
+```
 
-eval_env = make_env(4, asynchronous=False)
-returns = evaluate(eval_env, policy, num_episodes=100)
-eval_env.close()
-print(mean(returns))
+```sh
+DEV=CPU BEAM=0 python train_lander.py
 ```
 
 `total_timesteps` and `rollout_steps` count transitions across all environments.
@@ -68,10 +77,12 @@ Continue training by passing the loaded PPO to `learn` with a new transition bud
 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`.
+## Repository examples
+
+From a checkout, run the full LunarLander trainer:
 
 ```sh
+pip install -e . -r examples/lunar_lander/requirements.txt
 python -m examples.lunar_lander.train
 ```
 

{tinybaselines-0.1.0a1 → tinybaselines-0.1.0a2}/README.md

@@ -2,42 +2,51 @@
 
 PPO in tinygrad. Continuous actions clipped to `[-1, 1]`.
 
-Python 3.11 and clang.
+## LunarLander from PyPI
+
+Use Python 3.11 and clang. The library install does not include Gymnasium.
 
 ```sh
 python3.11 -m venv .venv
 source .venv/bin/activate
-pip install -e . -r examples/lunar_lander/requirements.txt
+python -m pip install tinybaselines==0.1.0a2
 ```
 
-Train, save, and load.
+Install Gymnasium and Box2D to run this example:
+
+```sh
+python -m pip install 'gymnasium[box2d]==1.4.0'
+```
+
+Save this as `train_lander.py`. It runs one PPO update and evaluates the reloaded checkpoint.
 
 ```python
+from contextlib import closing
 from statistics import mean
+
+import gymnasium as gym
+
 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)
+                             actor_layers=4, critic_layers=4,
+                             observation_scale=[10, 6.666, 5, 7.5, 1, 2.5, 1, 1])
 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()
+with closing(gym.vector.SyncVectorEnv(
+    [lambda: gym.make('LunarLander-v3', continuous=True)] * 4,
+    autoreset_mode=gym.vector.AutoresetMode.SAME_STEP,
+)) as env:
+    for step in learn(env, ppo, total_timesteps=20_000, rollout_steps=20_000):
+        print(step.timesteps, step.rollout_seconds, step.update_seconds)
+    ppo.save('ppo.safetensors')
+    policy = PPO.load('ppo.safetensors').actor_critic
+    print(mean(evaluate(env, policy, num_episodes=10)))
+```
 
-eval_env = make_env(4, asynchronous=False)
-returns = evaluate(eval_env, policy, num_episodes=100)
-eval_env.close()
-print(mean(returns))
+```sh
+DEV=CPU BEAM=0 python train_lander.py
 ```
 
 `total_timesteps` and `rollout_steps` count transitions across all environments.
@@ -56,10 +65,12 @@ Continue training by passing the loaded PPO to `learn` with a new transition bud
 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`.
+## Repository examples
+
+From a checkout, run the full LunarLander trainer:
 
 ```sh
+pip install -e . -r examples/lunar_lander/requirements.txt
 python -m examples.lunar_lander.train
 ```
 

{tinybaselines-0.1.0a1 → tinybaselines-0.1.0a2}/examples/lunar_lander/train.py

@@ -12,7 +12,7 @@ from tinygrad import Tensor, TinyJit, nn, Context
 from tqdm import trange
 from .config import ENV_NAME, OBS_SCALE
 from tinybaselines import GaussianActorCritic, PPO, evaluate, learn
-from tinybaselines.ppo import clip_grad_norm, normal_log_prob
+from tinybaselines.ppo import clip_grad_norm
 from tinybaselines.rollout import RolloutBatch, RolloutState, rollout, track_episode_returns
 
 warnings.filterwarnings('ignore', message='pkg_resources is deprecated')

{tinybaselines-0.1.0a1 → tinybaselines-0.1.0a2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tinybaselines"
-version = "0.1.0a1"
+version = "0.1.0a2"
 description = "PPO in tinygrad."
 readme = "README.md"
 license = "MIT"

{tinybaselines-0.1.0a1 → tinybaselines-0.1.0a2}/tests/test_core.py

@@ -12,6 +12,24 @@ from tinybaselines.rollout import RolloutBatch
 
 
 class CoreTest(unittest.TestCase):
+    def test_ppo_gets_actor_parameters_and_device_from_policy(self):
+        model = GaussianActorCritic(3, 1, 4, 1, 1)
+        weight = model.actor.layers[0].weight
+        policy = SimpleNamespace(actor_parameters=[weight], device=weight.device, critic=model.critic)
+        ppo = PPO(policy, pi_lr=1e-3, vf_lr=1e-3, gamma=0.99, lamda=0.95, K_epochs=1,
+                  eps_clip=0.2, batch_size=4, vf_coef=1.0, entropy_coef=0.0)
+        self.assertEqual(ppo.device, policy.device)
+        self.assertIs(ppo.pi_optimizer.params[0], weight)
+
+    def test_rejects_nonpositive_update_settings(self):
+        policy = GaussianActorCritic(3, 1, 4, 1, 1)
+        settings = dict(pi_lr=1e-3, vf_lr=1e-3, gamma=0.99, lamda=0.95, K_epochs=1,
+                        eps_clip=0.2, batch_size=4, vf_coef=1.0, entropy_coef=0.0)
+        for name in ('K_epochs', 'batch_size'):
+            for value in (0, -1):
+                with self.subTest(name=name, value=value), self.assertRaisesRegex(ValueError, f'{name} must be positive'):
+                    PPO(policy, **(settings | {name: value}))
+
     def test_terminals_do_not_require_final_observations(self):
         observations = np.zeros((2, 3), np.float32)
         reset_observations = np.full((2, 3), 100, np.float32)

{tinybaselines-0.1.0a1 → tinybaselines-0.1.0a2}/tests/test_parity.py

@@ -5,6 +5,7 @@ import numpy as np
 import torch
 
 from examples.lunar_lander import train as lander, reference as reference_lander
+from tinybaselines.policies import normal_log_prob
 
 
 class ParityTest(unittest.TestCase):
@@ -205,7 +206,7 @@ class ParityTest(unittest.TestCase):
         torch_mean = torch.tensor(means, requires_grad=True)
         torch_log_std = torch.tensor(log_std, requires_grad=True)
         with lander.Context(TRAINING=1):
-            actual = lander.normal_log_prob(tiny_mean, tiny_log_std, lander.Tensor(actions))
+            actual = normal_log_prob(tiny_mean, tiny_log_std, lander.Tensor(actions))
             actual.sum().backward()
         expected = torch.distributions.Normal(torch_mean, torch_log_std.exp()).log_prob(torch.tensor(actions)).sum(-1)
         expected.sum().backward()

{tinybaselines-0.1.0a1 → tinybaselines-0.1.0a2}/tests/test_pendulum.py

@@ -1,6 +1,7 @@
 from types import SimpleNamespace
 import unittest
 
+import gymnasium as gym
 import numpy as np
 
 from examples.pendulum.env import make_env
@@ -8,6 +9,14 @@ from tinybaselines import GaussianActorCritic, PPO, learn, rollout
 
 
 class PendulumTest(unittest.TestCase):
+    def test_rejects_default_vector_autoreset(self):
+        env = gym.vector.SyncVectorEnv([lambda: gym.make('Pendulum-v1')] * 4)
+        self.addCleanup(env.close)
+        actions = np.zeros((4, 1), np.float32)
+        policy = SimpleNamespace(act=lambda obs, deterministic: (actions, actions))
+        with self.assertRaisesRegex(ValueError, 'autoreset_mode=SAME_STEP'):
+            rollout(env, policy, num_steps=8)
+
     def test_normalized_actions_cover_torque_range(self):
         env = make_env(4)
         try:

{tinybaselines-0.1.0a1 → tinybaselines-0.1.0a2}/tinybaselines/policies.py

@@ -3,6 +3,14 @@ import json
 import numpy as np
 from tinygrad import Tensor, TinyJit, nn
 
+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)
+
 
 class MLP:
     def __init__(self, input_dim, output_dim, hidden_dim, layer_count):
@@ -31,12 +39,27 @@ class GaussianActorCritic:
         self.actor_jit = TinyJit(lambda state: self.actor(state).realize())
         self.refresh_inference()
 
+    @property
+    def actor_parameters(self):
+        return nn.state.get_parameters(self.actor) + [self.log_std]
+
+    @property
+    def device(self):
+        return self.log_std.device
+
     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 evaluate_actions(self, states, actions):
+        """Return values, action log probabilities, and mean entropy for preprocessed states."""
+        means, log_std, values = self(states)
+        log_probs = normal_log_prob(means, log_std, actions)
+        entropy = (log_std + NORMAL_ENTROPY_CONSTANT).sum()
+        return values, log_probs, entropy
+
     def preprocess(self, observations):
         return Tensor(observations.astype(np.float32), device=self.observation_scale.device) * self.observation_scale
 

{tinybaselines-0.1.0a1 → tinybaselines-0.1.0a2}/tinybaselines/ppo.py

@@ -7,15 +7,6 @@ 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:
@@ -29,9 +20,13 @@ def clip_grad_norm(parameters, max_norm):
 
 class PPO:
     def __init__(self, actor_critic, pi_lr, vf_lr, gamma, lamda, K_epochs, eps_clip, batch_size, vf_coef, entropy_coef):
+        if K_epochs <= 0:
+            raise ValueError('K_epochs must be positive')
+        if batch_size <= 0:
+            raise ValueError('batch_size must be positive')
         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.device = actor_critic.device
+        self.actor_parameters = actor_critic.actor_parameters
         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)
@@ -89,14 +84,12 @@ class PPO:
         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)
+        state_values, action_logprobs, entropy = self.actor_critic.evaluate_actions(batch_states, 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):
@@ -123,8 +116,8 @@ class PPO:
         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_value, old_logprobs, _ = self.actor_critic.evaluate_actions(obs_flat, raw_act_flat)
+        old_logprobs_np = old_logprobs.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)

{tinybaselines-0.1.0a1 → tinybaselines-0.1.0a2}/tinybaselines/rollout.py

@@ -33,6 +33,10 @@ def rollout(env, actor_critic, num_steps=None, num_episodes=None, deterministic=
     if num_episodes is not None and num_episodes <= 0:
         raise ValueError('num_episodes must be positive')
 
+    mode = (getattr(env, 'metadata', None) or {}).get('autoreset_mode')
+    if getattr(mode, 'value', mode) not in (None, 'SameStep'):
+        raise ValueError(f'Expected autoreset_mode=SAME_STEP, got {mode!r}')
+
     N = env.num_envs
     if state is None:
         state = RolloutState(env.reset()[0], np.zeros(N))

{tinybaselines-0.1.0a1 → tinybaselines-0.1.0a2/tinybaselines.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tinybaselines
-Version: 0.1.0a1
+Version: 0.1.0a2
 Summary: PPO in tinygrad.
 License-Expression: MIT
 Requires-Python: >=3.11
@@ -14,42 +14,51 @@ Dynamic: license-file
 
 PPO in tinygrad. Continuous actions clipped to `[-1, 1]`.
 
-Python 3.11 and clang.
+## LunarLander from PyPI
+
+Use Python 3.11 and clang. The library install does not include Gymnasium.
 
 ```sh
 python3.11 -m venv .venv
 source .venv/bin/activate
-pip install -e . -r examples/lunar_lander/requirements.txt
+python -m pip install tinybaselines==0.1.0a2
 ```
 
-Train, save, and load.
+Install Gymnasium and Box2D to run this example:
+
+```sh
+python -m pip install 'gymnasium[box2d]==1.4.0'
+```
+
+Save this as `train_lander.py`. It runs one PPO update and evaluates the reloaded checkpoint.
 
 ```python
+from contextlib import closing
 from statistics import mean
+
+import gymnasium as gym
+
 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)
+                             actor_layers=4, critic_layers=4,
+                             observation_scale=[10, 6.666, 5, 7.5, 1, 2.5, 1, 1])
 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()
+with closing(gym.vector.SyncVectorEnv(
+    [lambda: gym.make('LunarLander-v3', continuous=True)] * 4,
+    autoreset_mode=gym.vector.AutoresetMode.SAME_STEP,
+)) as env:
+    for step in learn(env, ppo, total_timesteps=20_000, rollout_steps=20_000):
+        print(step.timesteps, step.rollout_seconds, step.update_seconds)
+    ppo.save('ppo.safetensors')
+    policy = PPO.load('ppo.safetensors').actor_critic
+    print(mean(evaluate(env, policy, num_episodes=10)))
+```
 
-eval_env = make_env(4, asynchronous=False)
-returns = evaluate(eval_env, policy, num_episodes=100)
-eval_env.close()
-print(mean(returns))
+```sh
+DEV=CPU BEAM=0 python train_lander.py
 ```
 
 `total_timesteps` and `rollout_steps` count transitions across all environments.
@@ -68,10 +77,12 @@ Continue training by passing the loaded PPO to `learn` with a new transition bud
 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`.
+## Repository examples
+
+From a checkout, run the full LunarLander trainer:
 
 ```sh
+pip install -e . -r examples/lunar_lander/requirements.txt
 python -m examples.lunar_lander.train
 ```