tinybaselines 0.1.0a3__tar.gz → 0.1.0a4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tinybaselines
-Version: 0.1.0a3
+Version: 0.1.0a4
 Summary: PPO in tinygrad.
 License-Expression: MIT
 Requires-Python: >=3.11
@@ -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.0a3
+python -m pip install tinybaselines==0.1.0a4
 ```
 
 Install Gymnasium and Box2D to run this example:
@@ -86,6 +86,12 @@ 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.
 

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

@@ -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.0a3
+python -m pip install tinybaselines==0.1.0a4
 ```
 
 Install Gymnasium and Box2D to run this example:
@@ -74,6 +74,12 @@ 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.
 

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.0a3 → tinybaselines-0.1.0a4}/pyproject.toml

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

{tinybaselines-0.1.0a3 → tinybaselines-0.1.0a4}/tests/test_cartpole.py

@@ -18,11 +18,20 @@ class CartPoleTest(unittest.TestCase):
         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])
+            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'

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

@@ -1,10 +1,12 @@
 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):
@@ -19,9 +21,15 @@ class MLP:
         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()
-        return self.layers[-1](x)
+        x = self.layers[-1](x)
+        return x[:rows] if padding else x
 
 
 class GaussianActorCritic:
@@ -38,7 +46,8 @@ class GaussianActorCritic:
         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._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
@@ -66,11 +75,25 @@ class GaussianActorCritic:
         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:
+        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 = state_tensor.shape
-        mean = self.actor_jit(state_tensor).numpy()
+            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
@@ -112,6 +135,7 @@ class CategoricalActorCritic:
         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):
@@ -130,12 +154,24 @@ class CategoricalActorCritic:
         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:
+    def _actor_logits(self, state):
+        if self.jit_shape != state.shape:
             self.refresh_inference()
-            self.jit_shape = states.shape
-        logits = self.actor_jit(states).numpy()
+            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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tinybaselines
-Version: 0.1.0a3
+Version: 0.1.0a4
 Summary: PPO in tinygrad.
 License-Expression: MIT
 Requires-Python: >=3.11
@@ -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.0a3
+python -m pip install tinybaselines==0.1.0a4
 ```
 
 Install Gymnasium and Box2D to run this example:
@@ -86,6 +86,12 @@ 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.
 

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

@@ -13,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