tinybaselines 0.1.0a2__tar.gz → 0.1.0a4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tinybaselines
-Version: 0.1.0a2
+Version: 0.1.0a4
 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.0a4
 ```
 
 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,24 @@ pip install -e . -r examples/lunar_lander/requirements.txt
 python -m examples.lunar_lander.train
 ```
 
+CPU batch padding is off by default. Set `TINYBASELINES_CPU_PADDING=1` before
+starting Python to try it. In three runs per variant at 420,000 transitions,
+mean training time was 296.5 s without padding and 282.1 s with it. Sampled
+evaluation cleared 200 in two unpadded runs and no padded runs.
+[Settings and both evaluation modes](examples/lunar_lander/results/cpu-padding.jsonl).
+
+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.0a4}/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.0a4
 ```
 
 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,24 @@ pip install -e . -r examples/lunar_lander/requirements.txt
 python -m examples.lunar_lander.train
 ```
 
+CPU batch padding is off by default. Set `TINYBASELINES_CPU_PADDING=1` before
+starting Python to try it. In three runs per variant at 420,000 transitions,
+mean training time was 296.5 s without padding and 282.1 s with it. Sampled
+evaluation cleared 200 in two unpadded runs and no padded runs.
+[Settings and both evaluation modes](examples/lunar_lander/results/cpu-padding.jsonl).
+
+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.0a4/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.0a4/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.0a4/examples/lunar_lander/results/cpu-padding.jsonl

@@ -0,0 +1,7 @@
+{"event":"config","platform":"macOS-27.0.1-arm64-arm-64bit","python":"3.11.14","versions":{"tinygrad":"0.14.0","numpy":"2.4.6","gymnasium":"1.1.1"},"environment":"LunarLanderContinuous-v3","num_envs":4,"seed":null,"training_vector":"async","evaluation_vector":"sync","device":"CPU","cpu_threads":10,"beam":"0","jit":"1","budget":420000,"rollout_steps":20000,"evaluation_episodes_per_mode":100,"policy":{"state_dim":8,"action_dim":2,"hidden_dim":128,"actor_layers":4,"critic_layers":4},"observation_scale":[10.0,6.665999889373779,5.0,7.5,1.0,2.5,1.0,1.0],"ppo":{"pi_lr":0.0010000000474974513,"vf_lr":0.0010000000474974513,"gamma":0.99,"lamda":0.95,"K_epochs":20,"eps_clip":0.2,"batch_size":10000,"vf_coef":1.0,"entropy_coef":0.001},"source_commits":{"unpadded":"18ce701612a3f1be0b4bd6f7b9894bb6b46c5d6e","padded":"655c1ba47ffeb1a6c605950983c33f503e18752c"},"run_order":["unpadded-1","padded-1","padded-2","unpadded-2","unpadded-3","padded-3"],"training_timer":"learn loop only"}
+{"event":"result","variant":"unpadded","trial":1,"training_seconds":297.1832159580081,"warm_rollout_seconds":2.2340254170121625,"warm_update_seconds":10.916082541982178,"action_std":[0.31564828753471375,0.3440491557121277],"evaluations":[{"deterministic":false,"mean_return":240.39174186825218,"evaluation_seconds":5.321839792013634,"episode_returns":[35.6561998363139,298.2771221830735,268.57741958732026,290.42095377360874,276.8006851632731,309.0616069626583,281.4620396299208,285.4142213202755,250.49292237984164,276.6037684694031,246.5324525017868,263.4837714551411,-28.77169000500193,261.2897574645251,298.52434313216577,269.9602674627298,298.4109806652238,269.61712372180654,311.92114194143113,265.6203014490535,294.53717418425066,236.36969855774544,265.69730374237975,224.81631223700268,258.23763008274403,287.43529122100733,291.55589296506975,274.1802587545968,273.5469665280905,304.6566413168965,295.76170940904683,265.8544325421157,72.36076580090472,283.2886311171228,263.24170215783715,268.24481209887267,179.26839473466765,247.92899934014076,280.4431757998325,287.04485232029026,287.8778125161239,268.3979651404212,300.164983629183,183.62616164277367,252.05631656603646,245.92920810940848,284.5719832174914,50.76827046711526,297.3993088995197,277.7782534402006,259.5722911802263,205.69340053797012,61.7794818992198,55.07359228928189,232.6359441891852,240.7444528189691,264.64346522388996,291.90872404091925,75.93059205369985,275.2869238365895,245.35374619578906,247.10961211741085,249.59255951650516,55.34726696007064,287.01953221093265,278.898300096003,270.0359830117993,252.51464057549236,276.8804616655278,302.17821749488934,244.8760419153127,312.87729474617237,240.03374520990891,173.30274859057525,59.99889564513838,256.02627156009567,281.8557025447259,258.4691514052356,19.055609252153815,269.8411328740911,276.6224877800033,282.18448247545473,280.5570219248604,277.0991605740026,275.5729371407476,288.1190334557397,253.3138976045431,274.5736882310889,275.54395066201073,54.556902554066994,290.5114628698576,282.83914845060787,296.8156560258374,263.8734723532773,268.69685152385796,264.30319044036185,259.8529275154999,122.18160703205469,18.8657923451385,256.15843657498806]},{"deterministic":true,"mean_return":94.52528108461252,"evaluation_seconds":3.682486500008963,"episode_returns":[46.24559016097538,-25.570391633405052,253.51603743227216,231.09290586823437,-43.68081142518827,45.09982496046595,-180.53476015463582,198.87334605374235,276.39900890509097,303.2218769793467,196.20694763377645,36.03221289485427,-0.8415263552276997,5.923853721805699,195.51132259138328,-16.362432242910046,-172.17480333066266,-55.52220334791498,-32.75020089326007,-40.50244746027197,251.73716126050368,199.42262278842168,184.6439986774401,-52.52783693246649,251.7204933336076,-23.866719434377202,-134.84626001582544,296.75022173524013,-60.726140638374744,-36.823590572525205,165.32212507785255,283.45516458950647,-51.25750660727029,203.39556411952532,190.486419196499,-13.222272925204578,243.46253378893732,275.5733915749505,-17.689642230540727,-46.63206068111274,-113.57553675541234,-10.506213085169392,232.33753686561684,-58.103691835192144,223.77951539910464,222.70188534289179,5.032766989781464,199.5530545791565,213.18567182514636,233.79379609829647,58.30959683754119,255.07445062068572,279.7201348022994,-20.3436002260525,249.80380261238207,26.71094537572212,230.9836318713347,62.00187823093211,49.298701571962994,263.0559526247364,262.3087685543362,260.3042174773011,239.01150171744052,9.053917706584102,-9.608981393628383,152.05778052948267,191.4052427042376,14.889552297234644,-61.073073987852645,172.95711316490605,8.653999936412333,251.3726824776721,182.38963691856165,-60.884975389116704,177.0722056283517,-46.55661422007239,221.75266884123567,311.1795407367281,-58.40952322270949,225.24341815456398,-39.570045259804346,-1.173470521161903,-58.873466910384714,274.6263780805373,-30.72487695764059,-57.1746846821751,295.0731049899574,-26.18556722577499,255.22696716127052,-15.98763360654526,279.25723817674066,-36.408218859472825,-90.79216409672452,-4.263630237586057,248.9643282498574,-31.241744292686647,-13.14671456319104,-26.682568029280034,145.94912888925253,11.16137331537324]}]}
+{"event":"result","variant":"padded","trial":1,"training_seconds":290.42981791700004,"warm_rollout_seconds":2.252680917037651,"warm_update_seconds":10.432943666004576,"action_std":[0.34810739755630493,0.4137430191040039],"evaluations":[{"deterministic":false,"mean_return":180.62778245182855,"evaluation_seconds":11.679198166995775,"episode_returns":[200.8066728790336,270.133917641077,139.39231521137037,146.2501979720148,241.57681362780022,282.4830751590748,261.9765563475175,133.95294801170508,132.27241746914865,226.2673230273345,226.965767068635,224.03082391505285,171.1370636003007,148.40102897509104,188.7563251560003,194.8019199251024,105.78234711735841,144.76375327830633,161.80114791808245,165.62427351717267,164.7074999098417,150.10123287046383,165.59097789890978,290.4533296369576,125.00912209645797,188.85911298487054,147.11662982255964,296.2982609969252,153.60129950925565,111.03670230876463,272.91287849441403,245.86834061281988,177.5469357761804,194.62527928292536,139.76904804072598,139.01124361473128,141.77055930762222,248.8975937403574,221.5788705159723,159.8867971503815,180.8001054494639,167.32166860394892,243.7965282776548,117.29919629987788,134.05914307337855,260.1621844745094,158.89120284227118,32.19286884334641,233.0295613473187,144.5702387973397,121.54157701606229,98.69923887954756,231.38220403719856,216.634377312762,194.6596005862297,117.17236694861546,124.51855789837315,104.81572057546124,288.6039328497314,161.1052538908393,178.29767465384924,153.35209200557114,205.39000412329392,129.9788672512846,274.9058291694384,142.8993219589483,135.83423533288777,117.99111878308078,247.20291813769285,290.94726292268604,48.644321065841005,279.11147805740234,228.78238616494716,145.06815809196584,176.83345656044844,156.18415622062147,149.4185977826854,139.44182483261665,174.71670310187986,270.32548201404796,260.3873549684042,125.43595390369936,194.92901489567788,198.26738611355967,234.21225477699764,258.90770743970216,140.3425631977819,242.69382750347168,260.54790072135586,292.8777661688539,138.74775172560393,129.11246035612342,-9.649273782287052,137.95253107964632,253.72381759758017,183.22842622608837,69.3149403635436,145.28666203172648,145.47757795250703,182.58180331938368]},{"deterministic":true,"mean_return":201.60909536765791,"evaluation_seconds":4.318445916986093,"episode_returns":[283.5519351551185,7.565182985865022,268.7141463838077,215.3685792362784,263.84940941016333,224.26113809071137,290.89160161969824,214.9322788109273,245.10058367877517,287.5178720808873,231.1523271965112,268.09373511105883,257.52799069569426,-103.17062167688175,224.80702557843037,122.14056292303295,242.4927280910718,230.90212860115145,237.30710078055867,236.73233576118704,179.25729830353336,217.05767639242526,256.42366549870104,227.29623968046775,300.98050161973725,252.45407092867092,213.48831480020442,227.4161499640053,6.261973210889863,268.07392693798226,-178.01705986519727,155.7084069519873,267.12520540188325,265.1181504052979,289.4511907721086,260.9643463950135,294.68891652233674,-17.066506471930495,209.48942944899335,-10.406261682256627,274.93144084343373,282.43529770795993,212.80074917469233,279.20241843657726,244.91466941989506,252.20357361411155,-237.9615596134741,234.98700945290122,260.6969308038177,294.06460681492933,124.66376196672707,210.28279547953605,245.06222850169019,204.41539813983093,242.1894833993472,9.977164695826161,275.7981949150983,38.24354004014532,223.44966258547666,259.3745610775684,179.7555268232456,307.77996557219683,280.57928010696173,98.5683457858491,255.8629159235172,-10.357443289066367,223.51698694462755,180.34908445852733,-16.264808111189765,249.44277914241354,10.327264851808565,292.48756981570705,67.81360256915603,238.21570371096436,285.6914900659341,110.26448353414692,266.60226416455106,258.55756096567967,203.3379088564289,59.582792447507984,209.79579864617295,286.41305914707164,197.7122176352411,209.09541014302343,246.06861914335934,235.29956587334797,249.02774639665498,222.45922726612918,273.26950043806556,297.55523800626673,267.9166535426658,153.58313080503467,197.60613006810476,285.75561003451134,237.84851827313148,246.64989219914938,251.42151143060497,247.93516521284832,171.37592189019554,264.77174309025395]}]}
+{"event":"result","variant":"padded","trial":2,"training_seconds":285.93468633299926,"warm_rollout_seconds":2.2986104589654133,"warm_update_seconds":10.917446707957424,"action_std":[0.3272717297077179,0.3935219943523407],"evaluations":[{"deterministic":false,"mean_return":177.65360059121434,"evaluation_seconds":17.95727033400908,"episode_returns":[238.02811623774426,254.44252867803814,290.993720930859,118.1142599925426,-10.011358049457286,137.10979109143818,71.52377864598019,147.00489355917475,272.64569625708054,148.65779180573486,156.7230687240771,160.11468340208089,218.74956305897007,205.6953570969976,110.94026437119443,288.3052052316488,113.6291392194332,272.3198016190979,277.37712630079596,140.2590377548422,211.03793661751743,294.6030967831282,159.35364553040708,91.60994192855973,240.88370459148274,281.7813474663927,164.09832999394132,157.431335376186,99.06635860848058,270.5627078641311,282.6405598121877,145.87950528994884,154.0988506123301,171.1828674095496,48.962332757377624,130.9471321485787,128.26947142572138,176.88975219178974,-54.80168331253724,233.32328343958267,90.28064323715444,274.1202404714016,151.05197396845404,72.68346327371665,265.4290872277129,159.95000136773115,115.5295302615297,263.9259309838652,278.2463780352097,298.4503517956605,99.68759108450716,170.10093432280618,282.9813710210318,100.72823675988491,288.74248741490544,120.39087380232773,248.6722840522633,273.508475347799,260.1364989764617,242.13948930409822,277.01423250806045,90.61015129053375,-127.1403737992002,132.2898011476761,259.9420349260626,169.96716538290946,297.1514302997484,119.1679268238808,150.6675218096853,260.5977387370089,132.04506584051546,173.9484805658499,142.01934786500937,136.88554289873213,228.8302175668057,192.24673436783115,263.330512395276,118.14483665925788,88.82760756413698,-129.78179555076503,246.23258511255048,113.74093776774997,254.57162505343055,203.97055837015583,109.11475775294103,204.2312914683799,166.07058620911607,255.98311782587223,135.74552364028455,291.0934489754007,249.71242388958345,114.18879737453342,149.0627390077175,269.04713611123884,226.35792021562838,37.35963070294246,185.67986426355185,121.91610965882624,272.2384476605404,125.07959359042326]},{"deterministic":true,"mean_return":156.195449765777,"evaluation_seconds":5.846900165954139,"episode_returns":[247.18733365629595,208.592246641428,207.68313521419765,135.74420322371083,264.6846083016675,-63.059248768333774,238.48316372645158,202.53415811448167,-24.348306108038486,169.73414540463745,43.37564954152276,-186.5131817888509,25.673710449205643,242.50806550999013,306.63942894392824,-63.23196324110689,276.68341013606164,-110.2582042989405,112.78219969205914,183.1918960405404,228.82473582963718,234.6216112971865,109.87214247762387,6.219558202781286,231.36216925063115,138.8739560590397,273.03091935738036,264.05746297987787,-59.36446124992322,279.42646685977525,169.26320211112835,256.18492491232587,-13.910687423444344,14.472589395525588,-311.02460308169265,192.1134388144414,256.1987365042543,271.580846509501,160.83991951299154,177.37689980539722,225.85263747446666,227.20416699592784,183.9514090439608,-79.37872115335368,-64.52923756551644,196.355333792494,29.263829336525447,7.134354598361284,182.84055417838994,192.5433341150897,183.65621076058375,163.41445326707597,95.85232824765757,246.65035607577758,228.05814894684337,238.79638846944204,225.35910434298188,263.528788976952,-39.172954889533926,205.60233578981706,254.19146037943273,281.7645418019709,192.5405784370601,261.18566753185144,194.08726428698066,237.66763462692307,251.175016187629,144.46404599757645,200.61494883186367,172.05324993423926,23.32088639343695,-21.765839910455014,292.47309521413695,253.13331180851003,204.5292074181018,151.6401814836903,119.92001888147976,279.45421257554443,295.69154118913053,284.8379258156058,178.70728147612073,-184.41045665669532,217.15412650450077,78.7681754414423,196.1481935594051,112.45535248919613,174.86718000175023,217.25567086704592,274.09879087540276,230.80450016205265,109.60676742515722,253.91952299187173,222.22223475368511,251.09447587308355,230.82852128310205,231.5028123554544,84.45234359548802,265.6960452800492,-77.45208245691644,267.75747255250735]}]}
+{"event":"result","variant":"unpadded","trial":2,"training_seconds":315.26719466701616,"warm_rollout_seconds":2.291924959048629,"warm_update_seconds":10.708665291022044,"action_std":[0.34436631202697754,0.3661150634288788],"evaluations":[{"deterministic":false,"mean_return":140.35923858111204,"evaluation_seconds":7.608739333983976,"episode_returns":[49.163896212709744,244.7570383407441,289.8278376682652,-172.81368894663308,260.42064488709894,155.0629258814065,195.28898868296577,-236.9348662005794,192.70458241296853,125.74505153791743,254.22298993882822,39.543714605075195,-321.98758475004945,276.84838797627435,278.6520271523449,-327.59681673096543,150.9774977212928,196.8695511251004,170.82902357800916,274.2172322387131,-207.27957031527188,-332.54617198422454,292.31893127876435,119.50169670190407,296.7313090405485,233.79261108769882,253.870463662763,234.5629291935849,274.24434269884,161.5989748543225,287.3700023825976,138.08145964120027,155.66694801776146,112.21038087830298,259.1573754169377,-458.2922620938234,246.65236635115585,280.68774430269445,140.91142936930908,142.68708147421503,276.1503603050985,244.85157857599356,156.6846179390582,239.39231935075784,254.45070736177124,242.91842235503336,114.88360637608395,265.5620163211979,-128.21224945745954,-128.434872101101,87.98934570380638,-288.14737247969435,104.78699462626933,249.9536922579371,235.41621041961173,267.9269485308845,145.09147406958868,-259.05615336724657,266.4692896353575,159.40072009008438,146.2193214790688,266.709102525499,145.56420179061863,159.4397563287525,166.62467952317058,165.55970119815188,194.82630442441473,241.4401474300707,-305.94871547875545,258.66187318926603,164.8603439841439,150.74350247917556,48.638936038573746,164.1132532543796,254.9557904382674,270.6628185884957,47.41526292332111,231.163714691463,275.5821582416867,274.67541465662714,243.07187556041524,-60.91706213797891,285.17274625898006,156.0415009554551,74.58587786848429,266.65912732776405,100.26360989669723,128.16060257289897,151.0520172204392,255.3337612843528,239.94494285709231,296.2734918810148,249.69917683613988,160.27312030147885,-356.24334271404535,224.61358325282015,251.0493663283602,256.0379268603998,289.7500263766336,267.389809813611]},{"deterministic":true,"mean_return":-50.42661224618667,"evaluation_seconds":4.130432250036392,"episode_returns":[18.056001713577402,-332.5395143331592,-42.21278705746653,196.30696717068528,-136.25911887177452,-177.25221074705965,-34.144491755751545,-404.70716941113864,-275.25869182377875,-82.81107469509104,-35.33520135270945,194.54635326630688,229.5616260236157,196.9509953136202,-358.4129646648589,-244.32485034041846,29.511491153932525,-348.4037112036338,-244.78340669904793,185.70319927556258,-142.38477436869746,234.79698482656838,9.031598625546351,-66.68971114380989,-360.8593915005681,-50.432235651097656,-416.8057586222503,-300.7687628163699,-431.49566194179397,218.1356714041064,202.85036727892577,-170.19112716893883,148.25559588168113,-224.22673272761563,-7.401637120914415,20.427512514707473,-344.59153841617746,237.03129261413375,207.5499270881166,-134.50100380985677,-312.5420171020564,192.17960884595087,-182.8209625861136,-53.20689521582478,164.92039789937874,163.3545847392461,-68.23768855374519,-366.2389223307547,236.15901557175235,-101.72574027960351,-43.11780775741551,-373.29941918425084,257.0123911566659,-185.92372317778558,-185.2365849025739,-262.36047911687217,-305.4324403029591,-320.959394375992,284.2734968775063,273.73401887645184,213.11314506176222,287.3672206137177,-294.9594633703225,-196.34252542595829,167.68928178879412,197.69200583561118,-235.78000244758582,172.94593096649328,251.86542560790028,-37.702778701508244,-537.8762928261815,-36.85726136469903,257.02420693434135,-200.12866083362667,-197.63478349106012,-282.71410476840026,-372.2902898273705,266.3025384813959,-236.77883869843873,170.46132032720755,142.88797286124083,-228.66110897499095,311.75768260301044,196.93041297231645,-98.37798119106043,-84.7039401496894,-456.0707290795589,275.1059191492493,-234.1883770702415,-49.32262692369152,102.91692900986513,270.19254051296787,-252.78420173014925,-282.4293081559099,-18.700836062356842,181.41114512402123,-287.9978256878992,-58.20292140267202,163.16532069210675,236.5611400345906]}]}
+{"event":"result","variant":"unpadded","trial":3,"training_seconds":276.90296391700394,"warm_rollout_seconds":2.203918707964476,"warm_update_seconds":10.565118417027406,"action_std":[0.32887423038482666,0.35615432262420654],"evaluations":[{"deterministic":false,"mean_return":213.91487196771166,"evaluation_seconds":6.289956416003406,"episode_returns":[264.7238893480639,238.2115505489199,271.47439686667104,270.76729582073824,141.4249396636067,-6.759665434775343,297.5949156874836,278.4432691565372,152.61082825379972,42.90110883265771,237.2195646155548,-210.91777588264225,146.80902031524286,13.703676043593134,280.2119513094106,267.25821927067756,157.20179131879016,267.18487712587194,248.74550034597155,267.0616418488378,285.2768817725367,262.0532704771724,226.04675796971424,226.11366904328182,278.71401541146406,278.81392431869915,234.46514565078158,228.67375814599222,251.96673250628777,274.6101725298088,304.2655488431226,303.0323319979064,306.0680812632123,268.9756721622638,283.5954460301258,75.98243565713197,260.7043817369259,273.5295325219098,277.1709131758135,200.1990074122339,291.32188187838153,251.3017904242135,-100.87498641473941,58.075558792436965,52.00912147782002,274.51187285464005,223.79752623610648,210.67353654185908,240.3037898131425,250.23240167567488,-1.9321624833745688,-94.22085497410416,296.52004520518835,241.75890546079393,100.37930435530966,301.2120972085928,235.46446169027325,256.1325234454688,144.97529570692865,286.4058478068753,268.55929419636834,282.6435022604559,57.084888844657584,248.57507996765685,255.72006995770926,81.81023558502343,219.42546108810518,279.1725460158242,152.75090337693288,244.60909733361902,270.3887206048977,300.370629161455,251.34935402216706,260.332787809976,68.17518463731727,297.2712291223319,262.448628879198,238.89957432432655,222.82120046521982,176.53265943128935,264.8956598560327,253.49646220155665,143.56311232329122,277.8600220524968,267.17218418509947,241.24036240159117,261.88410501535475,263.76251377086345,282.66223529544914,252.01055042770264,155.66266731739262,275.8133314997292,240.74596129867862,248.39451432973325,224.29054049948616,78.12247984821279,232.9695683531001,243.49518350089173,201.1419336288777,299.18613149821067]},{"deterministic":true,"mean_return":86.94532477681285,"evaluation_seconds":3.8000413750414737,"episode_returns":[-63.1711425736914,214.66099103061427,-144.40776499539822,204.60129411468336,282.4418280371915,295.80324507071964,-239.8465054498658,219.56601123912174,277.24295998393023,-105.98277193627031,146.32050045436216,216.95898848336833,196.07959847620276,162.1404976046349,290.05459542201936,196.60634837679422,213.4441591908494,-24.28743195946143,-241.06386966878125,8.992083019142683,154.66176002213652,178.41054646673757,-262.66717224627985,254.5782578874513,214.7172458838357,283.4667983909241,237.8673148193027,285.93335449299695,186.15582530822502,282.05287767173763,8.80689041819734,190.42033918627413,230.24410673131797,229.5790004133884,-240.76061918327105,-177.46285887510123,262.43463677857875,226.38987682926347,186.7353176927346,-15.678579814638795,207.3925881912519,-31.26654851217728,183.17869348738338,198.85544597294964,228.6674026759536,252.9790142551533,277.61615626453977,-340.1216751649053,36.35939115964675,252.75222764219123,-210.62317409741797,-194.98526499808912,-24.916136368940343,-81.65182944677835,223.97544900805892,-290.5518218414191,-138.13090595635944,-134.26746223683293,-26.43899502034789,228.48682280786178,229.34729960880142,-97.17222750942148,-27.25942872547921,-357.9964754037846,243.23000437219335,-127.77269849664978,-36.07316026760722,278.06138859162536,189.62597643142735,271.8890073676895,293.63922268031484,273.8510108527378,245.92178295362777,208.09163709430868,-100.65607161019446,-321.60294037870176,225.61838828615643,219.92939035943812,280.4085098738958,-140.1969544235032,5.579733888643588,2.365309019693555,-251.17680369261672,-161.2933381192879,195.01023010382767,219.7725377062962,189.2567754447725,168.520273450098,227.70881250862067,266.8032197796158,-40.383901201192515,-20.571849648129387,152.30267179042175,159.01238840488696,253.06974822378638,15.561497039197135,99.91556153799574,209.27472636067205,-172.44830254940825,-13.978384639183105]}]}
+{"event":"result","variant":"padded","trial":3,"training_seconds":270.03262262500357,"warm_rollout_seconds":2.140024667023681,"warm_update_seconds":10.27281816600589,"action_std":[0.33396226167678833,0.32078221440315247],"evaluations":[{"deterministic":false,"mean_return":179.45960097269645,"evaluation_seconds":4.86975579202408,"episode_returns":[-118.11456773017545,246.1054745968426,264.9358672111564,64.12927449889864,229.4998045666191,287.00557432208075,-172.96698179963877,238.92036839498093,257.331700138339,270.19081680556167,303.9540551021845,234.39106321790268,30.080025610379316,61.269089949349905,-245.7382055737917,281.8743716773191,59.42335339662182,263.22751376088036,275.4722533779137,304.0307676869045,243.80770174402593,303.49675380860583,317.0826125377796,240.06579106840556,51.772001234463744,261.8911909882079,291.35836024400373,16.109090254096255,57.71653089574977,285.62869338600564,242.49374878851438,278.2289099839926,267.2745009138299,255.23597273856885,-326.23216454572366,228.96385346051466,271.92315873348133,45.06570837059644,310.61059390461537,240.0385164584546,280.5199756284766,233.4636573364511,293.2763309039013,-263.7987808286979,275.64139922395407,-132.03515152354254,-280.94495925554486,298.73755937732426,252.40748549049596,269.0620868414263,27.177874080127538,313.0019736222263,262.1890877067102,281.94305004755563,253.82443840039957,120.2179840599,239.33158611862487,14.820230713947844,32.75592514302994,233.73943248342732,284.1460828170958,275.4713030387658,-102.65334301595814,14.926641179912744,241.10802062803563,58.98743924493252,252.36816713271213,287.34244924342806,298.62487567454264,255.54278444067512,307.9249450160218,-183.68056600591933,290.56955362660403,258.04782202829796,269.71305208094867,-159.39252096334278,242.18413608192245,68.44110474474022,262.5277856938795,249.3308994316934,298.10333290499517,294.28197888013824,295.6740558884866,-224.95783147043454,242.85537782084273,258.43741169911493,283.07287388027964,56.02153676742763,252.73875855200583,300.7177327525815,282.1152601724966,72.21367422232075,260.54930731022745,262.502151617158,213.20495626463295,216.392258758922,261.51237713541235,295.4310337599301,271.25856024156934,259.4203282438372]},{"deterministic":true,"mean_return":1.753567138012233,"evaluation_seconds":3.346837874967605,"episode_returns":[-87.74093374500694,210.27931924332043,244.30974758202694,-275.28102796441857,-323.9091033650674,-284.17471941882997,182.39788605909627,-128.1287506115657,-343.63113262486206,-86.50031704507728,-92.24156787504894,-92.54661261283205,-70.73698105359358,234.9110310009836,213.33120541566302,251.95592900666298,248.68838161816052,-86.07349608982437,-115.00426326051995,229.6970970227546,-297.9118967733441,-331.0409850194543,-150.68466540435523,-430.61605850134583,228.1512125413847,225.29202283801965,-314.4990193547063,-269.682020424557,-105.74264951433939,206.7858650659574,242.47090166094685,241.22624358035748,220.43106730181864,262.5719284276547,222.74199746298376,221.8530651298715,-117.3507343561551,295.2473619678986,-160.63306288950167,244.8924235481182,-68.13859685102751,-307.93729555857897,-93.1149668929982,187.44834228013553,211.5707079219465,-108.58182406398205,220.37577272281828,-318.7416982596958,-124.78498603078492,248.00512055030887,242.95815045918874,199.9140829740999,214.1270595174207,-113.01566424882503,216.03674958290003,-275.422389739779,-155.0525934916269,-98.90369112820356,-216.8013885079865,-159.86456296335,-113.82987150396768,207.91821683085806,-248.21046883392984,254.15813283867632,-94.85631675606854,-103.06415272422933,-84.04054961873126,-0.5860762675011557,199.57719249910963,212.15696303702646,217.73997118704747,268.11664400752954,-88.02910002355833,-100.71591650798715,-96.14994131468335,-246.91552106357545,-125.0204688448493,278.67812336090884,-212.97493131418437,-93.88933882828456,-379.93290509535564,192.25556152093503,-259.64689128463397,236.74968392852267,-15.85970538109001,-136.6473859075787,-86.86073511508346,271.0379425636085,-64.02010338682697,-364.6459248842471,-47.13686589353394,-321.832057372445,-98.60865720911539,-120.43767845049821,201.94449860266553,216.29508980097543,238.3023128936397,237.88844443509507,200.53083952495632,282.7376435103732]}]}

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

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

tinybaselines-0.1.0a4/tests/test_cartpole.py

@@ -0,0 +1,48 @@
+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:
+            observations, _ = env.reset()
+
+            def check_inference():
+                for states in (observations, observations + 1, observations, observations[:1], observations):
+                    expected = policy.actor(policy.preprocess(states)).numpy()
+                    np.testing.assert_allclose(policy._actor_logits(states), expected, rtol=1e-5, atol=1e-6)
+
+            check_inference()
+            actions, _ = policy.act(observations)
+            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))
+            check_inference()
+            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.0a4}/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.0a4}/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.0a4/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.0a4}/tests/test_parity.py

@@ -5,10 +5,23 @@ 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
+from tinybaselines.policies import MLP, normal_log_prob
 
 
 class ParityTest(unittest.TestCase):
+    def test_cpu_mlp_padding_preserves_outputs_and_gradients(self):
+        model = MLP(8, 2, 16, 1)
+        states = lander.Tensor(np.random.default_rng(0).standard_normal((10_000, 8), dtype=np.float32), device='CPU')
+        params = lander.nn.state.get_parameters(model)
+        with lander.Context(TRAINING=1):
+            with patch('tinybaselines.policies._CPU_PADDING', False):
+                expected = model(states)
+            with patch('tinybaselines.policies._CPU_PADDING', True):
+                actual = model(states)
+            np.testing.assert_allclose(actual.numpy(), expected.numpy(), rtol=1e-5, atol=1e-6)
+            for a, b in zip(actual.square().mean().gradient(*params), expected.square().mean().gradient(*params)):
+                np.testing.assert_allclose(a.numpy(), b.numpy(), rtol=1e-5, atol=1e-6)
+
     def test_rejects_nonfinite_policy_std(self):
         model = lander.ActorCritic(8, 2, 8, 1, 1)
         model.log_std.assign(lander.Tensor([float('nan'), 0])).realize()

{tinybaselines-0.1.0a2 → tinybaselines-0.1.0a4}/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.0a4/tinybaselines/policies.py

@@ -0,0 +1,183 @@
+import json
+from os import getenv
+
+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))
+_CPU_PADDING = getenv('TINYBASELINES_CPU_PADDING') == '1'
+
+
+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):
+        self.layers = [nn.Linear(input_dim, hidden_dim)]
+        self.layers.extend(nn.Linear(hidden_dim, hidden_dim) for _ in range(layer_count - 1))
+        self.layers.append(nn.Linear(hidden_dim, output_dim))
+
+    def __call__(self, x):
+        rows = x.shape[0] if _CPU_PADDING and x.device == 'CPU' and x.ndim == 2 and x.shape[0] >= 8192 else 0
+        padding = -rows % 256
+        if padding:
+            # Align large CPU dense kernels, then discard the padded rows.
+            x = x.pad(((0, padding), (0, 0))).contiguous()
+        for layer in self.layers[:-1]:
+            x = layer(x).relu()
+        x = self.layers[-1](x)
+        return x[:rows] if padding else x
+
+
+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)
+        self.actor = MLP(state_dim, action_dim, hidden_dim, actor_layers)
+        self.critic = MLP(state_dim, 1, hidden_dim, critic_layers)
+        self.actor.layers[-1].weight.assign(self.actor.layers[-1].weight * 0.01).realize()
+        self.actor.layers[-1].bias.assign(Tensor.zeros_like(self.actor.layers[-1].bias)).realize()
+        self.log_std = Tensor.full((action_dim,), -0.7).realize()
+        scale = np.ones(state_dim, np.float32) if observation_scale is None else np.asarray(observation_scale, np.float32)
+        self.observation_scale = Tensor(scale.copy()).is_param_(False).realize()
+        self.jit_shape = None
+        self._cpu_input = None
+        self.actor_jit = TinyJit(lambda state, scale=None: self.actor(state if scale is None else state * scale).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
+
+    def act(self, state, deterministic=False):
+        preprocess = self.preprocess
+        buffered = self.device == 'CPU' and getattr(preprocess, '__func__', None) is GaussianActorCritic.preprocess
+        if not buffered:
+            state_tensor = preprocess(state)
+        signature = (state.shape if buffered else state_tensor.shape), buffered
+        if self.jit_shape != signature:
+            self.actor_jit.reset()
+            self.jit_shape = signature
+            self._cpu_input = None
+        if buffered:
+            if self._cpu_input is None:
+                # Keep the NumPy owner alive for the lifetime of the tensor.
+                host = np.empty(state.shape, dtype=np.float32)
+                self._cpu_input = host, Tensor.from_blob(host.ctypes.data, host.shape, device='CPU', dtype='float32')
+            host, state_tensor = self._cpu_input
+            host[:] = state
+            mean = self.actor_jit(state_tensor, self.observation_scale).numpy()
+        else:
+            mean = self.actor_jit(state_tensor).numpy()
+        if deterministic:
+            action = np.clip(mean, -1, 1)
+            return action, mean
+        sample = np.random.normal(mean, self.action_std).astype(np.float32)
+        action = np.clip(sample, -1, 1)
+        return action, sample
+
+    def state_dict(self):
+        return nn.state.get_state_dict(dict(actor=self.actor, critic=self.critic, log_std=self.log_std,
+                                            observation_scale=self.observation_scale))
+
+    def save(self, path):
+        nn.state.safe_save(self.state_dict(), path, metadata={'config': json.dumps(self.config)})
+
+    @classmethod
+    def load(cls, path):
+        _, _, header = nn.state.safe_load_metadata(path)
+        model = cls(**json.loads(header['__metadata__']['config']))
+        nn.state.load_state_dict(model.state_dict(), nn.state.safe_load(path), verbose=False)
+        model.refresh_inference()
+        return model
+
+    def refresh_inference(self):
+        self.actor_jit.reset()
+        self.action_std = np.exp(self.log_std.clip(-5, 2).numpy())
+        if not np.isfinite(self.action_std).all():
+            raise FloatingPointError('non-finite policy standard deviation')
+
+
+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
+        self._cpu_input = 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 _actor_logits(self, state):
+        if self.jit_shape != state.shape:
+            self.refresh_inference()
+            self.jit_shape = state.shape
+            self._cpu_input = None
+        if self.device == 'CPU':
+            if self._cpu_input is None:
+                # Keep the NumPy owner alive for the lifetime of the tensor.
+                host = np.empty(state.shape, dtype=np.float32)
+                self._cpu_input = host, Tensor.from_blob(host.ctypes.data, host.shape, device='CPU', dtype='float32')
+            host, states = self._cpu_input
+            host[:] = state
+        else:
+            states = self.preprocess(state)
+        return self.actor_jit(states).numpy()
+
+    def act(self, state, deterministic=False):
+        logits = self._actor_logits(state)
+        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.0a4}/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.0a4}/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.0a4/tinybaselines.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tinybaselines
-Version: 0.1.0a2
+Version: 0.1.0a4
 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.0a4
 ```
 
 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,24 @@ pip install -e . -r examples/lunar_lander/requirements.txt
 python -m examples.lunar_lander.train
 ```
 
+CPU batch padding is off by default. Set `TINYBASELINES_CPU_PADDING=1` before
+starting Python to try it. In three runs per variant at 420,000 transitions,
+mean training time was 296.5 s without padding and 282.1 s with it. Sampled
+evaluation cleared 200 in two unpadded runs and no padded runs.
+[Settings and both evaluation modes](examples/lunar_lander/results/cpu-padding.jsonl).
+
+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.0a4}/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
@@ -11,6 +13,7 @@ examples/lunar_lander/requirements-reference.txt
 examples/lunar_lander/requirements.txt
 examples/lunar_lander/run_learning.py
 examples/lunar_lander/train.py
+examples/lunar_lander/results/cpu-padding.jsonl
 examples/lunar_lander/results/2026-10-26/phase-initial-tinygrad-thread4.jsonl
 examples/lunar_lander/results/2026-10-26/phase-initial-torch.jsonl
 examples/lunar_lander/results/2026-10-26/phase-protocol.json
@@ -27,6 +30,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))

tinybaselines-0.1.0a2/tinybaselines/policies.py

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