tinybaselines 0.1.0a2__tar.gz → 0.1.0a3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tinybaselines
-Version: 0.1.0a2
+Version: 0.1.0a3
 Summary: PPO in tinygrad.
 License-Expression: MIT
 Requires-Python: >=3.11
@@ -12,7 +12,7 @@ Dynamic: license-file
 
 # tinybaselines
 
-PPO in tinygrad. Continuous actions clipped to `[-1, 1]`.
+PPO in tinygrad for discrete and continuous actions. Continuous actions are clipped to `[-1, 1]`.
 
 ## LunarLander from PyPI
 
@@ -21,7 +21,7 @@ Use Python 3.11 and clang. The library install does not include Gymnasium.
 ```sh
 python3.11 -m venv .venv
 source .venv/bin/activate
-python -m pip install tinybaselines==0.1.0a2
+python -m pip install tinybaselines==0.1.0a3
 ```
 
 Install Gymnasium and Box2D to run this example:
@@ -72,10 +72,10 @@ actions by default; `deterministic=True` uses clipped policy means.
 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.
+PPO checkpoints contain the policy, 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.
+`GaussianActorCritic.save(path)` exports a Gaussian policy for inference.
 
 ## Repository examples
 
@@ -86,6 +86,18 @@ pip install -e . -r examples/lunar_lander/requirements.txt
 python -m examples.lunar_lander.train
 ```
 
+CartPole uses `CategoricalActorCritic`. The example trains, saves a checkpoint,
+and evaluates the reloaded policy over 10 episodes.
+
+```sh
+pip install -e . 'gymnasium==1.4.0'
+DEV=CPU BEAM=0 python -m examples.cartpole.train --checkpoint cartpole.safetensors
+```
+
+Three independent four-environment CPU runs averaged 500.00, 496.67, and 494.37
+after 102,400 transitions each. All cleared 475 over 100 fresh stochastic episodes
+per reloaded checkpoint. [Run records](examples/cartpole/results/cpu.jsonl).
+
 Pendulum uses three observations and one action. Its example rescales policy actions
 from `[-1, 1]` to torque in `[-2, 2]`.
 

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

@@ -1,6 +1,6 @@
 # tinybaselines
 
-PPO in tinygrad. Continuous actions clipped to `[-1, 1]`.
+PPO in tinygrad for discrete and continuous actions. Continuous actions are clipped to `[-1, 1]`.
 
 ## LunarLander from PyPI
 
@@ -9,7 +9,7 @@ Use Python 3.11 and clang. The library install does not include Gymnasium.
 ```sh
 python3.11 -m venv .venv
 source .venv/bin/activate
-python -m pip install tinybaselines==0.1.0a2
+python -m pip install tinybaselines==0.1.0a3
 ```
 
 Install Gymnasium and Box2D to run this example:
@@ -60,10 +60,10 @@ actions by default; `deterministic=True` uses clipped policy means.
 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.
+PPO checkpoints contain the policy, 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.
+`GaussianActorCritic.save(path)` exports a Gaussian policy for inference.
 
 ## Repository examples
 
@@ -74,6 +74,18 @@ pip install -e . -r examples/lunar_lander/requirements.txt
 python -m examples.lunar_lander.train
 ```
 
+CartPole uses `CategoricalActorCritic`. The example trains, saves a checkpoint,
+and evaluates the reloaded policy over 10 episodes.
+
+```sh
+pip install -e . 'gymnasium==1.4.0'
+DEV=CPU BEAM=0 python -m examples.cartpole.train --checkpoint cartpole.safetensors
+```
+
+Three independent four-environment CPU runs averaged 500.00, 496.67, and 494.37
+after 102,400 transitions each. All cleared 475 over 100 fresh stochastic episodes
+per reloaded checkpoint. [Run records](examples/cartpole/results/cpu.jsonl).
+
 Pendulum uses three observations and one action. Its example rescales policy actions
 from `[-1, 1]` to torque in `[-2, 2]`.
 

tinybaselines-0.1.0a3/examples/cartpole/results/cpu.jsonl

@@ -0,0 +1,4 @@
+{"event": "config", "source_commit": "536bfa324af5ce2793a2332844af36b4033a6cb2", "environment": "CartPole-v1", "num_envs": 4, "seed": null, "requested_timesteps": 102400, "rollout_steps": 1024, "evaluation_episodes": 100, "deterministic": false, "acceptance_mean": 475, "policy": {"state_dim": 4, "action_dim": 2, "hidden_dim": 64, "actor_layers": 2, "critic_layers": 2}, "ppo": {"pi_lr": 0.0010000000474974513, "vf_lr": 0.0010000000474974513, "gamma": 0.99, "lamda": 0.95, "K_epochs": 4, "eps_clip": 0.2, "batch_size": 256, "vf_coef": 0.5, "entropy_coef": 0.01}, "device": "CPU", "cpu_threads": 10, "beam": "0", "jit": "1", "python": "3.11.14", "platform": "macOS-27.0.1-arm64-arm-64bit", "versions": {"tinygrad": "0.14.0", "numpy": "2.4.6", "gymnasium": "1.1.1"}}
+{"event": "final", "run": 1, "timesteps": 102400, "mean_return": 500.0, "episode_returns": [500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0], "passed": true, "training_seconds": null, "evaluation_seconds": 14.94911529199453, "checkpoint": "run-1.safetensors", "note": "Repeated evaluation from the same checkpoint after a logging error."}
+{"event": "final", "run": 2, "timesteps": 102400, "mean_return": 496.67, "episode_returns": [500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 390.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 445.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 332.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0], "passed": true, "training_seconds": 56.7557016669889, "evaluation_seconds": 16.15585825004382, "checkpoint": "run-2.safetensors"}
+{"event": "final", "run": 3, "timesteps": 102400, "mean_return": 494.37, "episode_returns": [500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 96.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 341.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0, 500.0], "passed": true, "training_seconds": 55.03045945899794, "evaluation_seconds": 16.498802624992095, "checkpoint": "run-3.safetensors"}

tinybaselines-0.1.0a3/examples/cartpole/train.py

@@ -0,0 +1,33 @@
+import argparse
+from contextlib import closing
+from pathlib import Path
+from statistics import mean
+
+import gymnasium as gym
+
+from tinybaselines import CategoricalActorCritic, PPO, evaluate, learn
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Train PPO on CartPole-v1.')
+    parser.add_argument('--timesteps', type=int, default=40_960)
+    parser.add_argument('--checkpoint', type=Path, required=True)
+    args = parser.parse_args()
+
+    policy = CategoricalActorCritic(4, 2, 64, 2, 2)
+    ppo = PPO(policy, pi_lr=1e-3, vf_lr=1e-3, gamma=0.99, lamda=0.95, K_epochs=4,
+              eps_clip=0.2, batch_size=256, vf_coef=0.5, entropy_coef=0.01)
+    with closing(gym.vector.SyncVectorEnv(
+        [lambda: gym.make('CartPole-v1')] * 4,
+        autoreset_mode=gym.vector.AutoresetMode.SAME_STEP,
+    )) as env:
+        for step in learn(env, ppo, total_timesteps=args.timesteps, rollout_steps=1024):
+            print(f'{step.timesteps} transitions')
+        ppo.save(args.checkpoint)
+        policy = PPO.load(args.checkpoint).actor_critic
+        returns = evaluate(env, policy, num_episodes=10)
+        print(f'Mean return over {len(returns)} episodes: {mean(returns):.1f}')
+
+
+if __name__ == '__main__':
+    main()

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

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

tinybaselines-0.1.0a3/tests/test_cartpole.py

@@ -0,0 +1,39 @@
+from contextlib import closing
+from pathlib import Path
+from tempfile import TemporaryDirectory
+import unittest
+
+import gymnasium as gym
+import numpy as np
+
+from tinybaselines import CategoricalActorCritic, PPO, evaluate, learn
+
+
+class CartPoleTest(unittest.TestCase):
+    def test_categorical_policy_trains_with_ppo(self):
+        policy = CategoricalActorCritic(4, 2, 16, 1, 1)
+        ppo = PPO(policy, pi_lr=1e-3, vf_lr=1e-3, gamma=0.99, lamda=0.95, K_epochs=2,
+                  eps_clip=0.2, batch_size=64, vf_coef=0.5, entropy_coef=0.01)
+        weight = policy.actor.layers[-1].weight
+        before = weight.numpy().copy()
+        with closing(gym.vector.SyncVectorEnv([lambda: gym.make('CartPole-v1')] * 4,
+                                              autoreset_mode=gym.vector.AutoresetMode.SAME_STEP)) as env:
+            actions, _ = policy.act(env.reset()[0])
+            self.assertTrue(env.action_space.contains(actions))
+            steps = list(learn(env, ppo, total_timesteps=384, rollout_steps=128))
+            self.assertEqual(steps[-1].timesteps, 384)
+            self.assertFalse(np.array_equal(weight.numpy(), before))
+            observations, _ = env.reset()
+            with TemporaryDirectory() as directory:
+                path = Path(directory) / 'ppo.safetensors'
+                ppo.save(path)
+                restored = PPO.load(path)
+                self.assertIsInstance(restored.actor_critic, CategoricalActorCritic)
+                np.testing.assert_array_equal(restored.actor_critic.act(observations, deterministic=True)[0],
+                                              policy.act(observations, deterministic=True)[0])
+                state = restored.state_dict()
+                self.assertEqual(state.keys(), ppo.state_dict().keys())
+                for name, tensor in ppo.state_dict().items():
+                    np.testing.assert_array_equal(state[name].numpy(), tensor.numpy())
+                self.assertEqual(len(evaluate(env, restored.actor_critic, num_episodes=8)), 8)
+                self.assertEqual(next(learn(env, restored, total_timesteps=128, rollout_steps=128)).timesteps, 128)

{tinybaselines-0.1.0a2 → tinybaselines-0.1.0a3}/tests/test_checkpoint.py

@@ -1,3 +1,4 @@
+import json
 import tempfile
 import unittest
 from pathlib import Path
@@ -83,6 +84,17 @@ class CheckpointTest(unittest.TestCase):
             np.testing.assert_array_equal(first.actor_critic.log_std.numpy(), before)
             np.testing.assert_array_equal(second.actor_critic.log_std.numpy(), [0.25])
 
+    def test_loads_gaussian_checkpoint_without_policy_type(self):
+        ppo = self.make_ppo()
+        metadata = dict(format='tinybaselines.ppo.v1', policy_config=json.dumps(ppo.actor_critic.config),
+                        ppo_config=json.dumps(ppo.config))
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / 'legacy.safetensors'
+            lander.nn.state.safe_save(ppo.state_dict(), path, metadata=metadata)
+            restored = PPO.load(path)
+        self.assertIsInstance(restored.actor_critic, GaussianActorCritic)
+        np.testing.assert_array_equal(restored.actor_critic.observation_scale.numpy(), [2, 3, 4])
+
     def test_failed_save_preserves_previous_checkpoint(self):
         ppo = self.make_ppo()
 

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

@@ -6,12 +6,35 @@ import unittest
 from unittest.mock import patch
 
 import numpy as np
+from tinygrad import Tensor
 
 from tinybaselines import GaussianActorCritic, PPO, rollout
 from tinybaselines.rollout import RolloutBatch
 
 
 class CoreTest(unittest.TestCase):
+    def test_action_dtype_and_shape_survive_minibatching(self):
+        observations = np.zeros((4, 3), np.float32)
+        done = np.zeros(4, bool)
+        env = SimpleNamespace(num_envs=4, reset=lambda: (observations, {}),
+                              step=lambda actions: (observations, np.ones(4), done, done, {}))
+        policy = GaussianActorCritic(3, 2, 4, 1, 1)
+        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)
+        values = Tensor.zeros(8, 1, device=policy.device)
+        for actions in (np.arange(4, dtype=np.int32), np.arange(8, dtype=np.float32).reshape(4, 2) / 8):
+            with self.subTest(dtype=actions.dtype, shape=actions.shape):
+                with patch.object(policy, 'act', return_value=(actions, actions)):
+                    batch, _ = rollout(env, policy, num_steps=8)
+                self.assertEqual(batch.raw_actions.dtype, actions.dtype)
+                np.testing.assert_array_equal(batch.raw_actions, np.stack([actions, actions]))
+                with patch.object(policy, 'evaluate_actions', return_value=(values, values.squeeze(-1), None)), \
+                     patch.object(ppo, 'train_step') as step, patch('numpy.random.permutation', return_value=np.arange(8)):
+                    ppo.update(batch)
+                minibatches = np.concatenate([call.args[1].numpy() for call in step.call_args_list])
+                self.assertEqual(minibatches.dtype, actions.dtype)
+                np.testing.assert_array_equal(minibatches, np.concatenate([actions, actions]))
+
     def test_ppo_gets_actor_parameters_and_device_from_policy(self):
         model = GaussianActorCritic(3, 1, 4, 1, 1)
         weight = model.actor.layers[0].weight

tinybaselines-0.1.0a3/tests/test_learning.py

@@ -0,0 +1,89 @@
+import unittest
+
+import numpy as np
+from tinygrad import Tensor
+
+from tinybaselines import CategoricalActorCritic, GaussianActorCritic, PPO, learn
+
+
+class IdentityEnv:
+    num_envs = 4
+
+    def __init__(self, seed):
+        self.rng = np.random.default_rng(seed)
+
+    def reset(self):
+        self.observations = self.rng.uniform(-0.75, 0.75, (self.num_envs, 1)).astype(np.float32)
+        return self.observations, {}
+
+    def step(self, actions):
+        rewards = -np.square(actions[:, 0] - self.observations[:, 0])
+        observations, _ = self.reset()
+        return observations, rewards, np.ones(self.num_envs, bool), np.zeros(self.num_envs, bool), {}
+
+
+def identity_errors(seed):
+    Tensor.manual_seed(seed)
+    np.random.seed(seed)
+    policy = GaussianActorCritic(1, 1, 16, 1, 1)
+    ppo = PPO(policy, pi_lr=0.01, vf_lr=0.01, gamma=0.99, lamda=0.95, K_epochs=4,
+              eps_clip=0.2, batch_size=256, vf_coef=0.5, entropy_coef=0.0)
+    observations = np.linspace(-0.75, 0.75, 65, dtype=np.float32)[:, None]
+
+    def error():
+        actions, _ = policy.act(observations, deterministic=True)
+        return np.square(actions - observations).mean()
+
+    before = error()
+    for _ in learn(IdentityEnv(seed), ppo, total_timesteps=2048, rollout_steps=256):
+        pass
+    return before, error(), np.square(observations).mean()
+
+
+class DiscreteIdentityEnv:
+    num_envs = 4
+    observations = np.eye(4, dtype=np.float32)
+
+    def __init__(self, seed):
+        self.rng = np.random.default_rng(seed)
+
+    def reset(self):
+        self.targets = self.rng.integers(4, size=self.num_envs)
+        return self.observations[self.targets], {}
+
+    def step(self, actions):
+        rewards = (actions == self.targets).astype(np.float32)
+        observations, _ = self.reset()
+        return observations, rewards, np.ones(self.num_envs, bool), np.zeros(self.num_envs, bool), {}
+
+
+def discrete_identity_scores(seed):
+    Tensor.manual_seed(seed)
+    np.random.seed(seed)
+    env = DiscreteIdentityEnv(seed)
+    policy = CategoricalActorCritic(4, 4, 16, 1, 1)
+    ppo = PPO(policy, pi_lr=0.02, vf_lr=0.01, gamma=0.99, lamda=0.95, K_epochs=4,
+              eps_clip=0.2, batch_size=64, vf_coef=0.5, entropy_coef=0.0)
+    states = policy.preprocess(env.observations)
+    targets = Tensor(np.arange(4, dtype=np.int32), device=policy.device)
+
+    def score():
+        _, log_probs, _ = policy.evaluate_actions(states, targets)
+        return log_probs.exp().mean().item()
+
+    before = score()
+    for _ in learn(env, ppo, total_timesteps=1024, rollout_steps=128):
+        pass
+    return before, score()
+
+
+class LearningTest(unittest.TestCase):
+    def test_learns_identity(self):
+        self.addCleanup(np.random.set_state, np.random.get_state())
+        before, after, zero = identity_errors(seed=0)
+        self.assertLess(after, 0.25 * min(before, zero), (before, after, zero))
+
+    def test_learns_discrete_identity(self):
+        self.addCleanup(np.random.set_state, np.random.get_state())
+        before, after = discrete_identity_scores(seed=0)
+        self.assertGreater(after, max(before, 0.9), (before, after))

{tinybaselines-0.1.0a2 → tinybaselines-0.1.0a3}/tinybaselines/__init__.py

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

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

@@ -25,6 +25,8 @@ class MLP:
 
 
 class GaussianActorCritic:
+    _checkpoint_type = 'gaussian'
+
     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)
@@ -96,3 +98,50 @@ class GaussianActorCritic:
         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')
+
+
+class CategoricalActorCritic:
+    """Actor-critic for a single discrete action per observation."""
+
+    _checkpoint_type = 'categorical'
+
+    def __init__(self, state_dim, action_dim, hidden_dim, actor_layers, critic_layers):
+        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_jit = TinyJit(lambda states: self.actor(states).realize())
+        self.jit_shape = None
+
+    @property
+    def actor_parameters(self):
+        return nn.state.get_parameters(self.actor)
+
+    @property
+    def device(self):
+        return self.actor.layers[0].weight.device
+
+    def preprocess(self, observations):
+        return Tensor(observations.astype(np.float32), device=self.device)
+
+    def evaluate_actions(self, states, actions):
+        log_probs = self.actor(states).log_softmax(axis=-1)
+        selected = log_probs.gather(-1, actions.unsqueeze(-1)).squeeze(-1)
+        entropy = -(log_probs.exp() * log_probs).sum(axis=-1).mean()
+        return self.critic(states), selected, entropy
+
+    def act(self, state, deterministic=False):
+        states = self.preprocess(state)
+        if self.jit_shape != states.shape:
+            self.refresh_inference()
+            self.jit_shape = states.shape
+        logits = self.actor_jit(states).numpy()
+        scores = logits if deterministic else logits + np.random.gumbel(size=logits.shape)
+        actions = scores.argmax(axis=-1).astype(np.int32)
+        return actions, actions
+
+    def refresh_inference(self):
+        self.actor_jit.reset()
+
+    def state_dict(self):
+        return nn.state.get_state_dict(dict(actor=self.actor, critic=self.critic))

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

@@ -5,7 +5,10 @@ from tempfile import TemporaryDirectory
 import numpy as np
 from tinygrad import Context, Tensor, TinyJit, nn
 
-from .policies import GaussianActorCritic
+from .policies import CategoricalActorCritic, GaussianActorCritic
+
+_CHECKPOINT_POLICIES = {policy._checkpoint_type: policy for policy in (GaussianActorCritic, CategoricalActorCritic)}
+
 
 def clip_grad_norm(parameters, max_norm):
     gradients = [parameter.grad for parameter in parameters if parameter.grad is not None]
@@ -52,6 +55,7 @@ class PPO:
     def save(self, path):
         path = Path(path)
         metadata = dict(format='tinybaselines.ppo.v1', policy_config=json.dumps(self.actor_critic.config),
+                        policy_type=self.actor_critic._checkpoint_type,
                         ppo_config=json.dumps(self.config))
         with TemporaryDirectory(dir=path.parent) as directory:
             temporary = Path(directory) / path.name
@@ -64,7 +68,11 @@ class PPO:
         metadata = header.get('__metadata__', {})
         if metadata.get('format') != 'tinybaselines.ppo.v1':
             raise ValueError('unsupported PPO checkpoint format')
-        policy = GaussianActorCritic(**json.loads(metadata['policy_config']))
+        policy_type = metadata.get('policy_type', 'gaussian')
+        policy_class = _CHECKPOINT_POLICIES.get(policy_type)
+        if policy_class is None:
+            raise ValueError(f'unsupported PPO checkpoint policy type: {policy_type!r}')
+        policy = policy_class(**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()
@@ -114,7 +122,7 @@ class PPO:
     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_np = batch.raw_actions.reshape(T * N, *batch.raw_actions.shape[2:])
         raw_act_flat = Tensor(raw_act_np, device=self.device)
         old_value, old_logprobs, _ = self.actor_critic.evaluate_actions(obs_flat, raw_act_flat)
         old_logprobs_np = old_logprobs.numpy()

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

@@ -58,7 +58,7 @@ def rollout(env, actor_critic, num_steps=None, num_episodes=None, deterministic=
         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)
+                raw_act = np.empty((T, *raw_actions.shape), dtype=raw_actions.dtype)
             obs[t], raw_act[t] = states, raw_actions
 
         states, rewards, terminated, truncated, infos = env.step(actions)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tinybaselines
-Version: 0.1.0a2
+Version: 0.1.0a3
 Summary: PPO in tinygrad.
 License-Expression: MIT
 Requires-Python: >=3.11
@@ -12,7 +12,7 @@ Dynamic: license-file
 
 # tinybaselines
 
-PPO in tinygrad. Continuous actions clipped to `[-1, 1]`.
+PPO in tinygrad for discrete and continuous actions. Continuous actions are clipped to `[-1, 1]`.
 
 ## LunarLander from PyPI
 
@@ -21,7 +21,7 @@ Use Python 3.11 and clang. The library install does not include Gymnasium.
 ```sh
 python3.11 -m venv .venv
 source .venv/bin/activate
-python -m pip install tinybaselines==0.1.0a2
+python -m pip install tinybaselines==0.1.0a3
 ```
 
 Install Gymnasium and Box2D to run this example:
@@ -72,10 +72,10 @@ actions by default; `deterministic=True` uses clipped policy means.
 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.
+PPO checkpoints contain the policy, 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.
+`GaussianActorCritic.save(path)` exports a Gaussian policy for inference.
 
 ## Repository examples
 
@@ -86,6 +86,18 @@ pip install -e . -r examples/lunar_lander/requirements.txt
 python -m examples.lunar_lander.train
 ```
 
+CartPole uses `CategoricalActorCritic`. The example trains, saves a checkpoint,
+and evaluates the reloaded policy over 10 episodes.
+
+```sh
+pip install -e . 'gymnasium==1.4.0'
+DEV=CPU BEAM=0 python -m examples.cartpole.train --checkpoint cartpole.safetensors
+```
+
+Three independent four-environment CPU runs averaged 500.00, 496.67, and 494.37
+after 102,400 transitions each. All cleared 475 over 100 fresh stochastic episodes
+per reloaded checkpoint. [Run records](examples/cartpole/results/cpu.jsonl).
+
 Pendulum uses three observations and one action. Its example rescales policy actions
 from `[-1, 1]` to torque in `[-2, 2]`.
 

{tinybaselines-0.1.0a2 → tinybaselines-0.1.0a3}/tinybaselines.egg-info/SOURCES.txt

@@ -4,6 +4,8 @@ README.md
 pyproject.toml
 requirements.txt
 examples/__init__.py
+examples/cartpole/train.py
+examples/cartpole/results/cpu.jsonl
 examples/lunar_lander/__init__.py
 examples/lunar_lander/config.py
 examples/lunar_lander/reference.py
@@ -27,6 +29,7 @@ examples/pendulum/train.py
 examples/pendulum/results/cpu-2.jsonl
 examples/pendulum/results/cpu-3.jsonl
 examples/pendulum/results/cpu.jsonl
+tests/test_cartpole.py
 tests/test_checkpoint.py
 tests/test_core.py
 tests/test_evaluation.py

tinybaselines-0.1.0a2/tests/test_learning.py

@@ -1,47 +0,0 @@
-import unittest
-
-import numpy as np
-from tinygrad import Tensor
-
-from tinybaselines import GaussianActorCritic, PPO, learn
-
-
-class IdentityEnv:
-    num_envs = 4
-
-    def __init__(self, seed):
-        self.rng = np.random.default_rng(seed)
-
-    def reset(self):
-        self.observations = self.rng.uniform(-0.75, 0.75, (self.num_envs, 1)).astype(np.float32)
-        return self.observations, {}
-
-    def step(self, actions):
-        rewards = -np.square(actions[:, 0] - self.observations[:, 0])
-        observations, _ = self.reset()
-        return observations, rewards, np.ones(self.num_envs, bool), np.zeros(self.num_envs, bool), {}
-
-
-def identity_errors(seed):
-    Tensor.manual_seed(seed)
-    np.random.seed(seed)
-    policy = GaussianActorCritic(1, 1, 16, 1, 1)
-    ppo = PPO(policy, pi_lr=0.01, vf_lr=0.01, gamma=0.99, lamda=0.95, K_epochs=4,
-              eps_clip=0.2, batch_size=256, vf_coef=0.5, entropy_coef=0.0)
-    observations = np.linspace(-0.75, 0.75, 65, dtype=np.float32)[:, None]
-
-    def error():
-        actions, _ = policy.act(observations, deterministic=True)
-        return np.square(actions - observations).mean()
-
-    before = error()
-    for _ in learn(IdentityEnv(seed), ppo, total_timesteps=2048, rollout_steps=256):
-        pass
-    return before, error(), np.square(observations).mean()
-
-
-class LearningTest(unittest.TestCase):
-    def test_learns_identity(self):
-        self.addCleanup(np.random.set_state, np.random.get_state())
-        before, after, zero = identity_errors(seed=0)
-        self.assertLess(after, 0.25 * min(before, zero), (before, after, zero))