tinybaselines 0.1.0a1__tar.gz → 0.1.0a3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tinybaselines
-Version: 0.1.0a1
+Version: 0.1.0a3
 Summary: PPO in tinygrad.
 License-Expression: MIT
 Requires-Python: >=3.11
@@ -12,44 +12,53 @@ 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]`.
 
-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.0a3
 ```
 
-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.
@@ -63,18 +72,32 @@ 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
 
-The full LunarLander trainer, evaluation, and PyTorch reference live in
-`examples/lunar_lander`.
+From a checkout, run the full LunarLander trainer:
 
 ```sh
+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.0a1 → tinybaselines-0.1.0a3}/README.md

@@ -1,43 +1,52 @@
 # 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]`.
 
-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.0a3
 ```
 
-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.
@@ -51,18 +60,32 @@ 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
 
-The full LunarLander trainer, evaluation, and PyTorch reference live in
-`examples/lunar_lander`.
+From a checkout, run the full LunarLander trainer:
 
 ```sh
+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.0a1 → tinybaselines-0.1.0a3}/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.0a3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tinybaselines"
-version = "0.1.0a1"
+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.0a1 → 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.0a1 → tinybaselines-0.1.0a3}/tests/test_core.py

@@ -6,12 +6,53 @@ 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
+        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.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.0a1 → tinybaselines-0.1.0a3}/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.0a3}/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.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.0a1 → tinybaselines-0.1.0a3}/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):
@@ -17,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)
@@ -31,12 +41,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
 
@@ -73,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.0a1 → tinybaselines-0.1.0a3}/tinybaselines/ppo.py

@@ -5,15 +5,9 @@ from tempfile import TemporaryDirectory
 import numpy as np
 from tinygrad import Context, Tensor, TinyJit, nn
 
-from .policies import GaussianActorCritic
+from .policies import CategoricalActorCritic, 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)
+_CHECKPOINT_POLICIES = {policy._checkpoint_type: policy for policy in (GaussianActorCritic, CategoricalActorCritic)}
 
 
 def clip_grad_norm(parameters, max_norm):
@@ -29,9 +23,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)
@@ -57,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
@@ -69,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()
@@ -89,14 +92,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):
@@ -121,10 +122,10 @@ 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_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.0a3}/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))
@@ -54,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.0a1 → tinybaselines-0.1.0a3/tinybaselines.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tinybaselines
-Version: 0.1.0a1
+Version: 0.1.0a3
 Summary: PPO in tinygrad.
 License-Expression: MIT
 Requires-Python: >=3.11
@@ -12,44 +12,53 @@ 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]`.
 
-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.0a3
 ```
 
-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.
@@ -63,18 +72,32 @@ 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
 
-The full LunarLander trainer, evaluation, and PyTorch reference live in
-`examples/lunar_lander`.
+From a checkout, run the full LunarLander trainer:
 
 ```sh
+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.0a1 → 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.0a1/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))