Initial implementation: all models, methods, toy and CIFAR experiments

Debug phase. Toy LQ experiments (3 seeds) complete with terminal gradient matching.
Credit bridge matches state bridge on linear system (~0.94 cosine).
CIFAR experiments in progress.

9 files changed, 3135 insertions, 0 deletions

diff --git a/results/toy_lq/state_bridge_seed42.pt b/results/toy_lq/state_bridge_seed42.pt
new file mode 100644
index 0000000..a87e99d
--- /dev/null
+++ b/results/toy_lq/state_bridge_seed42.pt
diff --git a/results/toy_lq/sweep_results.json b/results/toy_lq/sweep_results.json
new file mode 100644
index 0000000..30c3f49
--- /dev/null
+++ b/results/toy_lq/sweep_results.json
@@ -0,0 +1,1070 @@
+{
+  "base": {
+    "best_cos": 0.28987203165888786,
+    "best_step": 500,
+    "final_cos": -0.0006980087685709199,
+    "final_rho": 0.00831739305673788,
+    "final_nudge": 0.0027815375166634717,
+    "history": [
+      {
+        "step": 500,
+        "avg_cos": 0.28987203165888786,
+        "avg_rho": 0.30755721777677536,
+        "avg_nudge": -0.10750458451608817,
+        "loss_term": 0.4910352826118469,
+        "loss_bridge": 0.21142145991325378
+      },
+      {
+        "step": 1000,
+        "avg_cos": 0.17731603110829988,
+        "avg_rho": 0.17737891773382822,
+        "avg_nudge": -0.06381315520654122,
+        "loss_term": 0.1414657086133957,
+        "loss_bridge": 0.16202926635742188
+      },
+      {
+        "step": 1500,
+        "avg_cos": 0.07300025007377069,
+        "avg_rho": 0.07345063022027414,
+        "avg_nudge": -0.02641519942941765,
+        "loss_term": 0.07946588099002838,
+        "loss_bridge": 0.08628662675619125
+      },
+      {
+        "step": 2000,
+        "avg_cos": 0.05074742032835881,
+        "avg_rho": 0.048574333622430764,
+        "avg_nudge": -0.016719758200148743,
+        "loss_term": 0.06007867306470871,
+        "loss_bridge": 0.016275471076369286
+      },
+      {
+        "step": 2500,
+        "avg_cos": 0.045629044994711876,
+        "avg_rho": 0.0567939051737388,
+        "avg_nudge": -0.016934571011612814,
+        "loss_term": 0.04579576849937439,
+        "loss_bridge": 0.008496936410665512
+      },
+      {
+        "step": 3000,
+        "avg_cos": 0.06145744491368532,
+        "avg_rho": 0.05695042138298353,
+        "avg_nudge": -0.020737762407710154,
+        "loss_term": 0.04881729558110237,
+        "loss_bridge": 0.013217507861554623
+      },
+      {
+        "step": 3500,
+        "avg_cos": 0.04631757119204849,
+        "avg_rho": 0.0402289762472113,
+        "avg_nudge": -0.015406294725835323,
+        "loss_term": 0.04255056008696556,
+        "loss_bridge": 0.01808304153382778
+      },
+      {
+        "step": 4000,
+        "avg_cos": 0.07397006265819073,
+        "avg_rho": 0.07162392682706316,
+        "avg_nudge": -0.025235373992472887,
+        "loss_term": 0.05486675351858139,
+        "loss_bridge": 0.02182137593626976
+      },
+      {
+        "step": 4500,
+        "avg_cos": 0.07251314166933298,
+        "avg_rho": 0.07301904633641243,
+        "avg_nudge": -0.027378834628810484,
+        "loss_term": 0.027670202776789665,
+        "loss_bridge": 0.013123266398906708
+      },
+      {
+        "step": 5000,
+        "avg_cos": -0.0006980087685709199,
+        "avg_rho": 0.00831739305673788,
+        "avg_nudge": 0.0027815375166634717,
+        "loss_term": 0.021854262799024582,
+        "loss_bridge": 0.0122066093608737
+      }
+    ]
+  },
+  "noise_0.1": {
+    "best_cos": 0.28885648772120476,
+    "best_step": 500,
+    "final_cos": 0.02330129671220978,
+    "final_rho": 0.03084004654859503,
+    "final_nudge": -0.005932313855737448,
+    "history": [
+      {
+        "step": 500,
+        "avg_cos": 0.28885648772120476,
+        "avg_rho": 0.306567445397377,
+        "avg_nudge": -0.1071426725635926,
+        "loss_term": 0.4925232529640198,
+        "loss_bridge": 0.20981840789318085
+      },
+      {
+        "step": 1000,
+        "avg_cos": 0.17456108083327612,
+        "avg_rho": 0.17607956007122993,
+        "avg_nudge": -0.06276426805804174,
+        "loss_term": 0.1324017345905304,
+        "loss_bridge": 0.1835721731185913
+      },
+      {
+        "step": 1500,
+        "avg_cos": 0.08755840392162402,
+        "avg_rho": 0.08556619860852759,
+        "avg_nudge": -0.03154423305143913,
+        "loss_term": 0.07669232785701752,
+        "loss_bridge": 0.10240059345960617
+      },
+      {
+        "step": 2000,
+        "avg_cos": 0.05209386876473824,
+        "avg_rho": 0.04985952338514229,
+        "avg_nudge": -0.01653688432027896,
+        "loss_term": 0.06218753010034561,
+        "loss_bridge": 0.02190624177455902
+      },
+      {
+        "step": 2500,
+        "avg_cos": 0.038306045811623335,
+        "avg_rho": 0.047700356847296156,
+        "avg_nudge": -0.01320391776971519,
+        "loss_term": 0.06128765642642975,
+        "loss_bridge": 0.01958899199962616
+      },
+      {
+        "step": 3000,
+        "avg_cos": 0.059713449950019516,
+        "avg_rho": 0.054914615117013454,
+        "avg_nudge": -0.01939693869402011,
+        "loss_term": 0.04238733649253845,
+        "loss_bridge": 0.006958359386771917
+      },
+      {
+        "step": 3500,
+        "avg_cos": 0.04283027160757532,
+        "avg_rho": 0.028526597811530035,
+        "avg_nudge": -0.01381517636279265,
+        "loss_term": 0.024001002311706543,
+        "loss_bridge": 0.006152431480586529
+      },
+      {
+        "step": 4000,
+        "avg_cos": 0.0656419579560558,
+        "avg_rho": 0.06320031352030735,
+        "avg_nudge": -0.02216823499960204,
+        "loss_term": 0.04236245155334473,
+        "loss_bridge": 0.01628262549638748
+      },
+      {
+        "step": 4500,
+        "avg_cos": 0.05787398883452018,
+        "avg_rho": 0.0610881638713181,
+        "avg_nudge": -0.02213552314788103,
+        "loss_term": 0.024835661053657532,
+        "loss_bridge": 0.008629711344838142
+      },
+      {
+        "step": 5000,
+        "avg_cos": 0.02330129671220978,
+        "avg_rho": 0.03084004654859503,
+        "avg_nudge": -0.005932313855737448,
+        "loss_term": 0.035377949476242065,
+        "loss_bridge": 0.020307490602135658
+      }
+    ]
+  },
+  "noise_0.3": {
+    "best_cos": 0.28320933257540065,
+    "best_step": 500,
+    "final_cos": -0.007130145425132166,
+    "final_rho": 0.003303757131410142,
+    "final_nudge": 0.004990910800794761,
+    "history": [
+      {
+        "step": 500,
+        "avg_cos": 0.28320933257540065,
+        "avg_rho": 0.30105741570393246,
+        "avg_nudge": -0.10515742873152097,
+        "loss_term": 0.44116002321243286,
+        "loss_bridge": 0.210750013589859
+      },
+      {
+        "step": 1000,
+        "avg_cos": 0.1680816188454628,
+        "avg_rho": 0.1721916707853476,
+        "avg_nudge": -0.060552582144737244,
+        "loss_term": 0.14207889139652252,
+        "loss_bridge": 0.16882070899009705
+      },
+      {
+        "step": 1500,
+        "avg_cos": 0.08877109829336405,
+        "avg_rho": 0.08808179199695587,
+        "avg_nudge": -0.03221472934819758,
+        "loss_term": 0.12282441556453705,
+        "loss_bridge": 0.07336755841970444
+      },
+      {
+        "step": 2000,
+        "avg_cos": 0.04903770827998718,
+        "avg_rho": 0.04795513402981063,
+        "avg_nudge": -0.016069856472313404,
+        "loss_term": 0.06662434339523315,
+        "loss_bridge": 0.02520526573061943
+      },
+      {
+        "step": 2500,
+        "avg_cos": 0.05001032492145896,
+        "avg_rho": 0.053442057532568775,
+        "avg_nudge": -0.01783018947268526,
+        "loss_term": 0.04614880681037903,
+        "loss_bridge": 0.015394347719848156
+      },
+      {
+        "step": 3000,
+        "avg_cos": 0.03979540191357955,
+        "avg_rho": 0.03453826089389622,
+        "avg_nudge": -0.012916556559503078,
+        "loss_term": 0.04155049845576286,
+        "loss_bridge": 0.007860729470849037
+      },
+      {
+        "step": 3500,
+        "avg_cos": 0.06703585013747215,
+        "avg_rho": 0.04949762811884284,
+        "avg_nudge": -0.02297227829694748,
+        "loss_term": 0.059078700840473175,
+        "loss_bridge": 0.030750930309295654
+      },
+      {
+        "step": 4000,
+        "avg_cos": 0.06641693723698457,
+        "avg_rho": 0.06286459214364488,
+        "avg_nudge": -0.021883966866880655,
+        "loss_term": 0.025344880297780037,
+        "loss_bridge": 0.008583566173911095
+      },
+      {
+        "step": 4500,
+        "avg_cos": 0.03687915485352278,
+        "avg_rho": 0.028895257118468482,
+        "avg_nudge": -0.01225966370354096,
+        "loss_term": 0.020479349419474602,
+        "loss_bridge": 0.00925756897777319
+      },
+      {
+        "step": 5000,
+        "avg_cos": -0.007130145425132166,
+        "avg_rho": 0.003303757131410142,
+        "avg_nudge": 0.004990910800794761,
+        "loss_term": 0.03882830590009689,
+        "loss_bridge": 0.022515198215842247
+      }
+    ]
+  },
+  "lam_1.0": {
+    "best_cos": 0.2899630069732666,
+    "best_step": 500,
+    "final_cos": 0.0024576462844076255,
+    "final_rho": 0.010072723651925722,
+    "final_nudge": 0.0017782459035515785,
+    "history": [
+      {
+        "step": 500,
+        "avg_cos": 0.2899630069732666,
+        "avg_rho": 0.3076498980323474,
+        "avg_nudge": -0.10753695170084636,
+        "loss_term": 0.49082812666893005,
+        "loss_bridge": 0.21152979135513306
+      },
+      {
+        "step": 1000,
+        "avg_cos": 0.17747302974263826,
+        "avg_rho": 0.17881879458824793,
+        "avg_nudge": -0.06420790683478117,
+        "loss_term": 0.1407778263092041,
+        "loss_bridge": 0.15578678250312805
+      },
+      {
+        "step": 1500,
+        "avg_cos": 0.08065215001503627,
+        "avg_rho": 0.08029377926141024,
+        "avg_nudge": -0.029326035796354216,
+        "loss_term": 0.07120706140995026,
+        "loss_bridge": 0.09130121767520905
+      },
+      {
+        "step": 2000,
+        "avg_cos": 0.0563324602941672,
+        "avg_rho": 0.058026942308060825,
+        "avg_nudge": -0.01875296530003349,
+        "loss_term": 0.06783974170684814,
+        "loss_bridge": 0.01712188683450222
+      },
+      {
+        "step": 2500,
+        "avg_cos": 0.039659525423000254,
+        "avg_rho": 0.04613731553157171,
+        "avg_nudge": -0.013568413676694036,
+        "loss_term": 0.0634445995092392,
+        "loss_bridge": 0.02091756835579872
+      },
+      {
+        "step": 3000,
+        "avg_cos": 0.04386034506993989,
+        "avg_rho": 0.039356538482631244,
+        "avg_nudge": -0.01407914562150836,
+        "loss_term": 0.05031818151473999,
+        "loss_bridge": 0.01259439717978239
+      },
+      {
+        "step": 3500,
+        "avg_cos": 0.0481002099889641,
+        "avg_rho": 0.03689269566287597,
+        "avg_nudge": -0.01590493693947792,
+        "loss_term": 0.03421805799007416,
+        "loss_bridge": 0.014973677694797516
+      },
+      {
+        "step": 4000,
+        "avg_cos": 0.07999587121109168,
+        "avg_rho": 0.07683291751891375,
+        "avg_nudge": -0.02731190746029218,
+        "loss_term": 0.04726963862776756,
+        "loss_bridge": 0.017824556678533554
+      },
+      {
+        "step": 4500,
+        "avg_cos": 0.07246351769814889,
+        "avg_rho": 0.07343382605661948,
+        "avg_nudge": -0.027415843835721414,
+        "loss_term": 0.02768528461456299,
+        "loss_bridge": 0.013396943919360638
+      },
+      {
+        "step": 5000,
+        "avg_cos": 0.0024576462844076255,
+        "avg_rho": 0.010072723651925722,
+        "avg_nudge": 0.0017782459035515785,
+        "loss_term": 0.02114972099661827,
+        "loss_bridge": 0.011466547846794128
+      }
+    ]
+  },
+  "noise_lam": {
+    "best_cos": 0.28980905935168266,
+    "best_step": 500,
+    "final_cos": 0.000333610107190907,
+    "final_rho": 0.009229002442831794,
+    "final_nudge": 0.0020260093733668327,
+    "history": [
+      {
+        "step": 500,
+        "avg_cos": 0.28980905935168266,
+        "avg_rho": 0.30755361666282016,
+        "avg_nudge": -0.10748158146937688,
+        "loss_term": 0.4891643524169922,
+        "loss_bridge": 0.21093463897705078
+      },
+      {
+        "step": 1000,
+        "avg_cos": 0.1692807301878929,
+        "avg_rho": 0.17357065031925836,
+        "avg_nudge": -0.060928904761870704,
+        "loss_term": 0.12851648032665253,
+        "loss_bridge": 0.1991802453994751
+      },
+      {
+        "step": 1500,
+        "avg_cos": 0.08976124723752339,
+        "avg_rho": 0.08664474201699097,
+        "avg_nudge": -0.032548267083863415,
+        "loss_term": 0.08289942145347595,
+        "loss_bridge": 0.08361449092626572
+      },
+      {
+        "step": 2000,
+        "avg_cos": 0.0449219069754084,
+        "avg_rho": 0.03701058775186539,
+        "avg_nudge": -0.013820620176071921,
+        "loss_term": 0.07352523505687714,
+        "loss_bridge": 0.031066572293639183
+      },
+      {
+        "step": 2500,
+        "avg_cos": 0.0341803894067804,
+        "avg_rho": 0.04086895197785149,
+        "avg_nudge": -0.011834259377792478,
+        "loss_term": 0.05221429467201233,
+        "loss_bridge": 0.013445280492305756
+      },
+      {
+        "step": 3000,
+        "avg_cos": 0.043432267770792045,
+        "avg_rho": 0.035848827101290226,
+        "avg_nudge": -0.0138394293996195,
+        "loss_term": 0.08377102017402649,
+        "loss_bridge": 0.02437596395611763
+      },
+      {
+        "step": 3500,
+        "avg_cos": 0.04308730812044814,
+        "avg_rho": 0.02904196917855491,
+        "avg_nudge": -0.013877601828426123,
+        "loss_term": 0.028318829834461212,
+        "loss_bridge": 0.0096125528216362
+      },
+      {
+        "step": 4000,
+        "avg_cos": 0.07217423028002183,
+        "avg_rho": 0.07531639956869185,
+        "avg_nudge": -0.024754961564516027,
+        "loss_term": 0.04526568949222565,
+        "loss_bridge": 0.01877186819911003
+      },
+      {
+        "step": 4500,
+        "avg_cos": 0.041415012208744884,
+        "avg_rho": 0.039630077119606234,
+        "avg_nudge": -0.013629865366965532,
+        "loss_term": 0.0305576603859663,
+        "loss_bridge": 0.011333338916301727
+      },
+      {
+        "step": 5000,
+        "avg_cos": 0.000333610107190907,
+        "avg_rho": 0.009229002442831794,
+        "avg_nudge": 0.0020260093733668327,
+        "loss_term": 0.020293015986680984,
+        "loss_bridge": 0.00792770553380251
+      }
+    ]
+  },
+  "no_ln": {
+    "best_cos": 0.2994285201032956,
+    "best_step": 500,
+    "final_cos": -0.027601251068214577,
+    "final_rho": -0.03011056105606258,
+    "final_nudge": 0.013365049380809069,
+    "history": [
+      {
+        "step": 500,
+        "avg_cos": 0.2994285201032956,
+        "avg_rho": 0.3287110353509585,
+        "avg_nudge": -0.11226618165771167,
+        "loss_term": 0.5378487706184387,
+        "loss_bridge": 0.20209679007530212
+      },
+      {
+        "step": 1000,
+        "avg_cos": 0.23427631705999374,
+        "avg_rho": 0.24483238657315573,
+        "avg_nudge": -0.08577251620590687,
+        "loss_term": 0.14587438106536865,
+        "loss_bridge": 0.17536549270153046
+      },
+      {
+        "step": 1500,
+        "avg_cos": 0.10971186744670074,
+        "avg_rho": 0.1081712432205677,
+        "avg_nudge": -0.039541066934665046,
+        "loss_term": 0.09813931584358215,
+        "loss_bridge": 0.1144137978553772
+      },
+      {
+        "step": 2000,
+        "avg_cos": 0.17568999342620373,
+        "avg_rho": 0.18043102137744427,
+        "avg_nudge": -0.06037815675760309,
+        "loss_term": 0.12119113653898239,
+        "loss_bridge": 0.05170433223247528
+      },
+      {
+        "step": 2500,
+        "avg_cos": 0.1501085925847292,
+        "avg_rho": 0.1694059126699964,
+        "avg_nudge": -0.05651436994473139,
+        "loss_term": 0.07073387503623962,
+        "loss_bridge": 0.015867799520492554
+      },
+      {
+        "step": 3000,
+        "avg_cos": 0.10612630782028039,
+        "avg_rho": 0.09910948442605634,
+        "avg_nudge": -0.034841354160259165,
+        "loss_term": 0.05772688612341881,
+        "loss_bridge": 0.0220709890127182
+      },
+      {
+        "step": 3500,
+        "avg_cos": 0.12348712608218193,
+        "avg_rho": 0.12745930068194866,
+        "avg_nudge": -0.04433920063699285,
+        "loss_term": 0.04150111600756645,
+        "loss_bridge": 0.014408521354198456
+      },
+      {
+        "step": 4000,
+        "avg_cos": 0.012287488003494218,
+        "avg_rho": 0.038615713633286454,
+        "avg_nudge": -0.006448045140132308,
+        "loss_term": 0.04848453402519226,
+        "loss_bridge": 0.026737889274954796
+      },
+      {
+        "step": 4500,
+        "avg_cos": 0.0344570055603981,
+        "avg_rho": 0.043244189505154886,
+        "avg_nudge": -0.012965732564528784,
+        "loss_term": 0.04390523582696915,
+        "loss_bridge": 0.016968518495559692
+      },
+      {
+        "step": 5000,
+        "avg_cos": -0.027601251068214577,
+        "avg_rho": -0.03011056105606258,
+        "avg_nudge": 0.013365049380809069,
+        "loss_term": 0.05419892817735672,
+        "loss_bridge": 0.029718749225139618
+      }
+    ]
+  },
+  "big_vnet": {
+    "best_cos": 0.25947993124524754,
+    "best_step": 500,
+    "final_cos": 0.012725223108039549,
+    "final_rho": -0.0029713623225688934,
+    "final_nudge": -0.0007455882926781973,
+    "history": [
+      {
+        "step": 500,
+        "avg_cos": 0.25947993124524754,
+        "avg_rho": 0.2872927797337373,
+        "avg_nudge": -0.09759780826667945,
+        "loss_term": 0.24058431386947632,
+        "loss_bridge": 0.20110812783241272
+      },
+      {
+        "step": 1000,
+        "avg_cos": 0.11903308952848117,
+        "avg_rho": 0.10757205138603847,
+        "avg_nudge": -0.04069022353117665,
+        "loss_term": 0.1535366326570511,
+        "loss_bridge": 0.10731971263885498
+      },
+      {
+        "step": 1500,
+        "avg_cos": 0.04738504672423005,
+        "avg_rho": 0.04187516961246729,
+        "avg_nudge": -0.01667174060518543,
+        "loss_term": 0.12195796519517899,
+        "loss_bridge": 0.03276967257261276
+      },
+      {
+        "step": 2000,
+        "avg_cos": 0.05584627948701382,
+        "avg_rho": 0.06464844088380535,
+        "avg_nudge": -0.019484267104417086,
+        "loss_term": 0.07039390504360199,
+        "loss_bridge": 0.023653611540794373
+      },
+      {
+        "step": 2500,
+        "avg_cos": 0.11334512817362945,
+        "avg_rho": 0.13342145457863808,
+        "avg_nudge": -0.04256325137491027,
+        "loss_term": 0.09321287274360657,
+        "loss_bridge": 0.03603611886501312
+      },
+      {
+        "step": 3000,
+        "avg_cos": 0.07876436489944656,
+        "avg_rho": 0.08206061793801685,
+        "avg_nudge": -0.02565166230003039,
+        "loss_term": 0.03433217480778694,
+        "loss_bridge": 0.014776414260268211
+      },
+      {
+        "step": 3500,
+        "avg_cos": 0.059695989514390625,
+        "avg_rho": 0.043808821588754654,
+        "avg_nudge": -0.01833860871071617,
+        "loss_term": 0.07867280393838882,
+        "loss_bridge": 0.04651641845703125
+      },
+      {
+        "step": 4000,
+        "avg_cos": 0.008286381295571724,
+        "avg_rho": 0.01973852072842419,
+        "avg_nudge": -0.0019318210737158854,
+        "loss_term": 0.03502834588289261,
+        "loss_bridge": 0.011813998222351074
+      },
+      {
+        "step": 4500,
+        "avg_cos": -0.006770744492920737,
+        "avg_rho": 2.3505534045398235e-05,
+        "avg_nudge": 0.002897862965861956,
+        "loss_term": 0.04147114232182503,
+        "loss_bridge": 0.026934277266263962
+      },
+      {
+        "step": 5000,
+        "avg_cos": 0.012725223108039549,
+        "avg_rho": -0.0029713623225688934,
+        "avg_nudge": -0.0007455882926781973,
+        "loss_term": 0.038749027997255325,
+        "loss_bridge": 0.01941092312335968
+      }
+    ]
+  },
+  "ema_0.999": {
+    "best_cos": 0.10180055970946948,
+    "best_step": 1000,
+    "final_cos": -0.01584776126158734,
+    "final_rho": -0.01703926082700491,
+    "final_nudge": 0.007290713644276063,
+    "history": [
+      {
+        "step": 500,
+        "avg_cos": -0.005628384804974,
+        "avg_rho": 0.010925033983464042,
+        "avg_nudge": 0.0015428058492640655,
+        "loss_term": 0.5920301675796509,
+        "loss_bridge": 1.4890536069869995
+      },
+      {
+        "step": 1000,
+        "avg_cos": 0.10180055970946948,
+        "avg_rho": 0.10515290250380833,
+        "avg_nudge": -0.037424925404290356,
+        "loss_term": 0.5715007185935974,
+        "loss_bridge": 0.4665977954864502
+      },
+      {
+        "step": 1500,
+        "avg_cos": 0.021814276037427287,
+        "avg_rho": 0.003451728650058309,
+        "avg_nudge": -0.005690907438596089,
+        "loss_term": 0.2616257071495056,
+        "loss_bridge": 0.2758212387561798
+      },
+      {
+        "step": 2000,
+        "avg_cos": 0.03669632730695108,
+        "avg_rho": 0.036930523036668696,
+        "avg_nudge": -0.01006293793519338,
+        "loss_term": 0.11164076626300812,
+        "loss_bridge": 0.1407940685749054
+      },
+      {
+        "step": 2500,
+        "avg_cos": -0.020325756592986483,
+        "avg_rho": -0.027848235641916592,
+        "avg_nudge": 0.011153894321372112,
+        "loss_term": 0.15471391379833221,
+        "loss_bridge": 0.06181420385837555
+      },
+      {
+        "step": 3000,
+        "avg_cos": -0.0060501456415901584,
+        "avg_rho": -0.013405103546877703,
+        "avg_nudge": 0.005228333951284488,
+        "loss_term": 0.07506504654884338,
+        "loss_bridge": 0.08207326382398605
+      },
+      {
+        "step": 3500,
+        "avg_cos": -0.02149865326161186,
+        "avg_rho": -0.023167532014970977,
+        "avg_nudge": 0.010358475303898254,
+        "loss_term": 0.048137497156858444,
+        "loss_bridge": 0.03276998922228813
+      },
+      {
+        "step": 4000,
+        "avg_cos": -0.007104064881180723,
+        "avg_rho": -0.00720199760204802,
+        "avg_nudge": 0.005406570465614398,
+        "loss_term": 0.03773114085197449,
+        "loss_bridge": 0.037960827350616455
+      },
+      {
+        "step": 4500,
+        "avg_cos": -0.0034141440119128674,
+        "avg_rho": -0.003708663280121982,
+        "avg_nudge": 0.002461720102777084,
+        "loss_term": 0.0416095145046711,
+        "loss_bridge": 0.03200625628232956
+      },
+      {
+        "step": 5000,
+        "avg_cos": -0.01584776126158734,
+        "avg_rho": -0.01703926082700491,
+        "avg_nudge": 0.007290713644276063,
+        "loss_term": 0.023291591554880142,
+        "loss_bridge": 0.026996765285730362
+      }
+    ]
+  },
+  "K16": {
+    "best_cos": 0.3187306026617686,
+    "best_step": 500,
+    "final_cos": 0.012039402422184745,
+    "final_rho": -0.002023946028202772,
+    "final_nudge": -0.0010173787983755271,
+    "history": [
+      {
+        "step": 500,
+        "avg_cos": 0.3187306026617686,
+        "avg_rho": 0.3298306291302045,
+        "avg_nudge": -0.12337777391076088,
+        "loss_term": 0.3945310413837433,
+        "loss_bridge": 0.27078020572662354
+      },
+      {
+        "step": 1000,
+        "avg_cos": 0.15608959334592024,
+        "avg_rho": 0.14268888781468073,
+        "avg_nudge": -0.05713697926451763,
+        "loss_term": 0.13724187016487122,
+        "loss_bridge": 0.13912135362625122
+      },
+      {
+        "step": 1500,
+        "avg_cos": 0.07919560000300407,
+        "avg_rho": 0.08190769484887521,
+        "avg_nudge": -0.029933936273058254,
+        "loss_term": 0.08082282543182373,
+        "loss_bridge": 0.07666538655757904
+      },
+      {
+        "step": 2000,
+        "avg_cos": 0.06641554242620866,
+        "avg_rho": 0.055935436549286045,
+        "avg_nudge": -0.020354578581949074,
+        "loss_term": 0.07134468108415604,
+        "loss_bridge": 0.011993557214736938
+      },
+      {
+        "step": 2500,
+        "avg_cos": 0.0844428426741312,
+        "avg_rho": 0.0917752521733443,
+        "avg_nudge": -0.028687965202455718,
+        "loss_term": 0.0323462039232254,
+        "loss_bridge": 0.00667245127260685
+      },
+      {
+        "step": 3000,
+        "avg_cos": 0.0888429010907809,
+        "avg_rho": 0.06338833862294753,
+        "avg_nudge": -0.026075587142258883,
+        "loss_term": 0.04170331731438637,
+        "loss_bridge": 0.010068882256746292
+      },
+      {
+        "step": 3500,
+        "avg_cos": 0.016165781184099615,
+        "avg_rho": 0.014142975055923065,
+        "avg_nudge": -0.003436643397435546,
+        "loss_term": 0.03528498113155365,
+        "loss_bridge": 0.012309195473790169
+      },
+      {
+        "step": 4000,
+        "avg_cos": 0.06803990031282107,
+        "avg_rho": 0.055018783935035266,
+        "avg_nudge": -0.022906929564972717,
+        "loss_term": 0.027342472225427628,
+        "loss_bridge": 0.007336604408919811
+      },
+      {
+        "step": 4500,
+        "avg_cos": 0.024216643689821165,
+        "avg_rho": 0.04094697698019445,
+        "avg_nudge": -0.008364889615525803,
+        "loss_term": 0.027580715715885162,
+        "loss_bridge": 0.016118880361318588
+      },
+      {
+        "step": 5000,
+        "avg_cos": 0.012039402422184745,
+        "avg_rho": -0.002023946028202772,
+        "avg_nudge": -0.0010173787983755271,
+        "loss_term": 0.027855150401592255,
+        "loss_bridge": 0.010500052943825722
+      }
+    ]
+  },
+  "best_combo": {
+    "best_cos": 0.30479515840609867,
+    "best_step": 500,
+    "final_cos": -0.025737087552746136,
+    "final_rho": -0.01576789258979261,
+    "final_nudge": 0.011260819776604572,
+    "history": [
+      {
+        "step": 500,
+        "avg_cos": 0.30479515840609867,
+        "avg_rho": 0.33246727536122006,
+        "avg_nudge": -0.11396304952601592,
+        "loss_term": 0.5129790306091309,
+        "loss_bridge": 0.21324321627616882
+      },
+      {
+        "step": 1000,
+        "avg_cos": 0.24110793943206468,
+        "avg_rho": 0.24976263443628946,
+        "avg_nudge": -0.08793549550076325,
+        "loss_term": 0.14881110191345215,
+        "loss_bridge": 0.1860560178756714
+      },
+      {
+        "step": 1500,
+        "avg_cos": 0.12106851922969024,
+        "avg_rho": 0.11747027436892192,
+        "avg_nudge": -0.042971268917123474,
+        "loss_term": 0.10358402132987976,
+        "loss_bridge": 0.11228226125240326
+      },
+      {
+        "step": 2000,
+        "avg_cos": 0.19137668733795485,
+        "avg_rho": 0.2001398652791977,
+        "avg_nudge": -0.06597264126564066,
+        "loss_term": 0.0836295336484909,
+        "loss_bridge": 0.03368356451392174
+      },
+      {
+        "step": 2500,
+        "avg_cos": 0.1382010855789607,
+        "avg_rho": 0.15255235826286176,
+        "avg_nudge": -0.05140705577408274,
+        "loss_term": 0.058304790407419205,
+        "loss_bridge": 0.014804087579250336
+      },
+      {
+        "step": 3000,
+        "avg_cos": 0.0815443117171526,
+        "avg_rho": 0.07958398557578523,
+        "avg_nudge": -0.026181443439175684,
+        "loss_term": 0.059965550899505615,
+        "loss_bridge": 0.016236742958426476
+      },
+      {
+        "step": 3500,
+        "avg_cos": 0.09519872162491083,
+        "avg_rho": 0.09258671143713097,
+        "avg_nudge": -0.03418003007148703,
+        "loss_term": 0.03956954926252365,
+        "loss_bridge": 0.014235305599868298
+      },
+      {
+        "step": 4000,
+        "avg_cos": -0.011209671385586262,
+        "avg_rho": -0.007080828654579818,
+        "avg_nudge": 0.004587161975602309,
+        "loss_term": 0.03683673217892647,
+        "loss_bridge": 0.018107034265995026
+      },
+      {
+        "step": 4500,
+        "avg_cos": 0.04608155476550261,
+        "avg_rho": 0.05846519426753124,
+        "avg_nudge": -0.017059470837314922,
+        "loss_term": 0.043314699083566666,
+        "loss_bridge": 0.023334285244345665
+      },
+      {
+        "step": 5000,
+        "avg_cos": -0.025737087552746136,
+        "avg_rho": -0.01576789258979261,
+        "avg_nudge": 0.011260819776604572,
+        "loss_term": 0.03279898688197136,
+        "loss_bridge": 0.015516506507992744
+      }
+    ]
+  },
+  "noise_1.0": {
+    "best_cos": 0.2831856335202853,
+    "best_step": 500,
+    "final_cos": 0.010971122809375325,
+    "final_rho": 0.014127362189659229,
+    "final_nudge": -0.002180874968568484,
+    "history": [
+      {
+        "step": 500,
+        "avg_cos": 0.2831856335202853,
+        "avg_rho": 0.30480503539244336,
+        "avg_nudge": -0.1052918794254462,
+        "loss_term": 0.4858350455760956,
+        "loss_bridge": 0.20565161108970642
+      },
+      {
+        "step": 1000,
+        "avg_cos": 0.14922113033632436,
+        "avg_rho": 0.14803783098856607,
+        "avg_nudge": -0.052950371988117695,
+        "loss_term": 0.1246427372097969,
+        "loss_bridge": 0.1951574981212616
+      },
+      {
+        "step": 1500,
+        "avg_cos": 0.032916761857147016,
+        "avg_rho": 0.030339293957998354,
+        "avg_nudge": -0.011420852970331907,
+        "loss_term": 0.059441905468702316,
+        "loss_bridge": 0.09553220868110657
+      },
+      {
+        "step": 2000,
+        "avg_cos": 0.02735897192421059,
+        "avg_rho": 0.019464978327353794,
+        "avg_nudge": -0.008191200438886881,
+        "loss_term": 0.07519456744194031,
+        "loss_bridge": 0.017568862065672874
+      },
+      {
+        "step": 2500,
+        "avg_cos": 0.02141591941472143,
+        "avg_rho": 0.027642763530214626,
+        "avg_nudge": -0.007318320373694102,
+        "loss_term": 0.05027623474597931,
+        "loss_bridge": 0.014622311107814312
+      },
+      {
+        "step": 3000,
+        "avg_cos": 0.048597725903770574,
+        "avg_rho": 0.03585781451935569,
+        "avg_nudge": -0.015588213689625263,
+        "loss_term": 0.06457968056201935,
+        "loss_bridge": 0.01680811122059822
+      },
+      {
+        "step": 3500,
+        "avg_cos": 0.06355043267831206,
+        "avg_rho": 0.04236283013597131,
+        "avg_nudge": -0.021118728754421074,
+        "loss_term": 0.035970453172922134,
+        "loss_bridge": 0.015534179285168648
+      },
+      {
+        "step": 4000,
+        "avg_cos": 0.04567302499587337,
+        "avg_rho": 0.05423017560193936,
+        "avg_nudge": -0.01528523334612449,
+        "loss_term": 0.034011729061603546,
+        "loss_bridge": 0.009947280399501324
+      },
+      {
+        "step": 4500,
+        "avg_cos": 0.024254882785802085,
+        "avg_rho": 0.023739464891453583,
+        "avg_nudge": -0.009648310175786415,
+        "loss_term": 0.04320281371474266,
+        "loss_bridge": 0.01757601648569107
+      },
+      {
+        "step": 5000,
+        "avg_cos": 0.010971122809375325,
+        "avg_rho": 0.014127362189659229,
+        "avg_nudge": -0.002180874968568484,
+        "loss_term": 0.024934478104114532,
+        "loss_bridge": 0.011717695742845535
+      }
+    ]
+  },
+  "lr_3e-4": {
+    "best_cos": 0.6101815849542618,
+    "best_step": 500,
+    "final_cos": -0.008048781737064322,
+    "final_rho": -0.026018392760306597,
+    "final_nudge": 0.007571722225596507,
+    "history": [
+      {
+        "step": 500,
+        "avg_cos": 0.6101815849542618,
+        "avg_rho": 0.620405301451683,
+        "avg_nudge": -0.2231019102036953,
+        "loss_term": 3.107595205307007,
+        "loss_bridge": 2.7781143188476562
+      },
+      {
+        "step": 1000,
+        "avg_cos": 0.39656415830055874,
+        "avg_rho": 0.39799897621075314,
+        "avg_nudge": -0.14455867062012354,
+        "loss_term": 0.29117679595947266,
+        "loss_bridge": 0.13909414410591125
+      },
+      {
+        "step": 1500,
+        "avg_cos": 0.26989879210789997,
+        "avg_rho": 0.2640038679043452,
+        "avg_nudge": -0.09908402090271314,
+        "loss_term": 0.1437155306339264,
+        "loss_bridge": 0.06845193356275558
+      },
+      {
+        "step": 2000,
+        "avg_cos": 0.15282577524582544,
+        "avg_rho": 0.1327554533878962,
+        "avg_nudge": -0.05070468131452799,
+        "loss_term": 0.10841675102710724,
+        "loss_bridge": 0.044324424117803574
+      },
+      {
+        "step": 2500,
+        "avg_cos": 0.054395756063361965,
+        "avg_rho": 0.04450509278103709,
+        "avg_nudge": -0.01604464929550886,
+        "loss_term": 0.09302366524934769,
+        "loss_bridge": 0.016296718269586563
+      },
+      {
+        "step": 3000,
+        "avg_cos": 0.041961303912103176,
+        "avg_rho": 0.032989607813457646,
+        "avg_nudge": -0.01202150775740544,
+        "loss_term": 0.06656567752361298,
+        "loss_bridge": 0.0089980224147439
+      },
+      {
+        "step": 3500,
+        "avg_cos": 0.010103868204168975,
+        "avg_rho": -0.01982816867530346,
+        "avg_nudge": 0.0005857348442077637,
+        "loss_term": 0.0479322224855423,
+        "loss_bridge": 0.004413220100104809
+      },
+      {
+        "step": 4000,
+        "avg_cos": 0.028387469239532948,
+        "avg_rho": 0.012580555553237597,
+        "avg_nudge": -0.006568876715997855,
+        "loss_term": 0.06536682695150375,
+        "loss_bridge": 0.005300351418554783
+      },
+      {
+        "step": 4500,
+        "avg_cos": 0.015142000513151288,
+        "avg_rho": 0.019022303090120356,
+        "avg_nudge": -0.004821705476691325,
+        "loss_term": 0.037237197160720825,
+        "loss_bridge": 0.005878066644072533
+      },
+      {
+        "step": 5000,
+        "avg_cos": -0.008048781737064322,
+        "avg_rho": -0.026018392760306597,
+        "avg_nudge": 0.007571722225596507,
+        "loss_term": 0.02460472844541073,
+        "loss_bridge": 0.005403056740760803
+      }
+    ]
+  }
+}
\ No newline at end of file
diff --git a/results/toy_lq/toy_lq_seed42.json b/results/toy_lq/toy_lq_seed42.json
new file mode 100644
index 0000000..0a821be
--- /dev/null
+++ b/results/toy_lq/toy_lq_seed42.json
@@ -0,0 +1,335 @@
+{
+  "config": {
+    "d_hidden": 64,
+    "output_dim": 10,
+    "num_layers": 12,
+    "sigma": 0.03,
+    "batch_size": 256,
+    "num_steps": 5000,
+    "lr_fb": 0.001,
+    "lam": 0.1,
+    "K": 8,
+    "ema_momentum": 0.995,
+    "sigma_bridge": 0.03,
+    "eval_every": 500,
+    "seed": 42,
+    "gpu": 0,
+    "output_dir": "results/toy_lq"
+  },
+  "log": {
+    "steps": [
+      1,
+      500,
+      1000,
+      1500,
+      2000,
+      2500,
+      3000,
+      3500,
+      4000,
+      4500,
+      5000
+    ],
+    "state_bridge_loss": [
+      66.01814270019531,
+      2.058396816253662,
+      2.191567897796631,
+      2.1112077236175537,
+      2.2748100757598877,
+      2.1154427528381348,
+      2.0553503036499023,
+      1.966332197189331,
+      2.1147494316101074,
+      2.007577896118164,
+      2.0389387607574463
+    ],
+    "credit_bridge_loss": [
+      112.22884368896484,
+      0.5559965372085571,
+      0.36617255210876465,
+      0.22275042533874512,
+      0.13448813557624817,
+      0.0859837457537651,
+      0.06448937207460403,
+      0.03930438682436943,
+      0.033317387104034424,
+      0.03704077750444412,
+      0.02309737727046013
+    ],
+    "dfa_costate_cos": [
+      0.0011988391992277824,
+      0.007705975646296373,
+      -0.0006242827870524847,
+      0.0037568132393062115,
+      0.004209253507164808,
+      -0.0012397096635630499,
+      0.0033803660816584644,
+      0.003754911944270134,
+      0.002183924035097544,
+      0.0057398807973815845,
+      0.0028941635615550554
+    ],
+    "state_costate_cos": [
+      0.009303608121207127,
+      0.9451924241506137,
+      0.9427064611361577,
+      0.9456643370481638,
+      0.947175892499777,
+      0.9486432488148029,
+      0.9405259856810937,
+      0.9429925313362708,
+      0.9463646090947665,
+      0.9434385620630704,
+      0.9468063895518963
+    ],
+    "credit_costate_cos": [
+      0.03412332414434506,
+      0.3051199523302225,
+      0.2641584941974053,
+      0.17990898627501267,
+      0.11904546274588658,
+      0.03094297769264533,
+      0.022915477577883463,
+      0.008693795901938127,
+      0.0027211602920523058,
+      0.020937439484091904,
+      0.033342053540624104
+    ],
+    "dfa_rho": [
+      0.005011526596111556,
+      0.0020135376447190842,
+      -0.011304221504057447,
+      0.003935044708972176,
+      0.0159942601264144,
+      -0.011545649264007807,
+      0.01096861291443929,
+      0.0007782066240906715,
+      -0.015019190264865756,
+      0.007689292387415965,
+      0.006176682732378443
+    ],
+    "state_rho": [
+      0.011923154350370169,
+      0.9337877084811529,
+      0.9292550335327784,
+      0.9341330577929815,
+      0.9364036619663239,
+      0.9322425921758016,
+      0.9240961174170176,
+      0.9329939633607864,
+      0.9331585764884949,
+      0.9324665367603302,
+      0.9281178514162699
+    ],
+    "credit_rho": [
+      0.031667908265565835,
+      0.31275976697603863,
+      0.2433429310719172,
+      0.1816188059747219,
+      0.11667139704028766,
+      0.013198353117331862,
+      0.022044080891646445,
+      -0.007547003333456814,
+      -0.011566813724736372,
+      -0.003948230994865298,
+      0.016950203105807304
+    ],
+    "dfa_nudge": [
+      -0.0003799900102118651,
+      -0.0025626374408602715,
+      0.0017628272374471028,
+      -0.001205168974896272,
+      -0.0011821148606638114,
+      0.0014717701512078445,
+      -8.787959814071655e-05,
+      -0.0006076549955954155,
+      0.0005303900688886642,
+      -0.0014991160326947768,
+      -0.0005284918782611688
+    ],
+    "state_nudge": [
+      -0.002327537008871635,
+      -0.34574924657742184,
+      -0.3358767156799634,
+      -0.33698513607184094,
+      -0.3561149264375369,
+      -0.33514803399642307,
+      -0.3557068184018135,
+      -0.3218521823485692,
+      -0.33325668424367905,
+      -0.34358637283245724,
+      -0.34090926001469296
+    ],
+    "credit_nudge": [
+      -0.014598140881086389,
+      -0.11370646270612876,
+      -0.09322128010292847,
+      -0.06459770910441875,
+      -0.04436610918492079,
+      -0.008244700108965239,
+      -0.00636714743450284,
+      0.0006991980286935965,
+      0.0036097665627797446,
+      -0.0035287897723416486,
+      -0.007594603579491377
+    ],
+    "bridge_residual": [
+      0.06566914729773998,
+      0.37236853316426277,
+      0.3114783614873886,
+      0.26258066420753795,
+      0.20779911428689957,
+      0.13781529137243828,
+      0.10999641008675098,
+      0.07969006771842639,
+      0.06178054213523865,
+      0.07378745886186759,
+      0.0832565538585186
+    ]
+  },
+  "final_per_layer": {
+    "dfa_costate_cos": [
+      -0.039167389273643494,
+      -0.0378018394112587,
+      0.005690325051546097,
+      -0.023073989897966385,
+      -0.0005057593807578087,
+      -0.014485953375697136,
+      0.03301015496253967,
+      0.04401148855686188,
+      0.054177843034267426,
+      0.03981431573629379,
+      -0.04246171563863754,
+      -0.0012151142582297325,
+      0.01963176019489765
+    ],
+    "state_costate_cos": [
+      0.9444395303726196,
+      0.9453534483909607,
+      0.9460644721984863,
+      0.9466040134429932,
+      0.9469730257987976,
+      0.9472954273223877,
+      0.9476633667945862,
+      0.9478192925453186,
+      0.947782576084137,
+      0.9473406672477722,
+      0.947346568107605,
+      0.9471475481987,
+      0.9466531276702881
+    ],
+    "credit_costate_cos": [
+      0.04752141237258911,
+      0.04377397149801254,
+      0.04051002860069275,
+      0.03716123104095459,
+      0.0350164994597435,
+      0.03194836527109146,
+      0.02999947965145111,
+      0.02913709171116352,
+      0.027683185413479805,
+      0.0277146864682436,
+      0.027703404426574707,
+      0.027434173971414566,
+      0.027843166142702103
+    ],
+    "dfa_rho": [
+      -0.04803554713726044,
+      0.001050771214067936,
+      0.008967258036136627,
+      -0.0271889790892601,
+      0.02336559258401394,
+      -0.018210411071777344,
+      0.05891512706875801,
+      0.040720634162425995,
+      0.07478035986423492,
+      0.04802168905735016,
+      -0.06280035525560379,
+      -0.0254659466445446
+    ],
+    "state_rho": [
+      0.9311305284500122,
+      0.9222633838653564,
+      0.9287852644920349,
+      0.9287664890289307,
+      0.9245603084564209,
+      0.928197979927063,
+      0.9275168180465698,
+      0.9290561676025391,
+      0.9267844557762146,
+      0.927483081817627,
+      0.9315245151519775,
+      0.9313452243804932
+    ],
+    "credit_rho": [
+      0.05740518122911453,
+      0.035541512072086334,
+      0.002091987058520317,
+      0.024556485936045647,
+      -0.006993812508881092,
+      0.03284040838479996,
+      0.012268777936697006,
+      -0.004999782890081406,
+      0.014774687588214874,
+      -0.010628825053572655,
+      0.05940534919500351,
+      -0.012859531678259373
+    ],
+    "dfa_nudge": [
+      0.01363457553088665,
+      0.013758538290858269,
+      -0.0032786596566438675,
+      0.010209780186414719,
+      -0.0013850200921297073,
+      0.004463233053684235,
+      -0.012735363095998764,
+      -0.01801125332713127,
+      -0.019914839416742325,
+      -0.012350432574748993,
+      0.017156170681118965,
+      0.002111367881298065
+    ],
+    "state_nudge": [
+      -0.34207087755203247,
+      -0.3417494297027588,
+      -0.3413676619529724,
+      -0.3414176106452942,
+      -0.3412688076496124,
+      -0.34025871753692627,
+      -0.340381920337677,
+      -0.34054237604141235,
+      -0.3402785062789917,
+      -0.34061968326568604,
+      -0.34013575315475464,
+      -0.340819776058197
+    ],
+    "credit_nudge": [
+      -0.012597911059856415,
+      -0.011231275275349617,
+      -0.009987404569983482,
+      -0.008699672296643257,
+      -0.007977090775966644,
+      -0.006919408217072487,
+      -0.006156830117106438,
+      -0.005891043692827225,
+      -0.00546320341527462,
+      -0.005353953689336777,
+      -0.005423387512564659,
+      -0.005434062331914902
+    ],
+    "bridge_residual": [
+      0.10028190910816193,
+      0.09995798766613007,
+      0.09832010418176651,
+      0.09528136253356934,
+      0.09090456366539001,
+      0.0843343734741211,
+      0.07823348790407181,
+      0.0723348930478096,
+      0.06821676343679428,
+      0.06773027032613754,
+      0.06943590939044952,
+      0.0740470215678215
+    ]
+  }
+}
\ No newline at end of file
diff --git a/results/toy_lq/toy_lq_v2_seed123_lam0.1_sig0.1_tgw1.0_fm0.0.json b/results/toy_lq/toy_lq_v2_seed123_lam0.1_sig0.1_tgw1.0_fm0.0.json
new file mode 100644
index 0000000..56b3336
--- /dev/null
+++ b/results/toy_lq/toy_lq_v2_seed123_lam0.1_sig0.1_tgw1.0_fm0.0.json
@@ -0,0 +1,330 @@
+{
+  "config": {
+    "d_hidden": 64,
+    "output_dim": 10,
+    "num_layers": 12,
+    "sigma": 0.03,
+    "batch_size": 256,
+    "num_steps": 8000,
+    "lr_fb": 0.001,
+    "lam": 0.1,
+    "K": 8,
+    "ema_momentum": 0.995,
+    "sigma_bridge": 0.1,
+    "eval_every": 1000,
+    "seed": 123,
+    "gpu": 0,
+    "output_dir": "results/toy_lq",
+    "vnet_hidden": 256,
+    "vnet_layers": 3,
+    "term_grad_weight": 1.0,
+    "fm_weight": 0.0
+  },
+  "log": {
+    "steps": [
+      1,
+      1000,
+      2000,
+      3000,
+      4000,
+      5000,
+      6000,
+      7000,
+      8000
+    ],
+    "dfa_costate_cos": [
+      0.0061469420325011015,
+      0.007492478704079986,
+      0.006436596314112346,
+      0.001648913836106658,
+      0.005657717352733016,
+      0.010142655577510595,
+      0.005493079700196783,
+      0.008208209726338586,
+      0.002802101274331411
+    ],
+    "state_costate_cos": [
+      0.04875442975511154,
+      0.9345830877621969,
+      0.9331425180037817,
+      0.9344809154669443,
+      0.9360056469837824,
+      0.9349906196196874,
+      0.9401087015867233,
+      0.9377894500891367,
+      0.9381228536367416
+    ],
+    "credit_costate_cos": [
+      0.005350367398932576,
+      0.8715064575274786,
+      0.9082922885815302,
+      0.9268463253974915,
+      0.9348577807346979,
+      0.9338823060194651,
+      0.9398750513792038,
+      0.9371217240889868,
+      0.9403592944145203
+    ],
+    "dfa_rho": [
+      0.014851124413932363,
+      0.004832483362406492,
+      0.005500619454930226,
+      0.0014784028753638268,
+      0.0024716570042073727,
+      0.002679668522129456,
+      0.004042171291075647,
+      0.007841781480237842,
+      -0.003731151965136329
+    ],
+    "state_rho": [
+      0.0525277191773057,
+      0.9209283490975698,
+      0.9212760378917059,
+      0.9215241422255834,
+      0.9252830098072687,
+      0.9173514097929001,
+      0.9268933484951655,
+      0.9245105236768723,
+      0.9267788628737131
+    ],
+    "credit_rho": [
+      8.900166722014546e-05,
+      0.8210234741369883,
+      0.8804274102052053,
+      0.9100336680809656,
+      0.9228341629107794,
+      0.9162226617336273,
+      0.9262114216883978,
+      0.9239083131154379,
+      0.9273379941781362
+    ],
+    "dfa_nudge": [
+      -0.0020856610499322414,
+      -0.002391135785728693,
+      -0.0018826122395694256,
+      0.0002794961134592692,
+      -0.0017084906188150246,
+      -0.0027558018919080496,
+      -0.0014085437481602032,
+      -0.0023310642379025617,
+      -0.0009137461893260479
+    ],
+    "state_nudge": [
+      -0.01762500188002984,
+      -0.31959830472866696,
+      -0.3143775438268979,
+      -0.3134472444653511,
+      -0.33237058420976,
+      -0.3183835695187251,
+      -0.32320864746967953,
+      -0.32294898976882297,
+      -0.31496328860521317
+    ],
+    "credit_nudge": [
+      -0.0004618208234508832,
+      -0.2996478999654452,
+      -0.30627985050280887,
+      -0.3099561383326848,
+      -0.33090290675560635,
+      -0.31673717498779297,
+      -0.3218600004911423,
+      -0.32163529098033905,
+      -0.31442063798507053
+    ],
+    "bridge_residual": [],
+    "state_bridge_loss": [
+      64.66886901855469,
+      1.9055249691009521,
+      1.9802964925765991,
+      1.9747117757797241,
+      1.9333608150482178,
+      2.0244460105895996,
+      2.235288381576538,
+      1.8483705520629883,
+      1.833404541015625
+    ],
+    "credit_bridge_loss": [
+      129.2601776123047,
+      8.826760292053223,
+      8.817267417907715,
+      9.7413911819458,
+      8.333325386047363,
+      8.390682220458984,
+      8.381990432739258,
+      9.069635391235352,
+      8.738546371459961
+    ],
+    "term_loss": [
+      109.68403625488281,
+      3.380112648010254,
+      4.175432205200195,
+      4.737654685974121,
+      3.632157802581787,
+      3.3351938724517822,
+      3.776655912399292,
+      4.416824817657471,
+      4.120006084442139
+    ],
+    "bridge_loss": [
+      5.943464884694549e-07,
+      0.2433367669582367,
+      0.18713834881782532,
+      0.12417592853307724,
+      0.13950209319591522,
+      0.14192476868629456,
+      0.15276893973350525,
+      0.10663559287786484,
+      0.12137635797262192
+    ],
+    "term_grad_loss": [
+      19.57614517211914,
+      5.203310966491699,
+      4.4546966552734375,
+      4.879560470581055,
+      4.5616655349731445,
+      4.9135637283325195,
+      4.452565670013428,
+      4.546175003051758,
+      4.497163772583008
+    ],
+    "fm_loss": [
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0
+    ]
+  },
+  "final_per_layer": {
+    "dfa_costate_cos": [
+      0.04479089006781578,
+      -0.04426354169845581,
+      0.008874599821865559,
+      0.05599997937679291,
+      0.02961653470993042,
+      -0.022058574482798576,
+      0.027186769992113113,
+      -0.0337681919336319,
+      -0.020245034247636795,
+      -0.04076787084341049,
+      0.006424968130886555,
+      0.021834686398506165
+    ],
+    "state_costate_cos": [
+      0.9351104497909546,
+      0.9362065196037292,
+      0.9370632171630859,
+      0.9374789595603943,
+      0.9381210803985596,
+      0.9385455846786499,
+      0.9386894702911377,
+      0.9390624165534973,
+      0.9393105506896973,
+      0.9393242597579956,
+      0.9392982721328735,
+      0.9392634630203247
+    ],
+    "credit_costate_cos": [
+      0.9360430240631104,
+      0.9369803667068481,
+      0.9380120038986206,
+      0.9385954141616821,
+      0.9395070672035217,
+      0.9403176307678223,
+      0.9409258961677551,
+      0.9413450956344604,
+      0.9420279264450073,
+      0.9428501725196838,
+      0.9435971975326538,
+      0.9441097378730774
+    ],
+    "dfa_rho": [
+      0.041250549256801605,
+      -0.049739208072423935,
+      0.00025176629424095154,
+      0.04237007349729538,
+      0.040798719972372055,
+      -0.037202395498752594,
+      0.004211327061057091,
+      0.0017704367637634277,
+      -0.04704931750893593,
+      -0.02992381900548935,
+      -0.01541070081293583,
+      0.003898744471371174
+    ],
+    "state_rho": [
+      0.9230573177337646,
+      0.9292065501213074,
+      0.9277745485305786,
+      0.926990270614624,
+      0.930223822593689,
+      0.9262816905975342,
+      0.9238564968109131,
+      0.9235492944717407,
+      0.927101731300354,
+      0.929185152053833,
+      0.9233707189559937,
+      0.9307487607002258
+    ],
+    "credit_rho": [
+      0.924521803855896,
+      0.9213208556175232,
+      0.923781156539917,
+      0.9222273826599121,
+      0.9230118989944458,
+      0.9317750930786133,
+      0.9238618016242981,
+      0.9351733326911926,
+      0.9263075590133667,
+      0.9302359819412231,
+      0.9358923435211182,
+      0.9299467206001282
+    ],
+    "dfa_nudge": [
+      -0.01619502529501915,
+      0.016632290557026863,
+      -0.004560360684990883,
+      -0.018955951556563377,
+      -0.011058392003178596,
+      0.007517071440815926,
+      -0.008148249238729477,
+      0.01091383583843708,
+      0.006811931729316711,
+      0.015017258003354073,
+      -0.000978708267211914,
+      -0.00796065479516983
+    ],
+    "state_nudge": [
+      -0.31427979469299316,
+      -0.3147982358932495,
+      -0.3148774206638336,
+      -0.3150269389152527,
+      -0.31536394357681274,
+      -0.3156891465187073,
+      -0.3149397373199463,
+      -0.3145085573196411,
+      -0.3144562840461731,
+      -0.31494998931884766,
+      -0.3152415156364441,
+      -0.31542789936065674
+    ],
+    "credit_nudge": [
+      -0.31250905990600586,
+      -0.3131006360054016,
+      -0.31331712007522583,
+      -0.31366780400276184,
+      -0.3142547011375427,
+      -0.3149237632751465,
+      -0.31442567706108093,
+      -0.314262330532074,
+      -0.31445154547691345,
+      -0.3153681755065918,
+      -0.3161371052265167,
+      -0.3166297376155853
+    ]
+  }
+}
\ No newline at end of file
diff --git a/results/toy_lq/toy_lq_v2_seed42_lam0.1_sig0.1_tgw1.0_fm0.0.json b/results/toy_lq/toy_lq_v2_seed42_lam0.1_sig0.1_tgw1.0_fm0.0.json
new file mode 100644
index 0000000..14050eb
--- /dev/null
+++ b/results/toy_lq/toy_lq_v2_seed42_lam0.1_sig0.1_tgw1.0_fm0.0.json
@@ -0,0 +1,458 @@
+{
+  "config": {
+    "d_hidden": 64,
+    "output_dim": 10,
+    "num_layers": 12,
+    "sigma": 0.03,
+    "batch_size": 256,
+    "num_steps": 8000,
+    "lr_fb": 0.001,
+    "lam": 0.1,
+    "K": 8,
+    "ema_momentum": 0.995,
+    "sigma_bridge": 0.1,
+    "eval_every": 500,
+    "seed": 42,
+    "gpu": 0,
+    "output_dir": "results/toy_lq",
+    "vnet_hidden": 256,
+    "vnet_layers": 3,
+    "term_grad_weight": 1.0,
+    "fm_weight": 0.0
+  },
+  "log": {
+    "steps": [
+      1,
+      500,
+      1000,
+      1500,
+      2000,
+      2500,
+      3000,
+      3500,
+      4000,
+      4500,
+      5000,
+      5500,
+      6000,
+      6500,
+      7000,
+      7500,
+      8000
+    ],
+    "dfa_costate_cos": [
+      0.001022340264171362,
+      0.00324871806272616,
+      0.003955680275491129,
+      -6.335701133745412e-05,
+      0.0009047117297692845,
+      -0.0015635235855976741,
+      0.0050255006256823736,
+      0.003974012487257521,
+      -0.0024126653637116155,
+      0.0010036304011009634,
+      -0.0006338046708454689,
+      0.00017608713824301958,
+      -9.714129070440929e-05,
+      -0.0011926805212472875,
+      -0.0026168684319903455,
+      -0.003671417556082209,
+      0.0016770829679444432
+    ],
+    "state_costate_cos": [
+      0.009337988216429949,
+      0.943408285578092,
+      0.9517824550469717,
+      0.9448912839094797,
+      0.9432462950547537,
+      0.9444176405668259,
+      0.9435405482848486,
+      0.9443136354287466,
+      0.9458288550376892,
+      0.9465046127637228,
+      0.9465998113155365,
+      0.9461832990248998,
+      0.9424780905246735,
+      0.9474884221951166,
+      0.9449383119742075,
+      0.9425351719061533,
+      0.9405365288257599
+    ],
+    "credit_costate_cos": [
+      0.024892715892444055,
+      0.8234576731920242,
+      0.8861371924479803,
+      0.8940252065658569,
+      0.9119344304005305,
+      0.928387979666392,
+      0.9261951943238577,
+      0.9374307443698248,
+      0.939670259753863,
+      0.9423713783423106,
+      0.9447299987077713,
+      0.944857731461525,
+      0.9409955541292826,
+      0.9459459533294042,
+      0.9439971546332041,
+      0.9420813073714575,
+      0.9400773843129476
+    ],
+    "dfa_rho": [
+      0.015879416760678094,
+      0.004265802912414074,
+      0.008714484671751658,
+      0.001407407767449816,
+      0.007156838779337704,
+      -0.0010706717148423195,
+      -0.00500367038572828,
+      0.00037602245962868136,
+      0.00797194141584138,
+      0.012475600582547486,
+      0.006475616673318048,
+      0.006899521841357152,
+      0.008204833992446462,
+      -0.0002780493038396041,
+      -0.009471656677002708,
+      -0.00781721225939691,
+      0.025706250220537186
+    ],
+    "state_rho": [
+      0.0029325426245729127,
+      0.9265175064404806,
+      0.938806007305781,
+      0.9280826350053152,
+      0.9287765026092529,
+      0.9266663392384847,
+      0.9301863809426626,
+      0.9323704888423284,
+      0.9297124246756235,
+      0.9302101731300354,
+      0.9340693553288778,
+      0.9324037134647369,
+      0.9265123655398687,
+      0.9321505973736445,
+      0.927787164847056,
+      0.9323761165142059,
+      0.9224280416965485
+    ],
+    "credit_rho": [
+      0.02234963719577839,
+      0.7719902147849401,
+      0.8432890474796295,
+      0.8447659462690353,
+      0.877406562368075,
+      0.8958166440327963,
+      0.9122767845789591,
+      0.9254603485266367,
+      0.9236042300860087,
+      0.9263056516647339,
+      0.9323819329341253,
+      0.9314924577871958,
+      0.9266939163208008,
+      0.929329847296079,
+      0.9257104198137919,
+      0.9289217789967855,
+      0.9269137680530548
+    ],
+    "dfa_nudge": [
+      -0.0003799900102118651,
+      -0.0009698765352368355,
+      -0.0013957968913018703,
+      0.00035563452790180844,
+      -0.0003614289841304223,
+      0.0008372832089662552,
+      -0.0013908503266672294,
+      -0.0013163429684937,
+      0.0005854369762043158,
+      -0.000560786963130037,
+      0.00015371766251822314,
+      -1.0695696497956911e-05,
+      2.807355485856533e-05,
+      5.036763225992521e-05,
+      0.0008217894161740938,
+      0.0017250357971837123,
+      -0.0004334271264572938
+    ],
+    "state_nudge": [
+      -0.002327537008871635,
+      -0.34193428109089535,
+      -0.340603639682134,
+      -0.35042013972997665,
+      -0.3389856591820717,
+      -0.34554000198841095,
+      -0.3481511374314626,
+      -0.3558509051799774,
+      -0.3337947155038516,
+      -0.33715616663297016,
+      -0.33583804468313855,
+      -0.35106155276298523,
+      -0.34434546530246735,
+      -0.34641525397698086,
+      -0.3329972525437673,
+      -0.34870391835769016,
+      -0.3395152688026428
+    ],
+    "credit_nudge": [
+      -0.0079942528779308,
+      -0.30343081553777057,
+      -0.31936119496822357,
+      -0.33327333877484006,
+      -0.3284662067890167,
+      -0.3394550507267316,
+      -0.3424832498033841,
+      -0.3524467721581459,
+      -0.33071904132763547,
+      -0.3346964443723361,
+      -0.3337800477941831,
+      -0.34932391593853634,
+      -0.34230031818151474,
+      -0.3444634775320689,
+      -0.3310972551504771,
+      -0.34694332132736844,
+      -0.3377286195755005
+    ],
+    "bridge_residual": [],
+    "state_bridge_loss": [
+      66.01814270019531,
+      2.0012869834899902,
+      2.1027088165283203,
+      2.1019272804260254,
+      2.0727572441101074,
+      1.9770110845565796,
+      2.1761908531188965,
+      2.094480514526367,
+      1.9725749492645264,
+      2.0142102241516113,
+      2.0340821743011475,
+      1.9380583763122559,
+      1.989743947982788,
+      2.4057328701019287,
+      2.19437575340271,
+      2.1816155910491943,
+      2.0803794860839844
+    ],
+    "credit_bridge_loss": [
+      132.09298706054688,
+      11.307573318481445,
+      9.575517654418945,
+      8.768391609191895,
+      8.752981185913086,
+      9.271968841552734,
+      8.606213569641113,
+      8.753968238830566,
+      8.223488807678223,
+      9.935165405273438,
+      9.02967357635498,
+      8.346063613891602,
+      9.022041320800781,
+      8.36446762084961,
+      8.635725021362305,
+      8.760185241699219,
+      8.503408432006836
+    ],
+    "term_loss": [
+      111.63633728027344,
+      4.978545188903809,
+      3.8962953090667725,
+      3.81073260307312,
+      4.386394500732422,
+      4.507748603820801,
+      3.6740365028381348,
+      3.7450127601623535,
+      3.5060954093933105,
+      4.545898914337158,
+      4.322302341461182,
+      3.594371795654297,
+      4.038668632507324,
+      4.025404453277588,
+      3.9016177654266357,
+      3.7638583183288574,
+      3.807260036468506
+    ],
+    "bridge_loss": [
+      6.45359421014291e-07,
+      0.432157039642334,
+      0.20619139075279236,
+      0.24715952575206757,
+      0.1523856669664383,
+      0.12624874711036682,
+      0.13425695896148682,
+      0.16367560625076294,
+      0.14530189335346222,
+      0.18452416360378265,
+      0.11599670350551605,
+      0.11983858048915863,
+      0.11901542544364929,
+      0.16489851474761963,
+      0.16058529913425446,
+      0.09414370357990265,
+      0.12402483820915222
+    ],
+    "term_grad_loss": [
+      20.456655502319336,
+      5.8968706130981445,
+      5.4730305671691895,
+      4.710499286651611,
+      4.214200496673584,
+      4.637970924377441,
+      4.797920227050781,
+      4.845280170440674,
+      4.572091579437256,
+      5.204742908477783,
+      4.591374397277832,
+      4.631853103637695,
+      4.864356994628906,
+      4.174164772033691,
+      4.573522090911865,
+      4.9021830558776855,
+      4.5721235275268555
+    ],
+    "fm_loss": [
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0
+    ]
+  },
+  "final_per_layer": {
+    "dfa_costate_cos": [
+      -0.03974368795752525,
+      -0.03234883397817612,
+      -0.01280729565769434,
+      -0.03282422572374344,
+      0.005050511099398136,
+      -0.01393066719174385,
+      0.036346256732940674,
+      0.04377925395965576,
+      0.06324614584445953,
+      0.05324942618608475,
+      -0.04297419637441635,
+      -0.006917691323906183
+    ],
+    "state_costate_cos": [
+      0.9383341073989868,
+      0.9390445947647095,
+      0.9393846988677979,
+      0.9396730661392212,
+      0.9402315616607666,
+      0.940619707107544,
+      0.9413514137268066,
+      0.9420933723449707,
+      0.941848874092102,
+      0.9416857957839966,
+      0.9411274194717407,
+      0.9410437345504761
+    ],
+    "credit_costate_cos": [
+      0.937440037727356,
+      0.9379364252090454,
+      0.9384041428565979,
+      0.9385387897491455,
+      0.9392049312591553,
+      0.9398118257522583,
+      0.9403334259986877,
+      0.9412999153137207,
+      0.94156414270401,
+      0.9419326186180115,
+      0.9419490098953247,
+      0.9425133466720581
+    ],
+    "dfa_rho": [
+      -0.023802103474736214,
+      -0.039411380887031555,
+      -0.008970402181148529,
+      -0.0021191751584410667,
+      0.04573667049407959,
+      0.010564171709120274,
+      0.04995737224817276,
+      0.10094872862100601,
+      0.07801118493080139,
+      0.07188688963651657,
+      -0.02006501331925392,
+      0.045738060027360916
+    ],
+    "state_rho": [
+      0.9204449653625488,
+      0.917449951171875,
+      0.9193954467773438,
+      0.9267550706863403,
+      0.9280773401260376,
+      0.9225738048553467,
+      0.9190236926078796,
+      0.924782931804657,
+      0.9248216152191162,
+      0.9217315912246704,
+      0.9215092658996582,
+      0.9225708246231079
+    ],
+    "credit_rho": [
+      0.9222818613052368,
+      0.923028826713562,
+      0.924055814743042,
+      0.9263180494308472,
+      0.9288915991783142,
+      0.922785758972168,
+      0.9231191873550415,
+      0.9254114627838135,
+      0.9321990013122559,
+      0.9311563372612,
+      0.934211015701294,
+      0.9295063018798828
+    ],
+    "dfa_nudge": [
+      0.014356113970279694,
+      0.013086749240756035,
+      0.004036703146994114,
+      0.011657552793622017,
+      -0.0023863790556788445,
+      0.006091207265853882,
+      -0.013204541988670826,
+      -0.016102034598588943,
+      -0.022989045828580856,
+      -0.018843289464712143,
+      0.015521062538027763,
+      0.0035747764632105827
+    ],
+    "state_nudge": [
+      -0.34063079953193665,
+      -0.34022200107574463,
+      -0.33977580070495605,
+      -0.3401387929916382,
+      -0.34019535779953003,
+      -0.33964985609054565,
+      -0.3392874002456665,
+      -0.33898109197616577,
+      -0.3392573297023773,
+      -0.3390581011772156,
+      -0.33809930086135864,
+      -0.33888739347457886
+    ],
+    "credit_nudge": [
+      -0.3375471234321594,
+      -0.33736705780029297,
+      -0.33710652589797974,
+      -0.33756691217422485,
+      -0.3378610908985138,
+      -0.33758991956710815,
+      -0.33741605281829834,
+      -0.3374325633049011,
+      -0.33806926012039185,
+      -0.33822208642959595,
+      -0.33765909075737,
+      -0.3389057517051697
+    ]
+  }
+}
\ No newline at end of file
diff --git a/results/toy_lq/toy_lq_v2_seed42_lam0.1_sig0.1_tgw1.0_fm0.1.json b/results/toy_lq/toy_lq_v2_seed42_lam0.1_sig0.1_tgw1.0_fm0.1.json
new file mode 100644
index 0000000..2977e98
--- /dev/null
+++ b/results/toy_lq/toy_lq_v2_seed42_lam0.1_sig0.1_tgw1.0_fm0.1.json
@@ -0,0 +1,330 @@
+{
+  "config": {
+    "d_hidden": 64,
+    "output_dim": 10,
+    "num_layers": 12,
+    "sigma": 0.03,
+    "batch_size": 256,
+    "num_steps": 8000,
+    "lr_fb": 0.001,
+    "lam": 0.1,
+    "K": 8,
+    "ema_momentum": 0.995,
+    "sigma_bridge": 0.1,
+    "eval_every": 1000,
+    "seed": 42,
+    "gpu": 0,
+    "output_dir": "results/toy_lq",
+    "vnet_hidden": 256,
+    "vnet_layers": 3,
+    "term_grad_weight": 1.0,
+    "fm_weight": 0.1
+  },
+  "log": {
+    "steps": [
+      1,
+      1000,
+      2000,
+      3000,
+      4000,
+      5000,
+      6000,
+      7000,
+      8000
+    ],
+    "dfa_costate_cos": [
+      -0.003210080186060319,
+      0.006517285481095314,
+      0.002081584728633364,
+      0.002482607301014165,
+      0.0003743169557613631,
+      -0.0018058015266433358,
+      0.005970992535973589,
+      0.0013112322388527293,
+      -0.00010686229992037018
+    ],
+    "state_costate_cos": [
+      -0.008148168989767631,
+      0.942050834496816,
+      0.9447367936372757,
+      0.9457606921593348,
+      0.9468860725561777,
+      0.9432125687599182,
+      0.9417577485243479,
+      0.945394163330396,
+      0.9473433345556259
+    ],
+    "credit_costate_cos": [
+      -0.011870964197441936,
+      0.8804089923699697,
+      0.9169842700163523,
+      0.9352772484223048,
+      0.9427581528822581,
+      0.9391989062229792,
+      0.9398341725269953,
+      0.9456242968638738,
+      0.9460541109244028
+    ],
+    "dfa_rho": [
+      -0.012772266054525971,
+      0.015942092907304566,
+      0.008943260026474794,
+      0.004492536109561722,
+      -0.007791806012392044,
+      -0.007830069400370121,
+      0.014392409706488252,
+      0.0063066319562494755,
+      -0.001311147507900993
+    ],
+    "state_rho": [
+      -0.0018634579222028453,
+      0.9263874938090643,
+      0.9358506997426351,
+      0.9348766555388769,
+      0.9348586251338323,
+      0.9323162386814753,
+      0.9287882298231125,
+      0.9243461688359579,
+      0.9346484492222468
+    ],
+    "credit_rho": [
+      -0.023325924955618877,
+      0.8106876164674759,
+      0.8797721316417059,
+      0.9208102275927862,
+      0.9290666033824285,
+      0.9226995905240377,
+      0.9258128056923548,
+      0.9266915867726008,
+      0.9364602218071619
+    ],
+    "dfa_nudge": [
+      0.0015840742271393538,
+      -0.0023392424918711185,
+      -0.0002702907659113407,
+      -0.0010186488119264443,
+      -0.0001735797462364038,
+      0.000983052421361208,
+      -0.0020705487113445997,
+      -0.0009097627674539884,
+      0.0002721068449318409
+    ],
+    "state_nudge": [
+      0.0024757004963854947,
+      -0.33942782630523044,
+      -0.35110870252052945,
+      -0.35064691056807834,
+      -0.3403419057528178,
+      -0.3532385254899661,
+      -0.33996084084113437,
+      -0.35185040285189945,
+      -0.3480343023935954
+    ],
+    "credit_nudge": [
+      0.006069206205817561,
+      -0.31947339574495953,
+      -0.3404841795563698,
+      -0.3458903282880783,
+      -0.3376796667774518,
+      -0.3504715636372566,
+      -0.3379351521531741,
+      -0.3500200683871905,
+      -0.3462969238559405
+    ],
+    "bridge_residual": [],
+    "state_bridge_loss": [
+      66.01814270019531,
+      2.341611623764038,
+      1.9537389278411865,
+      1.9830116033554077,
+      2.0491604804992676,
+      2.0490386486053467,
+      2.1182405948638916,
+      2.373213291168213,
+      1.9122812747955322
+    ],
+    "credit_bridge_loss": [
+      132.09298706054688,
+      9.123503684997559,
+      8.516526222229004,
+      8.634014129638672,
+      8.720410346984863,
+      8.43734359741211,
+      10.247673034667969,
+      8.351385116577148,
+      8.474419593811035
+    ],
+    "term_loss": [
+      111.63633728027344,
+      4.155745029449463,
+      3.8754897117614746,
+      4.040826320648193,
+      4.010752201080322,
+      4.2646074295043945,
+      5.391899108886719,
+      3.27858829498291,
+      3.74959135055542
+    ],
+    "bridge_loss": [
+      6.45359421014291e-07,
+      0.22446846961975098,
+      0.18948684632778168,
+      0.14564603567123413,
+      0.1301368772983551,
+      0.1614246517419815,
+      0.19166265428066254,
+      0.19554881751537323,
+      0.15442883968353271
+    ],
+    "term_grad_loss": [
+      20.456655502319336,
+      4.743132591247559,
+      4.45129919052124,
+      4.447327136993408,
+      4.579296112060547,
+      4.011109352111816,
+      4.663890838623047,
+      4.877087593078613,
+      4.57021427154541
+    ],
+    "fm_loss": [
+      1.4957863925246784e-07,
+      0.001574160298332572,
+      0.0025115222670137882,
+      0.002142899436876178,
+      0.0022483705542981625,
+      0.0020201210863888264,
+      0.00219835271127522,
+      0.0016042720526456833,
+      0.0018535356502979994
+    ]
+  },
+  "final_per_layer": {
+    "dfa_costate_cos": [
+      -0.04993891716003418,
+      -0.03484980762004852,
+      0.008264334872364998,
+      -0.018589604645967484,
+      -0.019490346312522888,
+      0.0018934037070721388,
+      0.025425557047128677,
+      0.046417269855737686,
+      0.06304077804088593,
+      0.03301604464650154,
+      -0.06211906671524048,
+      0.005648006685078144
+    ],
+    "state_costate_cos": [
+      0.9454483985900879,
+      0.946227490901947,
+      0.9468960165977478,
+      0.9470230340957642,
+      0.9475629925727844,
+      0.9478208422660828,
+      0.9478966593742371,
+      0.9480605125427246,
+      0.9479074478149414,
+      0.9478251934051514,
+      0.9478766918182373,
+      0.9475747346878052
+    ],
+    "credit_costate_cos": [
+      0.9432812333106995,
+      0.9438751935958862,
+      0.9443729519844055,
+      0.9447157382965088,
+      0.9455237984657288,
+      0.9461240768432617,
+      0.9463104009628296,
+      0.9467931985855103,
+      0.9471821784973145,
+      0.9476308226585388,
+      0.9481172561645508,
+      0.9487224817276001
+    ],
+    "dfa_rho": [
+      -0.0704549178481102,
+      0.031227122992277145,
+      0.009932642802596092,
+      -0.015737976878881454,
+      -0.045219600200653076,
+      0.013525029644370079,
+      0.03365146368741989,
+      0.04508150368928909,
+      0.05846859887242317,
+      7.291417568922043e-05,
+      -0.06866942346096039,
+      -0.007611127570271492
+    ],
+    "state_rho": [
+      0.935008704662323,
+      0.9310771822929382,
+      0.9295646548271179,
+      0.938683032989502,
+      0.9329910278320312,
+      0.9313007593154907,
+      0.9357653260231018,
+      0.9376221895217896,
+      0.9326146841049194,
+      0.9372754693031311,
+      0.9379282593727112,
+      0.9359501004219055
+    ],
+    "credit_rho": [
+      0.9306489825248718,
+      0.9282867312431335,
+      0.9364583492279053,
+      0.9361423850059509,
+      0.9323122501373291,
+      0.9380875825881958,
+      0.93974369764328,
+      0.9351372122764587,
+      0.941116213798523,
+      0.9382349252700806,
+      0.9381735324859619,
+      0.9431807994842529
+    ],
+    "dfa_nudge": [
+      0.020365234464406967,
+      0.014575351029634476,
+      -0.002057630568742752,
+      0.006316322833299637,
+      0.007856002077460289,
+      0.0004078727215528488,
+      -0.010462434962391853,
+      -0.017474167048931122,
+      -0.02493620105087757,
+      -0.011619189754128456,
+      0.022399690002202988,
+      -0.00210556760430336
+    ],
+    "state_nudge": [
+      -0.34937936067581177,
+      -0.3488026261329651,
+      -0.3481972813606262,
+      -0.34827157855033875,
+      -0.3484804630279541,
+      -0.3478168845176697,
+      -0.34806668758392334,
+      -0.3481366038322449,
+      -0.3481329381465912,
+      -0.3476967513561249,
+      -0.3463967740535736,
+      -0.34703367948532104
+    ],
+    "credit_nudge": [
+      -0.34637150168418884,
+      -0.345956414937973,
+      -0.3455013632774353,
+      -0.34578341245651245,
+      -0.3462638854980469,
+      -0.3458409905433655,
+      -0.3462728261947632,
+      -0.3466273546218872,
+      -0.3469495475292206,
+      -0.34689074754714966,
+      -0.3459630608558655,
+      -0.34714198112487793
+    ]
+  }
+}
\ No newline at end of file
diff --git a/results/toy_lq/toy_lq_v2_seed42_lam1.0_sig0.3_tgw0.0_fm0.0.json b/results/toy_lq/toy_lq_v2_seed42_lam1.0_sig0.3_tgw0.0_fm0.0.json
new file mode 100644
index 0000000..78451b9
--- /dev/null
+++ b/results/toy_lq/toy_lq_v2_seed42_lam1.0_sig0.3_tgw0.0_fm0.0.json
@@ -0,0 +1,282 @@
+{
+  "config": {
+    "d_hidden": 64,
+    "output_dim": 10,
+    "num_layers": 12,
+    "sigma": 0.03,
+    "batch_size": 256,
+    "num_steps": 5000,
+    "lr_fb": 0.001,
+    "lam": 1.0,
+    "K": 8,
+    "ema_momentum": 0.995,
+    "sigma_bridge": 0.3,
+    "eval_every": 1000,
+    "seed": 42,
+    "gpu": 0,
+    "output_dir": "results/toy_lq",
+    "vnet_hidden": 256,
+    "vnet_layers": 3,
+    "term_grad_weight": 0.0,
+    "fm_weight": 0.0
+  },
+  "log": {
+    "steps": [
+      1,
+      1000,
+      2000,
+      3000,
+      4000,
+      5000
+    ],
+    "dfa_costate_cos": [
+      0.001022340264171362,
+      0.0024294707691296935,
+      0.000613357910575966,
+      0.002641987521201372,
+      0.0019003628791930776,
+      0.004648004221962765
+    ],
+    "state_costate_cos": [
+      0.009337988216429949,
+      0.9437081764141718,
+      0.9386202742656072,
+      0.9438605606555939,
+      0.9475679496924082,
+      0.9496726940075556
+    ],
+    "credit_costate_cos": [
+      0.023481125943362713,
+      0.0021192015459140143,
+      0.05908452393487096,
+      0.11570050132771333,
+      0.09644180163741112,
+      0.06304322431484859
+    ],
+    "dfa_rho": [
+      0.015879416760678094,
+      0.009290086299491426,
+      0.0009949249991526206,
+      -0.004670841522359599,
+      -0.0029721508423487344,
+      0.0010571565168599288
+    ],
+    "state_rho": [
+      0.0029325426245729127,
+      0.9371241927146912,
+      0.9219773809115092,
+      0.9298903445402781,
+      0.9345368842283884,
+      0.9332165767749151
+    ],
+    "credit_rho": [
+      0.02086908878603329,
+      -0.014479975526531538,
+      0.04267269264285763,
+      0.10674913817395766,
+      0.091526560485363,
+      0.04765695089008659
+    ],
+    "dfa_nudge": [
+      -0.0003799900102118651,
+      -0.0008909914953013262,
+      0.00031573620314399403,
+      -0.0008827850688248873,
+      -0.0003006396194299062,
+      -0.0016971436173965533
+    ],
+    "state_nudge": [
+      -0.002327537008871635,
+      -0.33619146794080734,
+      -0.3439306889971097,
+      -0.32351043323675793,
+      -0.33487510432799655,
+      -0.35304194688796997
+    ],
+    "credit_nudge": [
+      -0.007470574385176103,
+      0.002750888311614593,
+      -0.017584003585701186,
+      -0.037241545505821705,
+      -0.03454847944279512,
+      -0.02138534157226483
+    ],
+    "bridge_residual": [],
+    "state_bridge_loss": [
+      66.01814270019531,
+      2.143366813659668,
+      1.9674744606018066,
+      2.152421712875366,
+      2.0728020668029785,
+      2.0048255920410156
+    ],
+    "credit_bridge_loss": [
+      111.63633728027344,
+      0.35802197456359863,
+      0.10368431359529495,
+      0.12013768404722214,
+      0.07032017409801483,
+      0.051539346575737
+    ],
+    "term_loss": [
+      111.63633728027344,
+      0.2640027403831482,
+      0.05911973863840103,
+      0.06323631852865219,
+      0.03984691575169563,
+      0.03099522553384304
+    ],
+    "bridge_loss": [
+      3.0909704946679994e-06,
+      0.09401924908161163,
+      0.04456457495689392,
+      0.056901365518569946,
+      0.030473260208964348,
+      0.02054412104189396
+    ],
+    "term_grad_loss": [
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0
+    ],
+    "fm_loss": [
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0
+    ]
+  },
+  "final_per_layer": {
+    "dfa_costate_cos": [
+      -0.025100916624069214,
+      -0.025588383898139,
+      0.0002230944810435176,
+      -0.012649456970393658,
+      -0.003438096959143877,
+      -0.008626529946923256,
+      0.04746557027101517,
+      0.0495537593960762,
+      0.05357586592435837,
+      0.039518602192401886,
+      -0.06593793630599976,
+      0.0067804791033267975
+    ],
+    "state_costate_cos": [
+      0.94776451587677,
+      0.9486032724380493,
+      0.9489368200302124,
+      0.9489387273788452,
+      0.9493762850761414,
+      0.9498711228370667,
+      0.9505667686462402,
+      0.9509860277175903,
+      0.9503088593482971,
+      0.9503939151763916,
+      0.9501732587814331,
+      0.9501527547836304
+    ],
+    "credit_costate_cos": [
+      0.06872019171714783,
+      0.0653301253914833,
+      0.06267654895782471,
+      0.060872238129377365,
+      0.059321627020835876,
+      0.05934217572212219,
+      0.06014599651098251,
+      0.05938819795846939,
+      0.06065107509493828,
+      0.06321083009243011,
+      0.07017489522695541,
+      0.066684789955616
+    ],
+    "dfa_rho": [
+      -0.039396628737449646,
+      -0.062005721032619476,
+      0.005523890256881714,
+      0.014756813645362854,
+      -0.029910940676927567,
+      0.016373179852962494,
+      0.06027424708008766,
+      0.07274787873029709,
+      0.04506715014576912,
+      0.02043943479657173,
+      -0.07323547452688217,
+      -0.017947951331734657
+    ],
+    "state_rho": [
+      0.9350653886795044,
+      0.9352835416793823,
+      0.9357374310493469,
+      0.934054434299469,
+      0.9226828813552856,
+      0.930233359336853,
+      0.9396862983703613,
+      0.9324671030044556,
+      0.9295884966850281,
+      0.9347179532051086,
+      0.9344969987869263,
+      0.9345850348472595
+    ],
+    "credit_rho": [
+      0.07084883749485016,
+      0.0004139472730457783,
+      0.08371011912822723,
+      0.010933063924312592,
+      0.07074157148599625,
+      0.02977888286113739,
+      0.06011004000902176,
+      0.020937703549861908,
+      0.067134790122509,
+      0.04539356008172035,
+      0.06114630401134491,
+      0.050734590739011765
+    ],
+    "dfa_nudge": [
+      0.010068551637232304,
+      0.009215200319886208,
+      -0.002217007800936699,
+      0.007142472080886364,
+      0.0020765019580721855,
+      0.004165132530033588,
+      -0.019989464432001114,
+      -0.018945707008242607,
+      -0.01947549544274807,
+      -0.014030318707227707,
+      0.023684613406658173,
+      -0.0020602019503712654
+    ],
+    "state_nudge": [
+      -0.3548189401626587,
+      -0.3539973497390747,
+      -0.3535424768924713,
+      -0.35353928804397583,
+      -0.35366636514663696,
+      -0.35295385122299194,
+      -0.35282424092292786,
+      -0.35278239846229553,
+      -0.3528212308883667,
+      -0.3526480793952942,
+      -0.35135188698768616,
+      -0.35155725479125977
+    ],
+    "credit_nudge": [
+      -0.024559948593378067,
+      -0.023154649883508682,
+      -0.02181980386376381,
+      -0.020881079137325287,
+      -0.02010848931968212,
+      -0.019884146749973297,
+      -0.020017635077238083,
+      -0.01966693066060543,
+      -0.020179908722639084,
+      -0.02114756405353546,
+      -0.023488853126764297,
+      -0.021715089678764343
+    ]
+  }
+}
\ No newline at end of file
diff --git a/results/toy_lq/toy_lq_v2_seed456_lam0.1_sig0.1_tgw1.0_fm0.0.json b/results/toy_lq/toy_lq_v2_seed456_lam0.1_sig0.1_tgw1.0_fm0.0.json
new file mode 100644
index 0000000..8e7c7b0
--- /dev/null
+++ b/results/toy_lq/toy_lq_v2_seed456_lam0.1_sig0.1_tgw1.0_fm0.0.json
@@ -0,0 +1,330 @@
+{
+  "config": {
+    "d_hidden": 64,
+    "output_dim": 10,
+    "num_layers": 12,
+    "sigma": 0.03,
+    "batch_size": 256,
+    "num_steps": 8000,
+    "lr_fb": 0.001,
+    "lam": 0.1,
+    "K": 8,
+    "ema_momentum": 0.995,
+    "sigma_bridge": 0.1,
+    "eval_every": 1000,
+    "seed": 456,
+    "gpu": 0,
+    "output_dir": "results/toy_lq",
+    "vnet_hidden": 256,
+    "vnet_layers": 3,
+    "term_grad_weight": 1.0,
+    "fm_weight": 0.0
+  },
+  "log": {
+    "steps": [
+      1,
+      1000,
+      2000,
+      3000,
+      4000,
+      5000,
+      6000,
+      7000,
+      8000
+    ],
+    "dfa_costate_cos": [
+      -0.008305357536301017,
+      -0.0011448257913192113,
+      -0.011490714969113469,
+      -0.005118173896335065,
+      -0.0028971168600643673,
+      -0.007146042305976152,
+      -0.005135039333254099,
+      -0.006408803591815134,
+      0.0030692683843274913
+    ],
+    "state_costate_cos": [
+      0.010766413528472185,
+      0.9454400340716044,
+      0.9422353406747183,
+      0.945327232281367,
+      0.9411078443129858,
+      0.9500264773766199,
+      0.947187085946401,
+      0.9483370830615362,
+      0.9451590329408646
+    ],
+    "credit_costate_cos": [
+      0.010942678588132063,
+      0.8771380881468455,
+      0.9226242254177729,
+      0.9362819741169611,
+      0.9372822741667429,
+      0.9464249561230341,
+      0.9468309829632441,
+      0.9452949613332748,
+      0.9439582029978434
+    ],
+    "dfa_rho": [
+      -0.0028248391657446823,
+      -0.001316564545656244,
+      -0.00823917348558704,
+      -0.008288809360237792,
+      -0.009200356512640914,
+      -0.02464000484906137,
+      0.006656331475824118,
+      -0.014378171879798174,
+      0.00816792449525868
+    ],
+    "state_rho": [
+      0.02655455912463367,
+      0.9340365380048752,
+      0.930340642730395,
+      0.9328931520382563,
+      0.9261527856191,
+      0.9353549828131994,
+      0.9365303913752238,
+      0.930912658572197,
+      0.9325708796580633
+    ],
+    "credit_rho": [
+      0.015292729716748,
+      0.8341556191444397,
+      0.884218767285347,
+      0.9175956894954046,
+      0.9214809189240137,
+      0.9294924139976501,
+      0.9364803830782572,
+      0.9272431135177612,
+      0.9332626809676489
+    ],
+    "dfa_nudge": [
+      0.004384364855165283,
+      0.0014555706487347682,
+      0.005660574107120435,
+      0.003231095770994822,
+      0.0024609332904219627,
+      0.0033611954810718694,
+      0.0025722047624488673,
+      0.003350513521581888,
+      -0.00046457063096265
+    ],
+    "state_nudge": [
+      -0.005116054400180777,
+      -0.36389906456073123,
+      -0.37967536846796673,
+      -0.3594902977347374,
+      -0.3901909242073695,
+      -0.3562925284107526,
+      -0.34591514120499295,
+      -0.3550695503751437,
+      -0.3515613650282224
+    ],
+    "credit_nudge": [
+      -0.003232262640570601,
+      -0.34101804345846176,
+      -0.37139702836672467,
+      -0.355108546713988,
+      -0.38709117472171783,
+      -0.3536044582724571,
+      -0.34420712540547055,
+      -0.35209985077381134,
+      -0.3494671831528346
+    ],
+    "bridge_residual": [],
+    "state_bridge_loss": [
+      66.31975555419922,
+      2.3091068267822266,
+      2.0314407348632812,
+      2.110802412033081,
+      1.9804980754852295,
+      1.8128150701522827,
+      2.0147881507873535,
+      2.169416904449463,
+      2.0117833614349365
+    ],
+    "credit_bridge_loss": [
+      158.73072814941406,
+      11.18570613861084,
+      9.38658332824707,
+      9.429727554321289,
+      10.842954635620117,
+      10.344818115234375,
+      10.250753402709961,
+      10.136574745178223,
+      9.871820449829102
+    ],
+    "term_loss": [
+      132.93673706054688,
+      4.870186805725098,
+      3.9316887855529785,
+      4.335302829742432,
+      5.871437072753906,
+      4.5994768142700195,
+      5.077899932861328,
+      4.8271284103393555,
+      4.638625621795654
+    ],
+    "bridge_loss": [
+      7.166463547036983e-07,
+      0.3530547022819519,
+      0.20790576934814453,
+      0.15763649344444275,
+      0.18163490295410156,
+      0.154127299785614,
+      0.1578863561153412,
+      0.13594242930412292,
+      0.12283627688884735
+    ],
+    "term_grad_loss": [
+      25.793991088867188,
+      5.962464809417725,
+      5.246988296508789,
+      4.936788082122803,
+      4.789882183074951,
+      5.591214656829834,
+      5.0149664878845215,
+      5.173503875732422,
+      5.110358238220215
+    ],
+    "fm_loss": [
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0
+    ]
+  },
+  "final_per_layer": {
+    "dfa_costate_cos": [
+      0.0057477825321257114,
+      -0.05385865271091461,
+      0.031865090131759644,
+      -0.0450170561671257,
+      0.01222146674990654,
+      0.048112209886312485,
+      -0.0029807062819600105,
+      0.02491983398795128,
+      -0.03023526817560196,
+      -0.0383906215429306,
+      0.005803969223052263,
+      0.07864317297935486
+    ],
+    "state_costate_cos": [
+      0.9431054592132568,
+      0.9437580108642578,
+      0.9442368149757385,
+      0.9451977014541626,
+      0.9453786611557007,
+      0.9453213214874268,
+      0.9453005790710449,
+      0.945610761642456,
+      0.9459618330001831,
+      0.946029007434845,
+      0.9460264444351196,
+      0.9459818005561829
+    ],
+    "credit_costate_cos": [
+      0.9404189586639404,
+      0.9415367841720581,
+      0.9419513940811157,
+      0.9428710341453552,
+      0.943200409412384,
+      0.9437626004219055,
+      0.9443795680999756,
+      0.9448325037956238,
+      0.945106029510498,
+      0.9459753036499023,
+      0.9465078115463257,
+      0.9469560384750366
+    ],
+    "dfa_rho": [
+      0.04672882705926895,
+      -0.030157513916492462,
+      0.03139471262693405,
+      -0.02426629513502121,
+      -0.003147948533296585,
+      0.06672847270965576,
+      -0.0028628872241824865,
+      -0.0006099608726799488,
+      -0.06767985224723816,
+      -0.07952223718166351,
+      0.023338939994573593,
+      0.13807083666324615
+    ],
+    "state_rho": [
+      0.9322999715805054,
+      0.9332318902015686,
+      0.933842658996582,
+      0.9367775917053223,
+      0.9355518221855164,
+      0.9306066036224365,
+      0.9324768781661987,
+      0.9316605925559998,
+      0.9233090281486511,
+      0.9335967302322388,
+      0.9341922998428345,
+      0.933304488658905
+    ],
+    "credit_rho": [
+      0.929013192653656,
+      0.925915002822876,
+      0.9249007701873779,
+      0.9293269515037537,
+      0.9353340268135071,
+      0.9355732202529907,
+      0.9336704611778259,
+      0.9395359754562378,
+      0.934350848197937,
+      0.9350711107254028,
+      0.9423503279685974,
+      0.9341102838516235
+    ],
+    "dfa_nudge": [
+      -0.0025230227038264275,
+      0.01772259920835495,
+      -0.011610012501478195,
+      0.01727830246090889,
+      -0.002929902635514736,
+      -0.01679525338113308,
+      0.0033558662980794907,
+      -0.007272847928106785,
+      0.013050400651991367,
+      0.014434966258704662,
+      -0.0013312064111232758,
+      -0.02895473688840866
+    ],
+    "state_nudge": [
+      -0.3518102467060089,
+      -0.3520665764808655,
+      -0.35131847858428955,
+      -0.35165226459503174,
+      -0.35131698846817017,
+      -0.35155874490737915,
+      -0.3516783118247986,
+      -0.3522875905036926,
+      -0.35148167610168457,
+      -0.3511171340942383,
+      -0.35132837295532227,
+      -0.3511199951171875
+    ],
+    "credit_nudge": [
+      -0.34800052642822266,
+      -0.3485172390937805,
+      -0.3479737639427185,
+      -0.3486989140510559,
+      -0.3486787676811218,
+      -0.3492557406425476,
+      -0.3497552275657654,
+      -0.35073035955429077,
+      -0.35016047954559326,
+      -0.3501723110675812,
+      -0.3507159650325775,
+      -0.35094690322875977
+    ]
+  }
+}
\ No newline at end of file
diff --git a/results/toy_lq/value_net_seed42.pt b/results/toy_lq/value_net_seed42.pt
new file mode 100644
index 0000000..0cb6683
--- /dev/null
+++ b/results/toy_lq/value_net_seed42.pt