From c77633c38cadb8e99628d5de84d31c1eb06dd757 Mon Sep 17 00:00:00 2001
From: Simon Alibert <75076266+aliberts@users.noreply.github.com>
Date: Sat, 4 May 2024 16:20:30 +0200
Subject: [PATCH] Add regression tests (#119)

- Add `tests/scripts/save_policy_to_safetensor.py` to generate test artifacts
- Add `test_backward_compatibility to test generated outputs from the policies against artifacts
---
 .../common/policies/tdmpc/modeling_tdmpc.py   |   3 +-
 lerobot/configs/policy/tdmpc.yaml             |   1 +
 lerobot/scripts/train.py                      |  86 ++++++++-------
 .../aloha_act/actions.safetensors             | Bin 0 -> 5104 bytes
 .../aloha_act/grad_stats.safetensors          | Bin 0 -> 31688 bytes
 .../aloha_act/output_dict.safetensors         | Bin 0 -> 196 bytes
 .../aloha_act/param_stats.safetensors         | Bin 0 -> 33408 bytes
 .../pusht_diffusion/actions.safetensors       | Bin 0 -> 4600 bytes
 .../pusht_diffusion/grad_stats.safetensors    | Bin 0 -> 47424 bytes
 .../pusht_diffusion/output_dict.safetensors   | Bin 0 -> 68 bytes
 .../pusht_diffusion/param_stats.safetensors   | Bin 0 -> 98776 bytes
 tests/scripts/save_policy_to_safetensor.py    | 101 ++++++++++++++++++
 tests/test_datasets.py                        |   2 +-
 tests/test_policies.py                        |  39 ++++++-
 tests/utils.py                                |  47 ++++++++
 15 files changed, 236 insertions(+), 43 deletions(-)
 create mode 100644 tests/data/save_policy_to_safetensors/aloha_act/actions.safetensors
 create mode 100644 tests/data/save_policy_to_safetensors/aloha_act/grad_stats.safetensors
 create mode 100644 tests/data/save_policy_to_safetensors/aloha_act/output_dict.safetensors
 create mode 100644 tests/data/save_policy_to_safetensors/aloha_act/param_stats.safetensors
 create mode 100644 tests/data/save_policy_to_safetensors/pusht_diffusion/actions.safetensors
 create mode 100644 tests/data/save_policy_to_safetensors/pusht_diffusion/grad_stats.safetensors
 create mode 100644 tests/data/save_policy_to_safetensors/pusht_diffusion/output_dict.safetensors
 create mode 100644 tests/data/save_policy_to_safetensors/pusht_diffusion/param_stats.safetensors
 create mode 100644 tests/scripts/save_policy_to_safetensor.py

diff --git a/lerobot/common/policies/tdmpc/modeling_tdmpc.py b/lerobot/common/policies/tdmpc/modeling_tdmpc.py
index eab0f94e..1fba43d0 100644
--- a/lerobot/common/policies/tdmpc/modeling_tdmpc.py
+++ b/lerobot/common/policies/tdmpc/modeling_tdmpc.py
@@ -80,7 +80,8 @@ class TDMPCPolicy(nn.Module, PyTorchModelHubMixin):
         self.config = config
         self.model = TDMPCTOLD(config)
         self.model_target = deepcopy(self.model)
-        self.model_target.eval()
+        for param in self.model_target.parameters():
+            param.requires_grad = False
 
         if config.input_normalization_modes is not None:
             self.normalize_inputs = Normalize(
diff --git a/lerobot/configs/policy/tdmpc.yaml b/lerobot/configs/policy/tdmpc.yaml
index 71dfa9c9..eb89033b 100644
--- a/lerobot/configs/policy/tdmpc.yaml
+++ b/lerobot/configs/policy/tdmpc.yaml
@@ -1,6 +1,7 @@
 # @package _global_
 
 seed: 1
+dataset_repo_id: lerobot/xarm_lift_medium_replay
 
 training:
   offline_steps: 25000
diff --git a/lerobot/scripts/train.py b/lerobot/scripts/train.py
index 268185a3..f58dbd06 100644
--- a/lerobot/scripts/train.py
+++ b/lerobot/scripts/train.py
@@ -25,6 +25,51 @@ from lerobot.common.utils.utils import (
 from lerobot.scripts.eval import eval_policy
 
 
+def make_optimizer_and_scheduler(cfg, policy):
+    if cfg.policy.name == "act":
+        optimizer_params_dicts = [
+            {
+                "params": [
+                    p
+                    for n, p in policy.named_parameters()
+                    if not n.startswith("backbone") and p.requires_grad
+                ]
+            },
+            {
+                "params": [
+                    p for n, p in policy.named_parameters() if n.startswith("backbone") and p.requires_grad
+                ],
+                "lr": cfg.training.lr_backbone,
+            },
+        ]
+        optimizer = torch.optim.AdamW(
+            optimizer_params_dicts, lr=cfg.training.lr, weight_decay=cfg.training.weight_decay
+        )
+        lr_scheduler = None
+    elif cfg.policy.name == "diffusion":
+        optimizer = torch.optim.Adam(
+            policy.diffusion.parameters(),
+            cfg.training.lr,
+            cfg.training.adam_betas,
+            cfg.training.adam_eps,
+            cfg.training.adam_weight_decay,
+        )
+        assert cfg.training.online_steps == 0, "Diffusion Policy does not handle online training."
+        lr_scheduler = get_scheduler(
+            cfg.training.lr_scheduler,
+            optimizer=optimizer,
+            num_warmup_steps=cfg.training.lr_warmup_steps,
+            num_training_steps=cfg.training.offline_steps,
+        )
+    elif policy.name == "tdmpc":
+        optimizer = torch.optim.Adam(policy.parameters(), cfg.training.lr)
+        lr_scheduler = None
+    else:
+        raise NotImplementedError()
+
+    return optimizer, lr_scheduler
+
+
 def update_policy(policy, batch, optimizer, grad_clip_norm, lr_scheduler=None):
     start_time = time.time()
     policy.train()
@@ -276,46 +321,7 @@ def train(cfg: dict, out_dir=None, job_name=None):
 
     # Create optimizer and scheduler
     # Temporary hack to move optimizer out of policy
-    if cfg.policy.name == "act":
-        optimizer_params_dicts = [
-            {
-                "params": [
-                    p
-                    for n, p in policy.named_parameters()
-                    if not n.startswith("backbone") and p.requires_grad
-                ]
-            },
-            {
-                "params": [
-                    p for n, p in policy.named_parameters() if n.startswith("backbone") and p.requires_grad
-                ],
-                "lr": cfg.training.lr_backbone,
-            },
-        ]
-        optimizer = torch.optim.AdamW(
-            optimizer_params_dicts, lr=cfg.training.lr, weight_decay=cfg.training.weight_decay
-        )
-        lr_scheduler = None
-    elif cfg.policy.name == "diffusion":
-        optimizer = torch.optim.Adam(
-            policy.diffusion.parameters(),
-            cfg.training.lr,
-            cfg.training.adam_betas,
-            cfg.training.adam_eps,
-            cfg.training.adam_weight_decay,
-        )
-        assert cfg.training.online_steps == 0, "Diffusion Policy does not handle online training."
-        lr_scheduler = get_scheduler(
-            cfg.training.lr_scheduler,
-            optimizer=optimizer,
-            num_warmup_steps=cfg.training.lr_warmup_steps,
-            num_training_steps=cfg.training.offline_steps,
-        )
-    elif policy.name == "tdmpc":
-        optimizer = torch.optim.Adam(policy.parameters(), cfg.training.lr)
-        lr_scheduler = None
-    else:
-        raise NotImplementedError()
+    optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
 
     num_learnable_params = sum(p.numel() for p in policy.parameters() if p.requires_grad)
     num_total_params = sum(p.numel() for p in policy.parameters())
diff --git a/tests/data/save_policy_to_safetensors/aloha_act/actions.safetensors b/tests/data/save_policy_to_safetensors/aloha_act/actions.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..70c9b6d811a396f7e554b424beb0b0ced73048d0
GIT binary patch
literal 5104
zcmb7{`CHBD8^=#a!(=RJFv+MSVpLB1!gD{?L{ZXuI;|6(>XfCJP^pY%Letj_WyX+Y
zlD!Nj&gXtgS}Y+Wv`O1gX+71c(@DO^{0E-jKF@W%uIstp_ve1U?_130{daa52n_6Z
z8HB{_*u25O-oU}a+`!l%S{(H0%h%S}%+k-;AS5UzC}7jZjnNxoqE)v7V@pe0zj$K<
zv(F_}oUyH)^(WT+b6FK<Y-V9;`3bi8T=)}cY-VF;{)x8yY+3~xo10q-KG9a6OMe25
z&8=<!lg#?FX%%Q}VP^54Xf~fqe*%pytZe^N&Gxft6=-ZJu=r0kyU(RRfuH^qm1^<&
z`WA7fsJ5;O`G1S4!}J+AHnsyUDgsdFW&y0XabQ=@Y+=94y(w?aCT@6<64Fv+#Nb>t
zY5kuWDjvykX{(rSQs`m#r7jpbCBi>wF(jC-VplJQvAQG`H#LbHvw~0Cmv)gu&x*+}
z=gjfkuaB|8HHv1xn2A|BJz&%8gBIZhaPhJ;^IN@>osRdTg2GLlqlX&hjqW7XYwL(z
zfhE%MwP;uxNuLGk;koiI$ki8OMNI+Z|K`NDcg8awUql77^*QM>1q|lCCMU9M$^OGO
zI9K`{^B{t@e$dB#|MY|V8!rqH6o7M=3)3h}WH(HMY4@o=x%MCh)b6h!gUaV*jEN0i
z(r!V!SK;(Y_6$^R>xZfVZwwA#kPYitMo%2GpA$s2Yc6pu*9W0}dKoEitRl^`7h#6q
zdmJ2=(9@0D7?z@dIq!s6c(54m6s}|aE<dum_&}OBrGv}4AcxxI0W!7XEh$>F5W^hX
z@vvGXo&4`q<hd(gBJ;(s((l5|3|DrlE17+h8$d(m@-eHqAAX$KLJE0&cyntWx;$z`
z&*o6d?&#vtqusD{XCUfl7s8SY4$MpX1M^mT(X@CkZb8upP(LzEIv3ZIu#F~osN*B1
zw1!e$YXfv}l)<TJKdkO7hKkTtY}C!oY^$3W4djbCvHvI<QrJ&4M?E6TnoQBNr54*<
zBWa430Tyca0e3}+hXx8@xy*^}{vnag-7BJ+XBTmSKk#XN<^bVzTF6w989GF@p?8yn
z4%h2pORfwy8VK>si@T6z?8MZJl9{MmNEfH8arVcQ@TKS-Nw9cF4tZGM!d-80YjHTu
zduf0<w-jKX=7k!g^WkOTS~hddK4zj3M0w;OXRJ2_nN4*>tGk9oXj!0(PA`7r7*1<1
z>0yb#0vc5HvpmMZc70ct9J_~&@d%`I=cRFr?1msi@RGE=X(TUvY%$@k0yWo5sHw9K
zuCn?7BS9iOuU`lntuD;&mjmp7Cj;n#o+n(9=OC!>8zI%gAu{Jr8+^|c_;#~|o^;m6
zvEx3#;v5mKu3->Zxt48>+|TYm@~2<EY3F_&Hwc<Lx=Gt^4G5lRf(I*JB6^5vjX)p2
zZ|{T%r62a#mO|gvr7Z2|F!o@)4^8}j6zAS4hi}LA6On5HIkH`Vf3^3b#~v|d`Ll50
z-ENS)7NWZIL$Iu0!S2rrX0Nw<Q-R%7Zf7u$9-7@v0;?~R=m+LF$?Ykg|2dMHPB6qv
zt7Pzgo;SYtx(B~yJF@VYt;}<|A6<O4Q98PP1m?=RNmS!~!Y*1O-=_{6jUuUG{Y=bW
z)D6n%-gszs5!`QC$<n56XNxL*X<~?nv|z3rw0#=L@~nqsj+-@_3_Qp1U&3iwv;khZ
zD~Gdxc;KDDd+_&9PHg?7oh)o}FwJ?ig1daCA4G?$iEhJ95;SCmr?<CYqHQ=GcV7<)
zZvbo_c;Y1IyYSm7XSV+4cJ_2(5Y1V!iqjgc>S>o}q;4*f=JJKOV164;UnZfMBidM+
zs|2$dKDdQf3XK<>*%fIbi@O#`544=(PQ>(r;nzws_TCfn^qoDX-DtzDE)rV!W-5O3
zy%LN{d@$zZJ&659_0Au6GmCoxH0|9TPPU~N>Z99<@woRyZJhu$$2H@))1kB}QWu3W
zJ@9^90ID_J1wP+_)jyFi=WK5p*z3lbTpNbI<K5)k_9vw5e-^l;fyW%PLg^95nHXi*
z3%<{MQLV8E7F9bkndb&prRz;ID}uT9H+*^`y`97!=ZK=+68~Obg{|6=RFI*MS;9W>
zp5ucKXY%1{x-)B-xQ*#3L{vXk$gPj$(*mCklB!-q>Km+4qoM`5-zC)FT^IjqmcfY3
z8xN>_u&&UF-8%dO+xe?6%@ONx>RL)TU{XhxAFd=zP3_QkY#sh+5l)v5=;Jfd5OmG<
z#F>c*dG;=>WKSFm`Zb7d{`W99RX7BJ<OK;7KO!9;ZL#!>3{78!QF|?2ToOA3+PPl%
z*hC7;|8!+`F0m}ID3EU6aEv=vqX1EO8Cm-LA2PfLaQy3jOne+cwc@|Qm19O=<`N<5
zeSHtMrn)kxWywrtA4n^|e9o2M8-zpO%1Neu1KC_auslPK``<=THhl{AbbNq$=|W8T
zz+m8l3u}Cm$a?$(sP~pGZie9?)CRPW=A;plYB3+v%$o6%OiVX-&qD3WUNHYV02i=3
z@Xzk0tR^a!Et2@q0_k_$&n_RpdwDncE$l7{J#C7;i+SveX<@YU$sBy<+Yi4J5eD}#
zIFq@8`MPap!$saSaTjoTi$+n)m^Sh<>nhnd#te^F)Zzk-C@L#g_21|LaJLm=^>7|^
zPh7=r)g-Xn&i*vV-jMTj=2OA)R$^IGMtr7Q;OUZfTzNc_#?GIGC++&-y_yh%cz577
zab(eUNv!;&hz43+m+Eg&Lc!}Qa#{9}81Jz{XSZft`*j3$Cvy<IhM_9h3$^|7AaIux
zGn%r8?duMviJPOjywgLF?em<J+cHv}VuhMUgSbT+P7m&yg|hR*a4tvX6P>rfZ1x(q
z*F1?`{5F`zY9w;O(}zGm@d2^5DJHv1?D6TvVN`TVs2%q$PW0o^FHL=M$)FTs$E;!3
zChcQ!;vlO3Fqae08HA88`bhJ(=R~vG4mC~;<Msv#b-JR1A_pG*`H?T~%FF}p_f9N5
zXfHduN9ETK9&;C04nj`&TN1)`lP`7|;})yesQWUMR+UV{^ZVOj(zsx}BPs^vofWKj
z+Xf~-;7Pq9nF~%-`Spd5gzsNP_<I+ie7pkPCX1>0Azj?%(g_=1_~S|A`>=PdBU>>g
zfJG`jsD9Z#?pPU*w%i{gL+VAuyGwx1k<W2Z6iIET>0xG4H`G4x#jB$W;re4Imh2tL
z#J7cXa5#yp9M7Yd)s>{Iy^(mCSm3XZ8!%p#f2}vZ#Wn6-pg7`-9WjLvJ>H2~MQ&!&
zO)8)NXu}Dj2B0sXo$L&$Ba#JHI6=1tk8TL32ac=qT}uuvAHC6hK>=jMuVr6-wT(?Z
z7eMtV9_IpP%Rw{xEy*^iAZ6*6I26-~3k2a*uy;Bd&5(mUP>44+Fc^96!u}fA%sMLl
zY3$Uq+=mPOP^4E!E-EU?d8HjTRP^Ge{VM;gnu;aW3dp}H!p4{)@cYSy)h~}@V!0pv
zGQWYN4`g6CP908ec}KJ-*x~k(UYxa0LZ9SK#!dMOQ0(x-Wo-rEwq-5bGkX{F|JIMr
z-L2#<Pw549jUjS2R}C7b&&PkOYw*>yP-^|tH~7S~13rxV9+Se0AS`4h>+cY=_E=B4
z_?a_za)uIi_VMAz^)<xfttqBHRv^zSl*a7U!Fw^Ca6nHr_oSA>lo_kp-wT6S$W0IG
z+7`%FJXNEUV@61ddI@O@7vSDm71;AwLe<lBu*s?$+OkAwB`XB;kIpQeZeo`?A=Ny$
zh;t55qeZLKAS$boG(R=RH^K(IyjDUTOSSRc`%d_{MTA|6MPOy(%ybq;vnEj0FMU9o
zdRze^#=WHSRRcMgZHZ_5UZI*<IDLCy8YZXr!?{`?%$#x;R`Xq%#+V%}Co6!K*YD%*
z4-G=#+g9RU{+x_Avcy#H_h^e@^vLb0c&4u(=8p44Lnj8;jn=WGM=@-u$e%i@9p(l$
z4yy9Do=AFMlQfY%zSQc$nz0hvW;qFM#0nTG@x#{}ilHRRm5Drev+gbaG+R{0T|d?j
zA)V^LCo;kp+o99=9?Tvkp@tz7am7XjJl^S#vW~k@eb9yJwZ*fP{i?dpZsWv>{h&Fb
zBnG;p;jpJEM*Z82=A%RD;zP4A>O?PGuk^=0-9m5+U&bu2ZDr|3UbOyphqUL$5V$Lp
zq`nD>ph|#CT0dfhSt!koFv70aeQ?%Ym49I+u<36Hwtra!d&N9y&K3=BSiq-c^9RYX
z(2L~VHgkL!UxnB2NT~jdd1#s`hlyG~NKWL#(V|tXw<V5g{On5?yS|lfouWo<o8%<t
zw_@THW{GCB1t)n)XlnBuJS6P{_lMpnkru!c|5YsXXaf5xNl1CgK~lqM3WzglC6{*=
zlGYMyj9Ax%FVn;5KV#=&ZHp3YoITMqEFYq`IWx<*iHx@|kQTJLa~4(u5aUrtN|N)4
zg^M+A)E>a8%fjfsWJ8S0QG#BC7fyJ`fhn$GJ8s6aJAVYw#6A98`FuGHn!X}(YesZB
zRP)R55H^`cP}dCAo^$RaB;@#D>R*L$;{6(Co4beIdf-pRcmL#m3+{u}uXu1@T1G})
zx5x4+Lzpx#g1%ia0}l&7!X1$iAMVJ9xPE7Lv>}=4Jo2My6AHMAnSD?e*Gq)mZN&A$
z984~&#eMam6#ku!7az+YEh7N8W!!@|xP-YR#IaNpRsJ<JO0S(&0zZxqZXU%XD*rny
zema0-z7x~(jo;xvYx`hkf(RWAO5sS=a`tRg3<HsBe$g1iC2Z!?H0dy@dw7*x9Gs74
zgD>#%poBJd%){<0a@c!Yh`9xKLGpAZbNi6QvV;9-=EPE|)@(j)oW_H!jt8Xcu_@+f
zG~tg2B~;*NgpaQF!7f$)jW@gpX4Wg2`S?9-+8bZWJ2O)nzex##H9drz@{kNi3NYbX
z6>fuY8o{53BMnMux#opMZUrz>yPB0x+Rt)qgK4I}1$Rz51QnxNh&f>-`RxKc`a=ia
zR`uWYirM&kz7oo{eQ-xOgHa-9Ha>hWs|gLF^&Lw%H=7}t{HBgrJKrPKEmr7sU4~6W
zLPy!^q1KC!&~(ukf2a3A7Us-CU!^c{YM|<!nVd=T0Jvt1g2d_7<l1U0JX0#eH6WoW
zi)LWa!;f%U?u*(J3Sja|Cw7#V!bV>SpqgWlGrTbX1)&O}Z_`g^TFu2BwQsOnE~ays
z4!XCtL0WJiE=a!*j;B|!TSj55R8_yc+`U{yjshyT4U*({4~eI@F;3JU!Amp3=!s>z
z7;DoGO$YsOqR%5p-Ra0Cei6(v6TN9=_aQE8p&Hdp>LHS-8zkk0DXw#<!bK)gRJ~#b
zTHo$~8cQKga4LbE8Yfm)yoot9_)*@4ZJcNckH%`ti2nJ<L~dw?(tvlk9wX@_!|C|p
zk9If-LTve`1YG`H%_1L0vq}%u{%-XhcVzbf)Jq#k(a|R)ca;S;+^)f~e?-t?*>pq)
z861&%V(Ez@c(Glzzt8%CS>Ft%_b+5~?5Z5{vR)C*gaX1mEzs{p7fv#ZpvGTLLt9l=
z)%keib@M{F_P~{mv~FeZYlEm@{8g?XTMqiCo{;P3OUbDOTdY1lfDL0K=}#$NqhGun
zv~qm$snLDt6T7nfh6H9XIf!Ot)pEN$d*O2bAQ7&tA)*5|7#%)<0%O%&UN{LobLCKN
kBf|FMMPRMv!n8NUv#@!ARJ=sN9Zc<o@>gx-^7?-Af7^_iasU7T

literal 0
HcmV?d00001

diff --git a/tests/data/save_policy_to_safetensors/aloha_act/grad_stats.safetensors b/tests/data/save_policy_to_safetensors/aloha_act/grad_stats.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..2e8451891e477a10a94c36901734822c69470a04
GIT binary patch
literal 31688
zcmb`Qdvp}l9miLJBWi7xhXtQB#TN`|c4l{Wc4o5+9K^#>G!-~1D%*uD63BxlD-Xdo
zAc~c$D7Ju#MnElCut-8S3BF3SrL|ho8ho^+$EZ+WRq#}cw*3ytvUhjp-r3(>_{Rgs
z^ZopO^SgKEvy;irm-Bo_{wyk~sjEy@mnLSUs_JUvvyzF*(&<%+hImafQCngzDyd8@
zoSQ5$OU7$ziBi%qD`Efb+Uu2)%0wy=ubVluA(?8xzmzF~>u*+gX$`4LN?1TK+{5N4
zt8SQ;A{S>UepMJqtk171emxsFJuzd>^t#$)>5RJCdHzz#BK^9emU)E|3suV$m6^-Y
zLaP%CllA`6vf?UM1B%KF$|Idxx}j*y(A>jS*@LaAism-KoC9+-omjf2D*?Aj_J;e5
zJL&;!HNi~t(ut)90!lzI(w^|6$Ax;52pCFGFw?wrV(CFu335lO7Q5Dix}tMu>PRP+
zu9qn~cchtc>7%X(6rDR$M>@51LkV$5YKK?W&97}p)Xc3;;xCe?#*nIngrjwX5E}{U
ziXj+q-nOEye4|V;1QX6nCzfvDdMOxi-nJr_ZYW{Fg!9sgrH54|%$=-O?4}Ltioc9I
zR7X6qc0?ndJJn3Q^u9$l^0|X`<de&X7O`HGdx^E<-L@5V`2YcoaK>&JYBR7WzpFkS
zx#()8c{_~U1x(7%S1Ct6wS1_QpYKwcc=@HIV%B_}a^zFXhfewVJ{2f-A0u1w^Oeex
zPc0ux<>xz9CSLyN!>aszt#ag(%g1o}US-F-?J)8dfMEIghQ-5Bn}KNg`KA>pdLzTN
z{CwGR<WtLsZ29@Nm5G;MN;+oE7cNIWwS1VCpH;2O<P6*&s4qQo7Hz;vs;;U{Ch*Zw
z>LvlA6+dfP;&^KDkSsrIS=@9<4MjKr%mvvCKyCoTWnGJtuSy<$$d;eAEpa@xc*vH@
z$`&_W^5`QR07kdG0K^86jeOB`Yi_JgRoS=5;)ztMwzR4?KDWMZc3f&JfTXCbeL1=j
z(WtG#!Bp143Wdpz%D$Jy+8B{XZO^{V#mbo%uI|Rvk<rQRBcOmIER<Xvxj8HpR^1A2
zPIh>ZIhB<;Dv;VB<W6Pf&cUm#tFMu~;vg8|qLT&_PAwc-r?OhdDO`3AxLc;OS|^UD
z77v?KS(|fBm%P#-SW#Jt6NOU?hrp?<z;Uu=k2(ZSWd%+gPc9xaW%bQ9U6VX7LEcoM
zyxD8LH=J5H%uN-VTTiy^QHQ#zLUqfDrxp)yQ-$8vkgT40EH3H>#w!#!4v5+iOimS=
zoRcFv8EBm*)H*JT*c2!riB8p7QtlO%`bicPkVF@79*Nu%euyQMJQtAUpuho(xzFY}
zweVTgwxD;Ko8H-*RPNk`I7GS(KO0Fc5|*a9S=y0A*@KMxXquahv7yvL?Vm-u_*Q&w
zT|+!sGd)=;`(qQNDw>N~Ig_a+!>Tmas*<(G@-_jY5zp*TaXh(r%#igeZo2Fufm&&-
zT6qDe4ZyB6)~-1DlJbUV#4}q_98WDCex<Q~#Z8wz`nW@;v5Msdpf&)@(pbxK<lE;0
z$&V`tRy0<#MB&uJ;aD2$Se$Ix;||HvSjiH{lZ(exS+jCYmtFDDDUH=BQ8=}5_)~!O
zCr-BHQAaf5MVr-LjYn@c6N^VO(nTlYtWy+qHw+LZz>1RNC~-j4h9F4+R+0+f$W8{5
z6ksK(a1^yExKn_2CjqHAaRvpX!q!7BiQE#_1M69Zv?Mzy_*a1SF9M6&mi=THD`a^n
ze{nxvK{Vp|Cxg6rYVr27VFD}5rppcjH@pG@F>?T@4ZzC+0x!$tOI}3~tpo&GmKRSg
z9#$3*SXnk*_UJ>*0s=8}0LTqsx&klrmloko2u3*b;n%UTm~d+0aI%2F$ufn@9(QP2
zK%ixL@zmmBWdVVe^`y(Lc--y^2&610oLV@XEFf^QOt$P%hn58eT9y}2EFQ^77oCVX
zIzzpa1qeX=3$nkI10pwsCwB#Y=Ejkn3@|{6wz9LOUKF({7+O$ZXlzrtHbWj03@s=y
zG(L&i5(F(M5Hwei?4WQjEGTd^5sTUu6fG!FG_^QyLbMVTXjxu7wRl)rP+(=*blE{b
z%z^?ja{$N<V9)|D%j8R5MG%d6{)W$srxp(@3ks|(n=X6wA!b2=m^lE{2H<5uftRVp
zcoTw^pg_uU!l{MB$$|nW%Vf(QcW7Bqpk;aS)Z$@fL4lR^q|2^&NLf%IWjW#G!ZA^S
zlV!4Hk2<t0DA2OJcxv&mGF@P0Do<yq7i2&H;tLuVL~IBQaQJ@a#*v&1Fp4fvGy#g*
z6bwxl7#d3|&YM6{bb+DqNz|4gXu3eqTtTvff}`mIM-#EgZDCCa6iqA6n-Gn7{*^l~
zo?1MtOcz*LHeGfQ5HnpMW)1+g0eG1%@Ul$4<W&UGiZ0Nyym)Hyurgg>W!ZGuqYp9D
z1!Cp^P#b`k=>jj)it#1{Bb<L;&j}|Nj=2h)ER!vJ+@WQ<K+E#tsl~&}bb*!iq|2^&
zNSQ8>vYc>g;czlt;AEL>*`tnIWx7Di^5Ut*!^(7lm1)P!n;>W*fuM0g#D>5Chwo=@
z9LdQ51C(ehAFKOKpeP}Mp|MRB=S`r1B)TwXd&(tI=LA6u2?WixCD}p2(Lw@86S1gm
zLD51kiWaY`xgkC?nMlQ}YUdX9I|N8qLM~>;Cs12}n}uB5%oQLz5g1y?#n40yY8!C0
zkc*?ms}reY?J<}G94+MHXuN!K`53c{qq)S(&H$bka`7}if!YE*E#%^9@%p;yb*Xrw
zwlae%`guiQfO4Mo3qVnug1?1a{H>SNkizOJevjkkTF6D@gb->Yu(^<n&G8trs|;Ef
za?v^u47C~fU5NF&d5L5^?<cuM{st=Ifncuo-UxCNm^y2ELW1l(AbcTK_zHng8-eQ?
ztm_FFl0uJoAeasBjUYCGbPx)gw<`ObH2{jiik+~H-jP!qfzlbQ(z!C2ryQuaL=mnS
ztjc)-s0~2h3|8L+`I5T=@j!SDyXP25P666xu-e9LLUu<$+YDCQya3cGK-mmd*#!Br
zQ-HD=tg?9{s7*lH3|87)3E7W0MLn7dkYcb3M}birgUlJM%oWIzofOp0VAYO>qc#V-
zGg!OxpyhrrC7;e<cm`{DBpSIjtP$4q3WdoI4%TO|)<@(~+k^cXto^y*GLHvJ{^A4T
zfe@QO9)j8g%+K&Jzbrv^Fu2WUc*vgvLTv=@XLz_@h9P?ug8GF$)GrS~Yy#;Z6gF>Z
zF9QG&u@7Y1=p8w^5qv$Khx_$BOencC5DoxyFVA)JG2%yU0Nxk&@V-m~lDh)&KzI$i
z7%xLQ2ws!636eVk;Y!#;{4&ps9{D~P>J;F8VGr-i<jYP0>KFD<zdQuB37B8l!~EE%
z0*mr8fB^FvIY$D?jbZv84(QgH?4+Q9VGj)y;Hb^P0>d5_$kK}QGLS&?Iu~TemP@0y
z1_=y%NT6$LvV+4Nzp#e~ig?uapn+iz4a_`#D0}sR`h`8zFAqU&0_GR?FuyE8b}*2?
zu!sCPAmm0cd=K}_Fl4VnP`|K;`sE?0O~Cx{y|&DE=DO%c79lwfn_y$arvV3s*ozIY
zDdKBH&l8K1+rv%=!o&Y^5acG{e;|kjFhh`@2PDv6=3#+(5OO21Krp;!UWBCC={BN}
zf%on}dK31%Jw&kQKuLdWvGZ+6AtOHuAvXdW#1~9^ZS?GgxF}h(NjA=F_#7O$Ik+Kc
zUhBX~liokLN$4;0Fhusj&dky51(F+tA%f?%Lb46!Y83TmnvDX_LlW8ka(U$TAc?>g
zzFLmrqz4F71hMcHlTW0#Xs;L#SIj&EDSdT<AOa#@CV3EYBQQjKOTG9+l4VE_2$~3t
zcu6=g<YwTBpoo`2h9W6(c0M2;rkDpIHv&@xqp-;zy@=4y3$`vog0inIa?!<XE0SkB
z0Bnla?3sID<VKKgBfJLBAgB{TI0zo5*z*{u{41n-Ubf&Ni+K=oBepB%ylcUA=EZm%
z(rrW`1DEtTgxiEd<}JTFkPnE5EEf0VuN`kQyf%8)1Ty9!^{v%52{;dJWS_*$!I77m
zZH+l^W^i3<&gNvVGGN(2uVujn(i^l_h=(__hx4d*NRH3;#+;WpcpT7N9=SbR8*|>~
zP$*7%fVMg2yx)OHB)4dbW6m2N7+}2jwTi_)Eo=io;n*UNLT<&j#hkZ4h+C1Kk!_4Q
zuYo9tLvF`5#+>&-u$URuMSi)a9c|~kpDW^ye0urz$T^mndE^$`v$*HbZGw$)zkYBr
z<YsJ7%y~NmJB!|Dx!Go9*OvW~2!Pzp!L6NM5VaxuWf9<ThTrpO^s#x0*rou*5l7ww
zxH3P3v>`T$vlci_Nl(hw$efo%uuDP=BTvir$s=!wxOrI5pRcYPX`TE?P5F5f)F+=g
zescMFt1s-Rs7ppdukAK>J@mcEfT{i3|9)}b$du2T!e5=+H}cMhYucwyY_}Fqysd3b
z_xX|U?(b|}RJp{OHR4dXWA=>5p?lv6ue$nr>-2S-((l%tW8Jo7c=)`r=UZbNPDs~p
zI%Gck<;m%o?*~@*vH`(4k&nz<&bT$ywC7cG&Z_0D_s9QY4jp|?`Wt1nxuN^GwqKli
zjnzHmSIv7;>&=Vad#&}9->x?&tv#{zu0F4u=ia*ZiD{4h-dwY4>5kbS-*5TOnbLCb
z$-U;Rp>Kq`)~~ZhRlKeSmMyjRo;^i>@6G;^#>yqlBWG^3p8NO3ddXY;BRy~Z_4%GT
zWlY=sQ#M=sb`RW9f5#)1&(}DrW6T!(_v<p{rghc}x37&X9=<!4`TvbkE86uRTr@iK
z_hB~<PTw+hi?wXa*l=^}I_s+BH|l`Fe-1vUT@*^a{)(w-+NO%b51Y^Z<UmXJg2$|r
z-re0?_xs)E!plxcA6S;QzV^pkhfKQHY7L&$w!gmJ>Rui(YPQ{HjoiD&xMSF_tzB))
zTb^Dt*zy^lxBTtZA(6(DZql|K9BX}Z+riMY_x#bk?(XA9^gDfn+1PRC_UaKEtU+6c
zKehMFJ*N4@fVO@+(Z`%Mdh_hnR{gm9Lm&MwZCy5DUTEl!d#o{!k1{rQmRWBfIB~~M
zXZ5xAU;Fcx&MVhgqnjU8V_VO%PMguu`ebL9IcV$eHtjiKyV*2ygnom6leK8$w}Znj
z=ra4f|7LphwNF`%o9_%(Y`ohVK5d2m-)YZUjaS^RzdmoBwPgB>=?lIbWxd^6p<jMw
ze{0|Nhk}vst+jkzb!z2O)k^QXKre~yH5UzfWAnh1Tg}m%mb9$w+H9RWcyhD0s>}TR
z<YeoCXVaE%-^k$p{cEgI4W@DVr_Wfvxu*pacRyhD8~9CY`NxXYwLaN=`3v8&%BH>1
z^1_rR>*~)hd1}5Aw#MDOIo<SNd1PheuJkkQ|Fmu&_L;it)$zu%n*aG%>^RZbI_a*!
zi05lVKN>f*W$c}28rmOM2iE^&voZ11ZS6Zx|8M!$OA=caMh}%YRZi|$(>Wow;8)|@
zfBE<$QQy>rQMKZ-nD2nvHs)X?_OA_JnD=hDG*+>-$!w~zVjp&0W!`h$>ge;YY)n6T
z`PiuMi-@J}j>iVPbZ+{uzL%myuK$je4F0G5$K_{~uK4Ja@|8>fYQCMC5WC~?&hW<N
ztD`TMT^O0PabnCj=d1Mg=1A=0_Xk8?98wYM-nP-&`G6H`JpR$PYu;WOUHQp_ZGYP}
zE!ywt#*ROoH6|80;TEHN&j-=R-o4LEKK<`<th>^sgFi3#&3?ida@_dXu9j^(esf<_
z^sT;qc0PXg_?Yi-Wm})6Myz7-z@5t<oe=Xq@N09xY%{if;o<gmha02K^RH`v`HDpJ
zu0ao3&o7O}`b@euy?W9|(ZhqDHp|cXy8ItE_1&`hp)blif7sF9x!~g1uKvGi``Kkp
z(f;4RyW_K=$g3C^u3TIm8`p64&L_$y$6Bi9M^A0ND)#V|ldXY&SP)$`==RWams}kk
zHe#Q-aq;Nb{?DfB>)-n@TK4dNti=lgvENQx*51`=#eOt=S=%-5ofSLr$xCB{b}hiZ
eJG5<AOT*}cAKcKsa^%#=;rF}4KYQxa$o~OD9{RKZ

literal 0
HcmV?d00001

diff --git a/tests/data/save_policy_to_safetensors/aloha_act/output_dict.safetensors b/tests/data/save_policy_to_safetensors/aloha_act/output_dict.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..e8d537c82dfec39a503e6d8bbd8129dcae170827
GIT binary patch
literal 196
zcmdnM00GrX**PikIr+uKN><fMDJ7K!sY+H#ZpKDRI!eVEi6CBdtd3GjVo73rep*^_
xY6(zbw1JLEY^{z`jv>i9O>`{4I!Uz9LdVc3wzgJ5!MyZ<L)nul_A4z*odC!wJ?;Pi

literal 0
HcmV?d00001

diff --git a/tests/data/save_policy_to_safetensors/aloha_act/param_stats.safetensors b/tests/data/save_policy_to_safetensors/aloha_act/param_stats.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..6e33879f3432a7b83ad15f4101ceda162c1b43e3
GIT binary patch
literal 33408
zcmb`QeV9~L9miKe!+=QBFqMFXoRu!y&Ye5=+*$9;QihhPC0UVa(=0oSYx_dG11e_r
z;!~+*NiSxal@zFtVm`Y#vjfknPhceiyBdi+X-`=4WMKG`SQzv@J2T9gx%ZqqzccvX
zihjSp@A;i`&u7kF<}ACeVf>$K8+-aX(%sFeC7G_i-sBbOR7Z1rS85>HlTP(EwqD!V
zk-27By0NwK+=$s|G!9&m;y-=va-*>$l}RQ0Iy(o_nF0Kp7Q?#yIzy&5km;a^S%%{;
zc2&A-=@l7ragGr-m4d|j!ln_n%YyBxC0}gs>rFQ=>FZq)Zq_U^Y#U}vkSMWGv&Arl
zz8opEJ9SOEKiu3>U&W?nm_nlh)TyOAMnq_uzqqM8*b&o+_;s*iV2-F0OOM!w<=4rg
zc(}e%x3JU{J*`qFmTp;wrD$nE{K&Xa9}<>hL=`=)QYV%kHI1mWrDnZzJ!%`aw5J|*
zV(E5^VM|+D64xGe+cIovPd)0?(j9|ITgt^d`mX97NcAl1PU8p7M<X*0rfjXB2{8z>
z4M)-7s%1r;`A&=BC^}rFPAuKQeyM11)v_X&?ievehpW_yrN>MoChcsb-bEX;jc|*!
zsUCS^?U07Nw5ug~?R5)n=t~>x(I=M=EE2h>a*O5ie#?qFeP96_Wsm)As9``)VPAY2
z-*h9*RV$2~0wfidqm)OVT0T%JEN7{by#7=&(Q7$QdGx8}1E<1rp0es)$Iv#ya-{O;
zQ_BZZh2>0DlGh)7P*qrtRUUnE`Dm`3tGK-13L`H809IH|STY-G7=Tt-PFhyo3mL2x
zmcy1upISa(D=eq2lDz&@ve9cfaC!8p<%6`sB5HM{m*9M$zj^#AI)Itf-PM~;;iIG0
zApxX~ut-^Id1~>1tguK~QguxXg*Xg^4LQhw901J~SxcgCYK}f&D=gBMTAo@wV9OL?
zORBCp`VfbKfLlQZ!~oESzN))D%eymO{D>@>%4B+*yLyw$`una-YAppg8Ky{Io-`^q
zYBVsIDKc2CX1cB7XIUbR5&fv~_%W9VXF+y-%QNGx(_2Tdf*oO^RN9flVWNoWR&zMr
z<^kqR5$33d)PR6HQ-nK@UvFQ3kLHd8VTh|98iF{raA2J&Vx6RL-9F&7%oMRsEl({T
zG-rx5=c}%{(?HlTMTk?0Qws;cnIgbRv~@=v0B4E-r<NxdkDiM7=Bpmjd@cdrOeMVW
zxn2~f77lVVmE=~?)*W@Ao2f*%3VCYr;5JjqZ3F4<&ZDuZJ7}*G;3P)WKp;6&NpfBv
z-Od2(B1){2S`kCR1x~6*b=j$MN2Pv}g$tZi8@SAg90|X~QbL}Ok!GXd125s6&GTrH
zdr@Nn?;?J9#}}#IvkP*N^l5%sl3FAvE#gOM<BGb2jB~VzAB>5j)I$01B7Jx(xvXy>
zneJ&%cj*4w1gb{FhpfEH)RI9}5s|9Wy+`wy0Md{b)~8yYTs(Rxa+OqFcai|DA|hG^
z8Bha&t|B5`N%S@04bqSomZDmoT0HnFBJ!0~U3c_xLKYDbE69Kv0F)IGDa)hJp9Pv9
zS0HReM9fl&Qws-UMMTDuXzPwUAS)t5mRg=%Ji029m9M()j0a9dM4VEIQws+_St374
zv^7T^(vVjzRz4e#95)k-hceVv58^UY)O8vL6T=cg$umk4BWfT(k|jb?4L-V^0VG)>
zB-L(34Fz_xM0Qd*)d$XSfm3bsq0)&Q3G+eZtXd-JHVXV@iTp+IqQ>HHhKWE{WfiXP
z*DH{Qy!>HMB~LA$zZ<5Yva-5vBXHnlDZtEQKn(ycvlLub(%0NYAZ=I*EUS{I77r@3
z6jWAL*ByO;nWX?Tj{!LVx~t%_aC05p1Yw8^AAUUxOAx0P4koh{Ojasfcie$xmIBMF
z<f+Aj$}9zy71VWSJdV381(a2YQws-^SqdgAY3q(Uuq>*;vMPCE@lb}k>Om~Q8R|(E
zOkgD6ki*RqBXS_TxvStaKOfD`zy?;TrR<GVw4#Orp+yyhCWfk{8FHIIXi){B$(^W?
z0BBJKp!pi<HVS89Q3a!^yr{8&Xi){Cne}lKq>ZQo%c|t5#e>SC3MwnB>ox*l7FB?m
z$ABCFjaG13Nndjpfi&dh3%*L8T0E#Ms-Uv6y6)%$%%Tb~^B7PAfXku^E;H-lCI}l*
z1(a2YQws-^MHNg|($*b!U|CdwWmWRj;z4Cm1(g-lb!R-FEUJLA3UPAb=%|9pO4_=k
z4lIi*u&hd+T0E%CR#2HK!x`#^3=<g10ZnQ|3<MAO$ob6AN3%1qF>D2*DXgfWKxnpt
z&_t*DxCt(XtspeH6EzY5%~k-KuaRz}z-YFD(NtdKSeO$EM2pnNO^}AXeCMu`rxp(?
zvlUcUR@ZF=z|2;Fna6+{09<A(xU8hFxr;#BuoYNVB~L9LRAwuvtgNm(`T#Rq0cIWp
zY5;JVt>CgqJ=_Fgh|8br72@Q=(N_hNm9%xo9av^7u&hd+T0E%CR!~_%U3bO<%4`Lc
zRftmy2b0+fCM#*{jyjIYYz3B8$y195mDvg^iySp>0-!Mkph=C0f#3ljIiLCYXm$oR
zuu?7MqjlZ{7lSDXO$=2ZH^Bu?s*SmvsnUtsCjc5#0Ge+k-9~}Yn1az%Ues7XH0Fb7
z$*!KI$<A~tlkDnUR@dJlfV#na$V~1)jQ}=dK5XV|pxY4;8uKAEl?OEj7>)TbTCzKp
zN%tOwK7i4f52MNS$>pQXK8)rguiFD~8uQ^axdSx<IF0#mTC%^dy)TnY^>&o3>V97l
zHn8$O^=q)Ah62AaAAT!34P-F8>R;nHTw^{Ur!=7k0nITVnv;3x&N5&f^TE0R8)_Kv
z9TWL(MJk=F`b}<~e*+crun>0pU<+~xbX_DprGsui0DMdYzFJJELBM*B$a)G7O`wN7
zEQAFgY(WeGbvV@4Z&UX@YZw@g2zJUC#g<cp0O=eN>3ltew;ZTPq7XM65#@pmr~!a)
zj)-pxea%$?c~}HByD&zYT>!Q@BDP6G=&lH0n<HXdkO8#|KsHB2Hif=!7l3Szh-|?Y
z)DVC+M}#(Chw?|9x?W9%k>Q94M=_%Y1I#%h%+=_n+bN)(BcdJ6jv5Yh=ZJI{;8yv;
zl>9gY;W;AVk=)49Fh@kvtJO@maiBg&q&}h_H6G~C5$VszuJn4K=3jh(JS@~CP-Q_4
z0rGPK$gk|6+ZY`4IRWtJF`)(l`#AyZSMt!^g@Ass0Q9S}AclZC9BS*gHkZKwCTboi
z$0)X(90XsF7r=gnmkBjj2E<_?oaOm0J^}ow0l<B+0PZUV&|DRehec4c>)|q}!y%~C
z+(C0iK-`E0fM4mI(edwtp>_e>7YpFNlD=*ifPS$6^sBO<h5-4+0?1E%E3hsugAvSv
zS}uV=axip1fC2r2>2?Yj7z@Bag&j2<C@>a4fudV|Tm~n&1w9MOZL4&nMgs)K0wB;g
znr`E8;uj0xK$Rag9xyN#fPtmg4|R7RpkFKi{i-aeAwYhy0P-t4=r#uM7Yl$tj|n*l
znjgS^B@f+Q2<R6JK))&rY6y@YzSmax&RiedC|hW@gFA2=^<%(eL(IYrxQlv^D7>+#
zxjeW!ECTqi!h#$E{09d$0G1rI`vC|Hw**jNl?gcrC=eb&J+A}OxH`9}rQyXDNDsmL
z9RPxbfztli!u7dHEiFG{LJk5N#1~8lE%fq6To<fyCvF$ie2E=79M}+ULG!@troDb}
zNEmJjAVl%S&eCWW4aosPi0})VA?1LT7)3oy<5utsfJ8B0r5`ySAQ5)8&z2*0+6@FL
z!m;)lQ|(BP$Y)FdE0*4Y)ZU!{h%ix4lPVK(5D+51rC$9YDSK!)2$%>P^_1|~ki&oz
z;i8@fB^OPI<N7cOAjK*ZauARrJZkIwkxhjDT;R9}2r6G&<b#XlD4I_@7;u-M-V0}7
z<RDPz7D0^{ET|oUI2-~<vG5wG{uNS%%N7D)vC4!TgtKDBT?@WFuZQEH&Mj(bxF*Lz
zoIBLgZ~fH)eV7EmVtsG^a(V6%w9v~WP$CblZ>{D|unWLO@lDJMJMvWH)L3ydgYQ)H
zhSS|;;KdDtriIjy9+1zF0B#hA^9XimwvTgT#U%~_16t`vj>oaF;x>m`?X(-n$+6;o
z2cjc6A_vEc8y^J2Wbs<Xdfyi225^zkqRfRHh0|ii?GNfG+CAdLSaA(RO+Mr}oER(a
zgAi?&bl3Uinp~Re`#&p^9({WGyyX%~EWL7z<yqf*=-h$Z`0pQ58*&)Vi50g)h`lJj
z%Poe{om>2p2pIX3gI}DY5j7CMECN1~<`-U#K00p^a~D`iz)|)BuFmfuxe0fY%mqm(
z?M`uwthgjXoDym?a<@277H^2)T$<|cT3LL*d!V^}d1q(3KiStlknUfR8ovg&t0%QI
zJ&?@wEo<hVB*^#8axZQ;K0R~{=>hp1DUc@74^|raICT~;iO}#P$K%vlyd$C(yYZE$
z+d9si#VaF-UgTJuHj6h%)NZBMFn&N;yiI~;MGnOgvv{*a*(sCiUy6&OON*_*MY<j0
zR9U=Uq9zA&1P+nKJ0^k*^t!>Jv3ScwEe7-eyxRp_T;5xDDD1FDsSBJIi`Pxm=s}Lb
zL9ux2M6iWkCpaV)Z=R^lf*gWl;`r?o*Ug$WD>v(g{E}b&!~MnbnZuncM&0{AaDd&h
zePUu|*ZCXn&0ZRB+;v6%(K)gBu@Aj>L(KU^{JZ;pm5<M|;<v3G$;UrAFFxa2tMd<i
zVsiYI+s@js<Cou$FI@Ng{CRIoj{j!kn*8GVZ@91a4dx$zeV_Ym_@ew}^S8Q<E8ooT
z*?)X|>9%L{mo{hJ*ny4tpImaUd-kbY^V_!G?RGqJUw+5l@4B5YUY1{Q-tqB;EhleC
zyknJng0nk!*ZGUxDHBg0cBVb$9zXllhqs<N;<~3_l5@WHe)s&Lk7X~OKH?tt;O#@*
zdnd8|x9-j6|9RMb{09ezZhs`iK0E2S;j5mr;&-e)IMn!Ph&42PwT)jm|H{et41Q!^
zh#fk7ecS2pYhtUOJZb&Db%j4a@vJ8H)q4&O?tkPld*_J{WD_qOc2CKU4gGmrh<$YO
z)Zz6_(fFc|Pg*~m4Y9ZI|35mTi9N9T*Joe;v+LUy9ope6zTheM%0u(BzkSb$`}*a1
z*0Jv_H~;((+3js3Zo@CyKJfLILu?}Ic;!Hd^$c&c9{6L39gCkII<1L){FZBnzIR{}
z>l(Z%chlCx?uAdT8v4v%LTt~3H=N~d(Rk$Cn{#brWA1m3oyc5k#9i>W&DnFN&vWnn
zv|%?lj<_FvB4kZ@Da0l~Z~3$)w&R)5;9KzHx!;LgFT8m3VD^|3*wRy1*~d+7h_87e
zG`JLhe%gvH&MmvF_=K12;jYcw-R27}%H7>M;?CZb%D-^cJh!2LeYWF$BW~B|#-W2x
zhuD9S`z!Bgg6>|U<6i~8-N!Vsf6kpU^vky<vA>=+JvTbLA%5ewlhzNd3$YdFzLX1%
zS@9RAwdM9d{<OP#;=kAhlSkbAmz>yVZ<^y?^6G20dE$tB%l6RVA=Jb5d!T>(&qI;D
z_?hRqetHvoY~@aC*&ipd$4=RnyLv)HeCq18gImzY1s^_Tc=4QQJUjg#u@&uG+-={U
z#qOLo<__(i+q%5-nfM21PHvsEYApWlFQ2gS`Bh(Un>cjphN(BqN?dy9_gl}IIy2GR
z|DpVabI06c);z%$9~yP<x^rdh(D&xL3%>o%?D_>EcXsyTp{w^4$8la0`@+L(2VK-P
zKCbvV_xQ}g^k?4AezdhccjK1B?z`~cUxWYtKI2a;zHih$>5h!EFgNC=nnKR??;3M=
zJv*iK^DT|;eeN3jnuQ^Eeo_Az<UHg3P3+}|3x0fFBpRF8e}4KlYwBO7u=C#fRPJW<
zEw?!|xESLeKlW3(b#IQk7e+spYhCh`d+KHPu-X3}b$_?(!PsB7gxtY{yQ8iXa=-B8
zD(i=sN2Pl2-@Mj(U{_(@wZ6NFIXi!DedF56?3~a4A$#?-hWLzPJ<d8OH{rQa_t}Ft
z49{Dz-Mw;6gYyLR8#c`tzO#F-yK=`Q)*TJG4GllO?|#(7>s@eW6WjgZi`IePh1gE?
zaoVCe?Bet98fv|5GD|Kwj@>$`A-;NZXs8?WbWiTh?EU|R-mRPRGqZCN&g&=TyVstP
zXu4}^{MG$mZrgOx-5aj{`*wHH)AMti_Kdo#vTYk)e&&qC=!Yif7oRyJ@#KwXCQg3-
zYpwH>ixOv_@mlM1EuT)Dv*GuNt~VDZE==r*zvr8g#PF6}V!~^eCU*X>AyNJF{{j9S
BMs5HA

literal 0
HcmV?d00001

diff --git a/tests/data/save_policy_to_safetensors/pusht_diffusion/actions.safetensors b/tests/data/save_policy_to_safetensors/pusht_diffusion/actions.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..d9b2031708df0f0770c29ef3a44f5c6d0c1279cf
GIT binary patch
literal 4600
zcmb7HdsvNG`wmHrlH^>C4h}WDbdV;QrnPP&heU^Aexfi5A)+WTB;^>LL<vz+5hf{R
zR2p_Ug(itFHFoKw($=}X``&kL)9?FN*VVP&^{)GVp8Gkhwcgz_K#u+SYa40L^4DIs
zW4G6D+OxE6Oii=}+TI>+{Otc`&JdWm2(;I^?QmPY)!p6uw;kTxTO)y)v58B7K->6R
znH*1GY-BQnC!2gLnI{U2%}i(VY}0RMb3}oO(TpE?_Ka_3^F)D(*-RszZT8J<jwmoS
zGG|Gf`K@f8C@?kqk!8*No7o&uFvG}#%UZxB?$7+PF5)D{R8G#CY%TABdP)ygT`>he
zm;1dB3B_5cO_k85h7>wWrINi5IO+s`&GBXv&2Y7c?!<ajY?sIVFVRqpv?qDC4>5lH
zN3_i>B+-Um%+82_-SAv=gsyxGI4&KJ$n-Qk6G@<YXf<@^y(M*eWf-uwnkZGyLGFiE
z>^-)J<p0vl{JksgVB`GNF|?9o!!N6mEKc7^c?~goE^OrWii6)XU2<|l!p|oy?}bnB
zBAlP5Lo=)M@z^boo_jCF>*>01*kp<Yb2?EU(?*B<(E{C-sdU{-D>ToIAUhsQp}Jj2
z4R+O`v#SR|*7H$QFQLD!Q$<zb4Ni-YEP6Vh#g$cTN=sWO(}(~!(yY`DqZd+C33TDw
zYmBI|?O60l0V)58aBuTR*am%Maqzl%eG@PEBiyH#%e_pm=vF-Sv56*<>=hXPL=%Fn
zA~NLuVzfH!N1;<X?c3iA6;Bsx<LpX(YPxZ=vlsqPb7^HlGIeuk#2kqW{qFv0rmxr}
zkPMH`r5)=WQEvHw<wa;vLF@{2`t$WT{4EZ7ZzqejGN@=;E#3|kV7$I3-8i(I$Q4C1
zJGFs;qPT%Pe!Ud>VMDMr_byR#kHMU6&)GU3M^L5}b=$d>#y?lUFEw+C^X+by|MRgU
z>Ewf@$O-!oIrE)>;Zx?eVZX-$?%A3ct4W{~;M5=gIA0z1etp9`Qll2qZ{OS@8<GJL
z4A-8QkETXf?%A45xG0CcR;Q5gOo)+9i;*_D561iK7*EC56>Gh6Y0fXjh?n2X@b3+a
znf-tB{qHzIvl2m{<WbRROS`9xV1BkVxS>`xjaqDLhrOdR-GAaZvs)h118sESo!oeo
zy)wn&&8Nun>I+DNkSbf2vAE069L1}uTpZaIM6<40la=u=iBi4~F&(Z-)g)KRhTCr#
zZkzBi32gt;ALqc0f0IWU5k%dhn@lnM%-)-+tYUgj`su*rSrZ9~-a*V%x8skGm*A8v
zC9|K|kj%pkSXFs|jQ3u`;@dWI36xTW@C-@9;Q&X*KbxUXOI%(N;aMMKbuDB#!+p=l
z8Lf6Wr1wIj>M#!Y$@-G-&&JzLxbEYNL25TSjS@6G+RJGcl5ArO*zVC{`Q&k$2h9+?
z?H^oih?Dgd_s*wn6b;Ym{<;@s=Llaf9=~x@JHx3o-D0{%D~@Dy1^Zcu7@5iWl_E~v
z6cf08=b>FphR--iEPD7}l<~*zUW6?+iMTe?gS1r@;d}84l4~G^-iNbL{cscK=Y^nk
zS{4qL%Jg`wy+(y2>lokhuX?h3i7p12m!Z03In{L>Ob2Gj;kwyyv`=YedU!w2%RiyO
z@Dy>LUWjoO9}rcbNHu2vgtbB`y3MYzICxxTNd-A|(V5NjTYKnquSN_#OL2QiF&r-`
z_uDH!(}nVR33RDWJuTn$J&RW@HUftfp5s|=4$D)e(E$<^S&4>wMYta?!NB`n$g4A_
z3sl4m-(DI?fB19}!zLbJbx2zA7C93)V|4r)a=QH*N_6HBsg@23J|;3f4{OI_lV3N2
zxZJd^6482JOJ?U@x`gZ*90bc~uJ5>duz#+u>I<v690`fi!(KMpN4`g#Y8|P{e~-kU
z12Oc6sQ!5DXFB2zZD#Y9pN1mYd`!Qcgp^FiAJXPV-bHZdHs@o~Jq5ZYESJ7`ehxV{
zd2HR?Bpq(YdSDZ(PAx@lwEO%mGN$b#>~&@Ss~S%*ocN*a+$hUGM+$E!)7as%z8Sdu
zGFxYT@%8+4I(!hi$Ie4(@oXCBupadWigaAy7^+_|g%%&U&hUr(zSfzKo#QD!<L<Jr
z{5LP;Xxor3==r+PnAw~T4pUJmMZQ5ejmaUzB%z1x+cJY;kaxI*eE&zt95x-pYwV!u
zT_Cf=aWbQJ!u5wq^q)_qtQK8QeOT`4t2K1j*%Kt^yTQ0KQGw-8=R+0aT|1hF^?7@l
zPrb!qsGL#?qg7MU+&KzHM`T*&tSTdEKlkGEs|OfunT3ZDR&?rKC7R^?2;$9(_>%O9
zu51g2vV9n{V>|5@^u=+kUzY9U^2TWkH|oR&J!|&N$6X$Lh?=eKVfcUZAnS?$*|!|d
z)h?c#zHpq=Y)wuS@inq0u__Wgdw2$W=Jw)m<9KExbks%MZx!Q-r37)7BhVQq`D)vm
z3_ExizfSm#`Sz;ahxpfzP^y_kC+rr%>U9ONG48{vI9KX;q#7?1<@@V0ead_0FFm3Q
zA*%Jv-i`W9*xf0?`hB}eoEhJ9|F73+XC1s!G-%BQBQ(}1W7ZHUCan*}s}0joon685
ze){r{)T}^+&%vvRUL-+<*A-@0-*F906fQB1*;m&vySCfWpRYvED2BVYEsdJ0zsB+S
zWJ;szAmD2D=*vB9zqW@A98-#pr*W8eBp(xeTcDF#go`_lG8^-@-$AoQ3J&^o)2*4K
znf(Mc5wc3wLex=&bs_tSe$VRucz%fgooL3tVe`9B9Z;56jZ+tzSsxmGRG~pZ7eF?O
ziAVJbs+a1Cy>^t?c|XV7RC73KdNbbG>m^7m?Zp!B78F^f;CkRE_C5_eSS(WuTJXrU
z5|-v(P@E)1SEC#975Q&2Avt^e3z95)VL7W0w^F;&7{i_G-20*2T@m3KjqSt$7s|Qw
zmFumY8@XEZb>P-dm0V}NuivAFjl~XZy&$Zdt^ZZpM0QkcV0Bpds~@u`>p3BjRP~~~
z^cofnYan5cN_26p2><bFAvaAE>GTB`5gMukeFYJTP5+1SE}oR#%cFJL5j)VA^(XIt
z-^=y%^r8ajZxqA-cssU5E+H+mma+A%nx1G;w<Dvx+#uf6gHD&1OxI_tl~mzLE-Xek
zVN*vFvoBK4VDU&BuHtUm9D2mMgxX)a1eM8|5YIfrc>nGfgTXd*&&P0j+FZ<H*WJl~
zOIQ;H>s~wvK8B0BUI<%wofJLqBa?I2vwn~t7Kn7e<!~@9q*Ld=p_+2LSzg@#(!r^j
z8>o)fBp^8nVJlQw4&?;Ln65>xT{x6oNHer@iEZE$cxwN_bh?G`H4u_DX{|`{lc&cD
z!sxj9rSK3cA*-^Q1}|+z*Md1z>*@)JX&(`7@W=04xcil>pWVDn`t-nj@-QR-w!u>T
zY;gyk=?BTuFE%XBKTmCAqw^FgVx`@fdr|@)|HJ5Ud4roZ<B{)fLQUNwng98bhnP*V
zU;znv7bDZo`AAB-i@Gn(SY{-<8x|c)p{F$t;NG}q8v3pWS_QJRNT3$VI7*kfUgdav
z9@h+RBWFXB@p@7hsxC>OT$qb<fu+!xYX_ryB4Xn=nx;od@iade<2Qz)^Q8pGHeZC)
zPnn9hUtl#UeR&jvyrbyhNu@Bm@(H&GJ;I+O>S4IAjMbyk;55A=A54F$l|VRA1sT@d
zeard3YW1#P$AGd<1cWT1o2FlfReBg6XIwzc12HtFX47T6xbIXB@BgY3j<E@hv$C`b
zvow3~X@czB9HE<o8=jkSXF1nn9QK7@7c7^`zLy%_1E1QonGMl2d$#Ua8G)jcwd~m^
z?Ewxa{tb=gQfw+ehxDIK7?!v5L}L`KZI^wI2*wPd{%=1ZeCi9tcD3U1EDcnANyU>9
z^T~s{Z0KqABEY>B%^o!{Iw?<`|Ga{y?QdyI@c{VEh`{CGXp|V1P@|DA5cQoQ#IF2r
z7a{Qp;_fnTREz#Vi_|}p^MSH{<o%a>=%FLLgyE|CifKX2YI<nQ4dPaPoN6XVVf6B6
zMAQ*Ll~!z{d-6_@+q&Ht<DA3#map&mOUvMU(}fK3mJ+r1{23)A1-2(Kc4<C~eeB&l
z#QUyg{ujS}L*+tM==v?4sNI=DwH1X7Z<0~S#$}h3vCwiLy?pfp!|$~`BJ;!Pw0M(^
z1)(8``*jqY^|(Ic{4^C8vh}5nH(=15MGL|`Fk3SO$CE{fcN)N9Wak<m=gQJfmh;PJ
za%u8Z*?x21|BAY2cClXR_~iw?zvn%)EzYA+)B&@Wbh<LtjhK~hgx&XC9XWrAOKWi1
zAP9;tUz6U%4bUqrN5dg5rvLH7<9=7~!P!jN{^RGO@Lp8;9-!{=1m%^6pg~PIG9w?(
qFJwA4{Mv`YhwZd3ErS~qsPu6L^JkXD|HkKjXL9=a@%FB{Z2T|(-&TD9

literal 0
HcmV?d00001

diff --git a/tests/data/save_policy_to_safetensors/pusht_diffusion/grad_stats.safetensors b/tests/data/save_policy_to_safetensors/pusht_diffusion/grad_stats.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..4eed4aafa56b6780475d2721b33db0ce8ed129cb
GIT binary patch
literal 47424
zcmb`Qcbrw#wZ})WAc_s_Paida9EMv@yB8(s0}CK|Q7nnlrch^)0To08iiIai1j{f1
zR1igaaT+C}phU$OBQX*SYBb2l0zL(M%x}*v&OQ5_yVu%R|5NaM=XciH`+U!yJF}|S
z7d8EJQ^%6gWo47AMpsr8Pb?oXw6tPmWl8D8;t|6~jvY~1QCb`-zM*t<`KU=l$CnPT
z=;-3_Pnt5Jw4>XxUo_Iu=vXysxcIl<Tw`=B89r(F(8{v1s?teS_%9(NcFoQBuf)q&
zO)BABJ!aU1tA~q6j2>Ra{?ctD9HC#G`IUzwM%W}>UGv6=O(PQG0GI(HAtOR5Z4A9L
zqC&F-k(d#oR5%C14PqNnN}V$d&3Z&5MwC);4~7{gY8o+0)l-y+-ol7sQPWGGhA@M~
zVn&Q!`H1fJ6tj&uz4%!SGfX^U#HqC(Sy^#?ST|1GG)!vwb3n{MX2>w99UwT_8^?Wo
zm@&howm}Yv8^|^+S|4C1)7{Oih+)y%pc#Z2#4-(=)&VJo=5A((44c*fIT&Ub8_#!I
z^#w(9H?wWSp_g87H*+F}L$AE{YH>^>9HLi0%VCxU3~;FB&)H|V>mL}Ps6CK_Vuyn7
zk%TW+94^k;YuKTL{E_TWLUBXE0>dN?j2EJ5KmbE&h#U+z3~Vq=(!mTxb8`d42$P)9
z1L6jP6NX7rnBr(|a6lMgk`?A)xM5(0VUiYRD7w{w8iq-3=mD_<p$kZIC^$J=4|_v^
zAcjeXD51EaV2ELoB4$Y0hXgwmBr!~~L<z+W1xpN*G*OUBDsQN$8a{r)xB^Z!Kr_N5
zVGJU2Bf=fSBzg1zAx1i9A7B|_l1TdVxbfhVVUkcvINe<YZW$)YWe|}a5iLU!%{cr*
z`yB%aU?{GcgW-mOX+}t<nW1Qw2@rs2W>LtS1{mN_JTnzY^HB_h5g|!t4u)AB5I~{#
zfc6#-h!G*`W$Fzr-3Hu1;;@nEmnk>bK0yHnC~70*JoVY3&<jM%EbxT`7@(*Pk@IfD
z4F%VXkX+L&<b?wWU?|%{4u%^BrWqlbW`?4>%i$O@LQ+i+h#Lr=86kOQile*BK{6vG
z$;`oU!@w^iB)`m1bgKisjF9xw17Zh47m%b<aB{XDuA2h`6lFt@P~1=u$p}d#GbH_|
z5hOA~l1K@~4F!3OkmOO2evTIo(2NMl9)pP7i15e=$s;{L%@+=^j3~({{dvrIKmkX6
zfHbS1wpf4$B1J5ufO^sdm=Ptprdh}f2N1x}wvuM=HQai^G@~Tb%!JY0>;VCY))<An
zaDV|0tsPS8to0rbgb^i4X6nA_eR~50Q0P6Ny~P7!M2UKtx@($i2YwkP`elJH9KZlY
zZG@bsK6eelGNVMxEbxT`7@(*PQOFAiDEM3j<$GRMAuk+207KapaxmNq!8D^J)67tG
zcR8G2MoFsa0dWJtGovKWOmQ?{IDjysB+1OdaKpeaqa?q~P;{#Uy^NCd(gR`#LKl#v
zQgCv%9_|YVFhEf@1PR3r1(A%AL^4Ctei{J-6m?T5WS0X4B+5fVgleX%+MnA5&4`ig
zF^I^l8(;uZJEU1eJm+=3!vU5NBN?S9V4BB`2d|8gyi&qx?joRpsL?aoR~&Xkv<%5L
ztwLTnfB=TpD8d0X>m@!(Nixk$7~L{8eVUTCiKODBLEP$yPgIgTGZjelQ4E9;BS~fs
zh8qTc86){+hN4>?=w*zkm!dk^X9+hD{4z%L%K~3GfB}lOFL+SwQ0N7sWfu6t0Sr*o
zO`)JiF<dj2&owJ2PttrM!7pR^{4#rGZXMv0v3x!$F08via4;Fm=a9{=&Hb|BjIn&q
zIH9s~TyX_{&P4Nm!xZEBOwn_F=9iC4<MMen_Ma}VdGB#`5b~^KiNKD4kH>o@a5C<A
zarWnVhE`1~osj)`q4cws`*r{ZIG)u@;BezW`Ql#X%a75m3~Vp%wLKDv8wrvZ_ex$q
zP*iWNfMLYFZb#s7<G}6WUboAS(X9~tF7EX^5{VlLei!#1C8j|5AllIJl_itMl~&=c
zt?3QBq~N6l&5L_AkB(!9LtA)lk3`cg8^kZ}6+aG;8xZ~%_xfKT4?nmwVqE3Ov8bii
zJ1Iz@c^(;L8aEm|Fz)q03QV_jICG49ZIHp^#)A>Yy+)|U={~_=g>kPHl4;y%5XHDx
z6e%#>%E2GwUVmip*zwRPUZo7i;YGEeH3!R#do7bp<3@vWnqK9kz%(ld7I@?h<)>Sj
z<}u>|1)N7n)$ugf97u-gZ>K&_R{RWNs!!51ZZs&W>945yVY-!rtD63<O5$<jL0C<H
zVfErBmR5}}nLK>lQ1KS6_nZR>G*X}V(Ad$?CjRQ0Kbmgk;IO8@!;*O1co13BUu1J}
zO(!e``z27o5q1=R95)<n)}+|1MAI!EjxSA$&;|jy0pYVI#b=v%h598(hDotmUm7<W
zWY(m}tQV$RI{2(f@!0?#J03cPqO-X;-De$a)}+|1FO3@wGHX&~)(g|E9DLTK_-p`=
z8xJ~bQgk+j%Xow4I~_=dNwHa98aEnb)}+X+7p7Y|_^d_o*#I6h9#FthA3k2=<DN=_
z1RB9oa>MEQB|wHnky$THvvOd8=X>PpqYce&WtzvW9CX&A=qy&7&Uf>1ZCszw8R_eD
zzkCR*MG@BQ_y5?<h67fMVyRgIHv+WOQfq1WXF7)j%COX08ieD9gOyrpEA^r2E+2@g
zr54j*AUh!1L+zz$UZI`{5@^cJA%Mn>1}U}FQYyoAO9wBt)Lt6M<Hm!QT52uz;xs)G
zB*Ri$X#kBI4N_{UrBsILRt{cjsl7Ci$BhRqwbWXg!4>F<Ac3Z|jQ|=u8rnoHr7}#n
za_~}1?WKV{ZaiqIt=3X6PIJwH1e(%40%*)=00G8slk8{uxLyL5VXM7V4wvo}ZaiqI
zt=3Y3^Enx0hl7>cYAX$(aic*|ZMCAxFwHdw7I>si$-R}Dm4m?AYJm;oqn|i9F8CZ0
zDBuW+;*aCj4mN93Y*wP_77s#eQ-n4M$PEaewJAQ^#4FS<K{9NL&HB=~(IB%nMP|J)
z-O|BlZHmtZ@VN1yvo=L%b8))QI@qjDu~}alI~v-ABC}qYZsp*!HpOQHc-(l<S(~D>
zDO`bm2@+@oOYx;~qd{hEip+Xpx|PErt4;CQ03J6Ubk?TmY%WfB&B10Jip~1cn9%?N
zjNK;MehC(M1WWOTOLq!89)8;^MQ1h7A#rV7pP(gr=a9HGep#*Or(NX_ScHosx}WF9
z)j=TCKTF_N09xv(wKVcGokId;IBG2o!g0gFN*%S8`p`5tKEMn|EvCUhZa{dcqxRA?
zuTW0}2{h&A5I|!`L#L>vREFu64qoc0y)=-=jR!4t)LQDrX+E<+0!_J%1kkwAAf=94
zN@bXC<#5RAsJ%3h$BhRqb<|p#!4>F<AQ_I@N&{%zXpmAzEu}I{w{q}ONA0D7Ja#-Z
zidsv(INdb|D|OUX8bISlgOuX;z$;ZmrfHT=zyy)Jr)1VaR!x}+1xo-)kDuHVG_7x8
z3lw;2ITg3mhsTZwRRvDzsxnTyg0R<cNUg7dM0P~TEP#ZTX?nXT_)t=?1dz}&{y=^}
z(J~ahm3Z1kgyI58XqzA;J0u(zb07#?O%t(D=@clzqjrs#$BqZd1&-P?&2ZW!#PMr5
zL{VMu4NVi-5uv&OQhO#FQS(74P~cHyH#cCK$BqZt1&-Y?*(VeC`&I%bi0q!p68RBD
z%TSy*bw@J~U(F9q3Y6fHJI2Ff$Ak1@%nW&p`3a}J24TJ+l6yuXvLix$0VKCf2B^IT
zaU>fKQRFwbV|3Rb<QF)0!(>+ke+>!)hP36*Sd<UrBB}3q7W@%P!4g2fmMI%uk@<lf
zthhpc98LfP#jj6-Dcokn^#So!U0;aqDg#1r{JJ6^haE?Vtu$w;sTF)Mx+|(_%91o@
zNkL&p5e6$wSqg?|R|e;@VQI*clEaQ8Y*m`E<i%)LM3|{GS}EY_>@Qtts5DVYzPNS;
zgm+4Vlmf5M|MEqFeI;}J=#pl?_u%LFG>8BRn%_cK)7a63UP@Dxn(ka)n0BRvUrM8t
z$UJsDA(zrTC4XF2)pbR$X_AsON{I$z2NO~$O;J*^w96)>QW~N}#j(Q)p_FDQDQMao
zpO8sud=d@D4<<GeUy+nnTecnbwXSJ^k~Bd{87<u;>~KONrAbNww2YGLt~D`C2{1o_
zl<%Ym#tkM0DuLyvk@;A<YfTJR3QLofl=o4yYGS|=Xr!l2Gw;_s78XO6!qSu_MKsN~
zC}0GO*CG0AO$=KCj@Kkv9J?ija!M1Iq#N1vI$HQqu$Z|7nyN?qY3yi1M5Vb)>S)@_
zO}MEvfQigw#}le5O=9xLWgh|ByQDZBm8LP#VC-N*TcwFiN|tungt$tRnW#8+IH9f5
zgeC<|*GRw-Fx6I)I%@)?>5N2om&Et~!_uTCWiaik32l|8HBoWwstIeACN>4oJd(gZ
zLjy;^RL`5-6T<YuW(O0_Dot%t2Gg#ZIRBJJHc@fxaN_t=n%AVDX|FXQt<tb28jKrE
z3~B<4*CAO&U^k?g(gYlB8_9Hk7RL{V9+74=HQmU`6Pg{7{5-!JQji4BPe}+kb~xd!
z(#$48INcQ|MmGWI*DCqrrGe~#!d|5rP5}W6))FzvDJ+e0BGTB=guF@<ozyVx(g}Z+
z20M{??07<7rTI=OobGlhCOm=Wx2MQ`ne7uzHvtDQ-(%OjFzw0-gOx@-sc)t3niB>q
z&3g*MWj1KtR{=5g2{hV{l6}$F^b>FdjMyeWMboaFuvlpbl=@ccR!&%~Gz&_F(_VAJ
zVx^H#L>hZL5f&>=hEl_{D<@1=8W2V1@#Bd`@jZa`Y$-mUO5_`uFj;9-lwUX>8atXW
zS!r678m3!0F)<1}!iJI?PIJwPu~EP&w-o<FD0b<@%&4$5G)lgk{y1*!#K0)flr}5T
zw2LQ%RvH;ah^JdU;j_}rr~qEp8+4y_Vq6qxv|jPj*tHWfE6s{hN7F8y@L6eS6q(14
zCv;Yt9;L$RKI;GyH07p}Iz<7ckh3pv1dQ1x**9eVo+5l!8Yrc{mAY$A=&UqZDhT(J
zgz0^22O!XVKc|}WWKC7hLt{6kkXdQQlscOBniD=N4VohJ`0+%e_<lk)&%M-MbHZk&
zfm5n*X&O74kXdQwlp3a8IpMR?Bq}nG9Z%@2G@nYwq5pM$HKe#SuCMmj^!3?az7SSv
zPL<!6uU!vesnVpXJOpk8F{%mx#S<j`JDv*v)~J|U1)Sm;5`^Q2L;3xuST)-id}zAm
z64R^zRETLXkR4EXssCInIiT0v*biiY1dY}!UK%@^kW&BoR(Ydomri)8|C}odj~!2F
zssFsIJUHEF7C@l+{v4H>#*QYW)PL?(UYK^}gqQjczM}B>@kFDjj#<-d+<P*}Uvt7r
z{l{MUge%Y!!2wJqrE)aw$_X#^pL&%yT$;zOoX}GLxmS5`+G|c6vieWGQq$Pcgp~Tv
zyvhsHuAJ~v|EX6L9y^}UQvbPE9$fC>nVvNQBxnRl@uBgfi8fIws#tTmFzw0-XZ4?Y
zmG@Tat~nvF{&TN#aMDqoUg^ZtE8vuh;*aCjP7J*QP1z_CO}BVr?p4@-@Kyf1saZXs
z1W)O+sZ$E<lm-3#89<=<{*nMWU7>ynj)2iRr5Q}Sbi!x-Ctu|ar&~Fpv;MQM^5Arz
zbpQz(vrp1A_No&y>p%A@Z#3=737_>Jd`020;|ZPhpM2$m^B5o7F98BgZL_Hr?xC@x
z37Peud6hSscIAZ6`cJ*0@YwN$&ic>2%7fEhbHZl*r(UUP{Ai+0R7b4o-|(D$>Cr#?
z2TJgW8%m}vvpjY@p|k#TuX64sKw9UJxU{&w@?n#{KKsiT!s<WQ%Jb!GuOxAh>OafM
zpTLeFw6tk{)y?gj{#n`ki^eJ6exF!=c(pS*)+6yzdAG(7I*duQKk*~$wfX;ai>99+
zzIpc{ch&>XJ8#|fq1$fRN~h1vgYN7_HT4HNKj|KPe@NY@bC$RRte35=t)6j<`t(^>
zv2T^Te1ALVjZ5acFMbqiJg&IG?RN44Yu}bFZsgsY%%UN`b9eW<tG;jeKDX$=;?U0R
z>)o3!c(A^I%LVRfi~E~zjaly+AIzx#^`}p{&JL&Ui$P=EEBhv`yU%zi`}h8@^}dFi
z>K}0b{C%YH^i7YsMMc{~9}R!d9rxnkhM6x`yKP>6-~MLy5_j<CU3DXeRl1Ssd(9P*
zVeYwO#xyM3G6sKtcdtda%ygITN!X8nHP-!O%Wn4AU3a=iJU_2)=Y3P$`&&O}w*M~S
zF1WOnb<8O<vVX67p4)J!?I?Hl@l)*uZ(fS)j9GKr8&|t~XV0?U{d|Ob{q6UgT`rsG
z_J3u^Q?I|$)BRv;tJ+@Oi`~mNcW-#CM-R7X?c7kb+;Ph`46xT;c9MJCs}t?_21nf4
zWgV<`y%X+?zMboaZ~x92a>my9h*7`D{=IqsINpBXxc07H|7q=+6=%6cTQ6Ai&WC5a
zb)mDYK4W$`pLLy4|Hb+3++Ck^sI5HwYv-$Hm)1UC{=M_s?smynOEx42e|CJsg&&tC
ziYg|rIy-!*c=IXeHQe~x>cqP@tZMAHY)+!C&EAH&%R40o4V=2^)NW5FZ;ZBS+}3G+
zV%nqA+!?JqCI?z`t)u#PL*8LMe|PgFcfxgN)vu^)>sDR&u({;kzdA*IzKQHy_@Pr&
zw6)iuqn}7VJN8I(L#s#J>Yj5g<FIl!dF^fH_H#zNxBRh%-F5I<cWu11&Yga*d(8Cz
z-)R4~OY){)9o^8mxER;p+H*|#FWvZOkJethqP^Sd!t?4EZQbfjo!+JHuU9X{^%tz|
zu;hj0C$HaMYqeS7Cf}cHe*0pXd&|`et!QF~J7mo{?$Up^Nj!buT<4kX?GhLMexG~C
z4<nQ3zW<FiYVTNFKeYOaTlc%|k31>*tLw(Ni??6faN@R8+{ZiJ(_mDeDDs7doN##W
z+p9M|wQuA{?yYT4tp9e~GWWKrHTDnf7r8Aij@6%Aaa7{oE9RST-t~~Xyi<L{&f}8F
z^3VRKuI|Ua`2YK?>i%$zJEGIh+D{JkbC=Z~+c5N%1J1f<$2WXF?Et?16GH2zeVbh0
zA{rYs=PkG2H-EBPy}itBH?5n!z3vV7v}?QDy*51Nw7qJMRXwN1*?#aHd+8l56TkS^
zQRba@x5ocp)T3o+l{2GwTlDEr8|Rt7Z8B${yvbg$a8CTuH5=`sqEV}-4}3Tg+Ouc<
znED%>-4_fnI}SP588`FC+CA6yb&h(Y+?p`#9;a+To7!zlk9I=c``2G}(L3&*-`*Cz
z>(Wo~|8HG!T>k;izJXn%kG%bX-Q|~S>h@2YXRrTpVO^Ko@5Y~>Q2h4+^Aa@|*F=VP
zu5hZ4{f*V7o#7m2oo=n@aSpz|N7fAMH7!|s{KMw91s#%8XWr^=J92#Dyz$H2UF#l8
zoUroMWcj%L$s1=4PLAKTsQ1cQSJk(;<9>JitSwK?x?vHn-+o2;{vGzFjUzX_-C?nP
zQmgNtTJh3A`;jFJYfH|)8h`%os?(QWl9>0>J!=o0)zz8O<8pI%>wR{=OM1qqeD$8)
zweKK1cJh38arG7PMeUa0?;lw1EO^|m8*_N<=%XIByPjNR9&zTc?Hdm)iyyq|a{PIy
z_=&X_B|2U6%KFVa|I3-W>@)M+p?mEK8@@Ha9JJS7cKgKo%P$z6*s}1yYdc<Yy}18H
zZ!Mbce0OMl?bn-ccKW>caQx%Lx;S^d`8TuW^KJ3x(^n@Bo}K)|*RQWXvhk2}&c!>d
zZN0WT7xl2Lw`yK=p4e4mU3mVfZjbgOtS2`|+_%p@yz!punTd51+tt<neU`YsarKV5
z&Och-6aIKlx%2GZkJp~}^6AbaH&(?@J1K@gzpU5WrNzl*r<O&(n%vsm-)plq<HtWc
zcm3rQd*c@`Ih(IIV2@b&r90{1T5JA$_qq%2d~j3IhH&rJEgo%n?D}v5^;@%lc4cDH
zLo@46dA)t&{5~5RUU=?Zcle+$>+7$48U1@iq;pAq?}5<)wf+9RIMH+Cla1FE-;(HA
z@?*o+=Wa<Xxb4)92hLobD0`u8<E5*nC#IdXa#OgbRqtmuKG*O`pJQ?TzpVcKKf?*9
z^BLClXMFEokZ9vH)+}{v+I??-{)Yv){t4^u9$nM>f4djeS$Eu@XtD3r#*xYKiP^KZ
dG``nmSYq|*S2p}`OzULVDgEkS{@~=~e*h>$TSx!^

literal 0
HcmV?d00001

diff --git a/tests/data/save_policy_to_safetensors/pusht_diffusion/output_dict.safetensors b/tests/data/save_policy_to_safetensors/pusht_diffusion/output_dict.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..77472bb50b2896c48a0a739eee9848e821769021
GIT binary patch
literal 68
zcmcCufPiYHoc!WqC97(sl#<GVR3$4VH)A6u9i`%oL=Z1JR!1o%u_Q4*KP{~|wFIaz
P+Cax7wzk%Nd9ytLS|1bx

literal 0
HcmV?d00001

diff --git a/tests/data/save_policy_to_safetensors/pusht_diffusion/param_stats.safetensors b/tests/data/save_policy_to_safetensors/pusht_diffusion/param_stats.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..a26868dbdd2782b49aff295071cb6b42d7d6c617
GIT binary patch
literal 98776
zcmeI5eUM~VS;iL<<44L6Hwj7k85a=6&bZyv(;vItF)@?^RwNM>3W)1$W_B6bo!RWn
zh7S#<Obk-Rmx%}_DpDXwlwg4D#{^1|Do~7(AVd*R6pL7XL5!j(h|kk~x@Yb^=k|G@
z_ip)D_y@vf`#0}>-*fKop6==As*T6&;D3H**V4*?1IN}^)>h|^9^AigdG*@0rRAe@
z`xmc$;r_MN<+<H+FIrwX_=1gnhnE*uckRKyzwweI%e(gM`oa0uuEwtQ7cBCho_9@S
z*V5v~;=Z*52iBK2*75)BZnUqt0spV@<<~cslvi&zdWKhT&h1}WT$g|JUZdHvzq<I7
zH(QNn$MovX3*YQCTDuj1SU_ucqh)Dr8?&>2)!BkpyV0_AxByWL={4ppeV#C!{g`hx
z<}D2`VZ>tQJB_xb>tjmGw=mjR)b`c~h*(Iw-Dul8-|}ux?OtQS-ux*>EM}qASg`i~
z+O^f|o8EF3I*pFC{RNO%P-l0eW1RqUrWa1VeRSH5j&%$QAhn=gqidT3>}1~EtlMgI
zZDX(vA{Nr^G<voP7&Dx^S#NivXB&V5Ml7a>`@60Cq&Ro8Ua!%&x8Ap#^;?a;z4PwX
z((g2yyY1bda>TZP0gko(g?&a{|G)rcoq+;ME(*S962A6abFQ%0$VE~9O!gO{)S_U4
zO_K&L)X_8`fUyij0izZJ8*G|%a6)lzZh$nJCMPU`)PmrIO_LOkInE6Zh(^<7g$0aS
z46Lwe(!vSF+a0K3)8vLFkX#U^z$Ay{6t*6DLx3PQO@<hu)S_UBO_L%{NYk5yTofd+
zX|lu!r4|KCY??HYq@}eNt*$Q~K60qWsRn3`rb!s9h-!&&$EHahOF)Q`=j;QlM$;sc
zX`Wghe6nd0$_VG(Mc|fAlU!C2<q|P6Ced7gU$~z!fB?qgngxtn3{10SGR+Ca*(N{$
zvYADlCk-&bv3TaVAm^<Zh(^mKnFWm49S}gV&wzW22c*$5>gD(eEsTL$5Dy!TemQn?
z_X!FxKv@T&aMzcM!Ymjqv*rT_FhE%cqVQ~^76sRAnOw6|=YazVU@Y500izZJ(`=bc
zb3*a%ayW);nN+g`QVW7-woINm=6H8GNM_3<nFWkm4E(ZX^2-Uu+a2g-%cPejkX#U^
zz$BIA6t*7K&4B^RvLQq$wJ3;W%OsK$lK*Z5iENo9GD4|EK^|Ksc_is{JaB;4XqoJ>
zil~+dk8GJdvIKNKaDdgAHyLG`Czb~kaMl~hwhp?*0yGdUVi6VSlP16#^Cs8q)Op|l
z0vOv?vhBG>?H5dQ-ej7SVw{^jAOP70qs{{d7~t5(VXV%s=Xf9*^Crn0UpL>|8z6vU
zp8@w44@hI)sF&lb=3G1Q%Xy<;)_mXq1}N(w6z=-!8iHlc8!fZu0|zibSqGxd0|zMh
zT!rO5uWp?O4j_QBYzqa9+98<cyvZ~t6z?vF^UHaYYL-B1LGa9ZlV^@O&Ib-48uKQ}
zEMU}P;Ft3zznoCK-GN@toAj~-k_*BVn52@N!q%fcZ~z08WkZNiYEcl$wn-!>B=_A2
z7@(}1LY-X>6p$=82^O_YSKXi61g+6F*<%$^>>FSJS|?;%iMY>u-r)eN(KZ=nsldQf
z%Y#?8O<oz{oVy5UAUbAd`iLW!h><b5X1C4*2N1y6289mj>=!>tX)?`8G2S+fK22%c
zM8@R=klG!7qSEA<<AR*GVjvoAlVlb!YBBK3w#hFi6mNH+mu;h7vOCjziCPf+vTgLs
znhzYn0A<@3N+`J~%!1J}Yd&xQ1C(`BsC6rbYqn>(=Gw6h=N$=t*`DQ>(^ppe0H18n
z@=3n1cY)wwvOUWox4pLdW5XHSvz+nB+S;MHRs5cb^LoP+7iO7a>Gj1QA1{rUpRuuT
zy}a|<<JCc!v63T#Tmn8GUzWgQxZ-ou-{;x4zOj5{`ujrRUaLMGfB{Zv^%^*8IZ(cZ
zvhvNwcsm2zTPWL}iKLbU$y+E(-Yn3(Z>@mQSSY)lfuoiKw_7N?-E54vL-4zWvfr6V
zYDw_Bh4NA27=#a^?K`}-bnMXbI-c4_PuNSfml8B@p{#j!oLn5nqHKF6nzwBbzlE~+
zDL}PA_}@a=|0;R-#hv|!)~<aadg*#51qrm$EyGMxOM?e4ls(V_^R^CWjtgZQ6nJWR
zFv5kh5$18;I~c5Rp=^a_npzq}aiJ`V7MQnl@W+L+KMFj#JPb-%DXVdKP_4D*V3`YL
z%QVx}(x9B3vT|Br&dz}ao_Ryrx|D$@mIoAYB_Yk1=Uj6jH9BcKP25@WJBZ^s2{g4d
zC~7CIsIy_-&cRhXX;(Gz)bb##owTr)aYvWeSC)<~9@@vJXy1JfB+yKAlAy_@VNBBM
zI=eJ)=isoNw8NTsYIzXZPFiG(xX}qqZNCHxIKz&T#;L`@W;+&}jcDHH;rOy+5!x!C
zS|EJ3WAWJ$udZK$)aY1jHc3-UgUog;GFyguTL+)*SbVmECzppwvFL0O=e^g#W;+&}
zP14lTAhR8d%$8x^&cSCp7N4!)spUauI~JWC<0jtVe5M1b(XrTUlBSjhneA9)whZ%j
z4nEtp_-qAFEDtE)tT&&s@u{bhAc1DElwxr{zXYh!wa9E4=Ik6;;3bb-6WXxtQU;#d
zIp}QHqO(|So_F){+IW4#WQ5mOe|!jQ*CMRb_y6Q!!vSm8VyRPtS^~7Ro7d9jX*!1l
zs?p7BX%$W_4p!RDTWNykT|N-gZeC2QfpUQukGz)#UR_TF3AF6yP(f2mgOqmjQX0d&
zt%H|#^IlrXQ_F*vcJo?V#yLF^q((Pyr4=-_G)QSTFQqZe+c|h?H}9pDJheP%X*aK>
z6I@MC1PQdPZB)?Y(l91@DUD&?&cRE&c`vQxspUaSdwDG_<D6>_B+#<%Q9%<+0|+pA
zOs1ddQ@sSNMlbKBadBZ%)bgOEy}XtZmpB=ei-VQ+@>W_wQ%i%Q_VS7v!<=glEbvT|
zQoNMT&Ou;%d4X->qn|i9u6+&(6mSMbN#oSs!Df3Fn~iAR=0RwC7NM;Iss+MldlsJ^
z@#^{|NR6JwW|K6vG{|hvBC}<fw{`H@p2cS?cxrjj*`7sbi#YGS4mR7f*ldy}mxeL1
z$ZQ$r?HqizXYttzo?0GswrA1VF|MXxf&`kuQj#>aG{|hvBC}<fw{tjT?OA-bf~S@T
zo$Xn4wutktIoNF9VzWt_SQ<cp$z$U7OR&H*SW3CLFe!3*__kMz&N|N_@!EKOgO>Qt
zA@S1qvf9#byT%`|XwJ2~pXbJ_gJ9_Yl%RG1TH4QRY3nqdLju+4=e4v7rxph*?dPpD
zL33_=fHnGgF|7uw1;R`Fc`psTx}FFUXxYu7f+m-SNy$rT4D+@QUfR!lX(dlB4_eyK
zYiSwhyk~&~T6P<$psA%nO8a>!jbYx-;gGeT_tHw9S{}5tpV!g}uBIn~)ad7}w1TFV
z1}W|5r8I_lI|ncA=e@L&Czppo$!lpD=UsEK(th4bD`;wIkWzdPd{#wbnzMC;k!api
zChK6TrecL;5oo&o6i?7xZ($P!US3XlNfSJ|Jg6#gSyzp5?he9Uo4fP+T1k{kgv<hH
z7@5)2MeRdLvIsPcOd6;c$dR$=ZNzgo5sC|>VQi|9a*=Rc<Ukm<nh~+CbP`3Lb!y5y
zxjaZNaMqdGhI6+N$FI%Z7S%0Z&_I++gz5rlotbHg&Kn_7;8|q1SYY7E<w16VlLuyc
zXHws{B8)_NW~M~FM2?KbdB<0Da`Sb5Xp$)M%o9`M$>l+M5i?`nVzzMZH3;(s(L6H|
zQ7#ec3!r&qCP4QZ#F1=sw?%%76XRWjkYC{BftjudeGSrpA#KHs#rQC;HGapl_D3km
zA}}>F2Ub=mzmS6!SLfH^2p}ajC)E^n81edmBvm&F@vbr;$VpAbERI|b#Wu=W8jV5%
z<6Tjsl%*(UDMOJ<p}|HeOVto}XK*gtj6#;O9Jw6YYLu~5#<)8|GmWB^Dqdaw(W#+P
zqEh_g?hep9qadZq*Vli1HaIDnhgX)ieec2P_%sMXq@@;mo+g(@y^K<nMpteo%-t#a
zWfY}k=E>zzE~7l9G;XTuUeOyRDMe9AHke!(r7}uU%Cg*Tqf|yAN>-d)9ECE<P|DEU
z8y{sdichk^^uoB2Bt<e>ZPRh|YuzY7DN0bvmKKIYE{+lzB`H<VCQ9yIYYbBYEG3Zf
zl`O&3!WgIoR!SphvAk=I!Ai|2St)xRon2$V5@@Eo%{K4%9SbvLsTrj#WlD2CMFB&W
zHX;7C#;_&ev>}<|<dLLuMu|(&g&aMO);|hn<`QVR8A;RR(kP-)?oz%q_j03~MgdG_
zo?ISPHA-Si<EFO&_b$oP(I}0{29pb;wnm9eS(dwP6xS%3$%>PUqqar~O&OZkNWdX1
zx0Q^~ngAJ{k$87Wy#L>flA5xGxw}SfjnbN|IJs-I)+n*5f>x3Q`3wyl!gBYU;t63`
z*m7ZX)+n_pTbR3RJpYU$o2)pwI39mSc}*Fbd#zDgqp&6$Of8H-O<-vgGF1e5Aeqtx
zoNXJKO#c+87l#>%GMYvg^4O7WM<l1`t074uE~O*}oLn5eHOg!<6z5%WjBWxjH7c{q
z3qZL*+G~{IR8e4UEn$*VGm3ID(&W-8uTi2?9_DTx{WS`9GV|o}sIO7JQx50dE}8HI
zT53;;=W;qHPB#Guu;j7pOqjcKG}tKWDSs)wYmNpR<vms5CL6T(QGlsWpxJhm>4U!0
zPrxC}I3`=AxjRRTjY6REm(trgT5OaBmBYE$94$7Agfi0P?SvK^B}3(5?#|I<qkt$g
zPcM&yl01MdZ7GRQCF%x7lZ~RHQpF``a%nW#C@m@v^LCDjQQ#RilwxttHOJT};Idmv
z`VdNP9W$evQD{_rHPbk?cMOaIEo-w8&D}f-Z4?=0D9_tH`fQXLRl%EjgZEy?xG2zU
zvr?wXy`#)VSyB1Y+^wU}Mxjw=o?ISvHcF4m;k@@cfJn=3D&tcWQ0kn0fkT)$CetTm
zeNUm!MuAfKOX*#6)Y&LmstUIyVZOI^00J%fIn~0Qb=<iUO&&_hY?LvTFU`H?=(ADK
zl$ocO$3aPcLUqQqbgwzuY!o<^D=yIF(kQc0=2RZ$?i_tKN}@9J<npMqQ9f0S!~A=`
z8j>%K*U$TFczyZDr?5sjRjD7}-49x7lvFi?pq9X>DgZKfknkN(_1_w0ZWVBudq@>d
zFAnWbr&w*<7!ovZyO?GLU`9-<fpUTL(sZs>TwvMU<Qp<Tq}gVrOp{Bal&15oW|rn|
z9lbQ2b7kSl<xxx1c~>)V-g_26pe28gDo>M3qm-s|uV%vBouikggRd++y*v&|?wB>K
z@$$)_zUFA9>Da48aWy>=9Kdo?8kgqo9KAH1dNs4Sz>_;iEluZM&A_?W91mI3saJWL
zTpFb`oq07A=I$K5G@W{7;mPGuOVhbmC0z0F%x8@Nk!FyT1WhlEW0F%;uDK%2-8njI
zI`wMirSz^j3T!&}s(_1*>U^zZ>J@NVMM>k--ZAtFwCtcnG;i~md(}(_U(LRn&h7z4
zUe;&Frxe&JYyJKiK%gbRBp|M?u3v&fm~B$F!Q8E*&!&^FW)|n|9CbFGeKiB;z1IOm
znm8wcCa*fmY&!R9W@+xu(Pz`aR~DXJ9(6XIe3ihJj8FDUfI!RJ?0AKjXmV+k*>vXB
z%+lPQqtB*OuPi*dJnC#Z_i6^tz2<1M>C~${O)rgOk~?Az-|#$r=<%QZ14W*3Lz#@_
zlqZ)*olWOn6|N;fu5(DdG+#e^vk9*+|M(QvbgotD$9JzJ9;BwTtkMLz1Zru>uUbC5
zxNn*=KX>%t{%JDy{>5uwxPNVR88+;$ghCpEs}zu05Cf~g8MWNN@$Ldlssd)DpmvO2
z4B8(us!SAbcbH8D&}iK&Ah{r_W=N$naM~ES{j2UCQ8hzEl?sy!qh^M@DuLzh7d114
zSLrypIBI4{vNF@Wdp2XNnjyzZg~^3cF++-#z;btuiWxGjbevusXJo9HMjG>lr#Exi
zX_;XKpmh`B;w}K?0;!oH$4V%0`ZdaWffUUUWaR<M1yVIbl9h>vozlNs^4v58S*b9&
zFluJVu@YGBbw|w%F;+TGE{@U}Qmo80_ijlA4Ix%4OfQTRGCndDSnjS-O+$p0j+2X{
zpN0f0BdxG`$(so?tN^r%YCC+V)pUDS3uKHHcvfSb@cg?abF7*n$V#}PQ=nWR9X2Fc
z8F&k|+!7qX#%HD~OfHNb8*;1!mUrDT#R@p<piHh~!0E+tO2*%(YA(*ZTQb55Fsq@A
zOLHC$00Y+ey?>!Fch{)0A-_t;$z7w#hV&{Ut>$jY^U@Gt<pIhC(q%)2mB5>BUi$5T
zG8;myJV3cXx@<_XGV+!XQ@y@;_{gC;=RiP`tL!c*u3j$ZXMB*Vm+L)&0FYE=zu8u7
zpvtAvc0<aQr&RY2MDGoGS2ehDxisLAd}YGcc}M~WFtY+LK-#8HWO8Bj-;jAF6y|Op
z1vrFW={U7GhF*bYwcqiD^xpQFcm<e!&tcb`w|!v18h_8K`7-)%8%)0f&c4HpOLNK^
z!?5seZKE&XH5OOzw$E%V0F^N*0Ohfy`-ZG6p}-pb<3Vmn%klu_0%^b@FU!E|)Otu_
zR8})&WvMW^F#2!E$r4!J-I6(3z*)DH$#wLuJKAqZ$}$({-Yu!VAtp<O$z7xMhKwwM
z<?b4#H-u#AIJr2gZb->8(h8fG{0spWfM&_%P!@o4fz;cOlO+^bW4EN*hNvtLP%e;4
z8`83jylwvOf3qA&a+TdB#nsE@{Fv>Q^Ro>Lvv7S#13*%(3X`VFrNVac<*ZhdiJ;CM
zh)8m+LQ};R%O#Om`#^1bZc(=(asV?caHqyY5;0`ihi6=vyL}MgR>;#bZ3qRN+%?!Q
zX!c<t&3W4=hcInKd=E*)fMq`P9Z!h+wm}kc+K^0fa@SzJ_%3Gq)-w4Nw|gj|z5ps?
zQs=f$6nV;|kT<=c`i27iZH4qK<EkqF<qpDt!L!awx6VTnIe?9S`O8pT0VWp){{_oB
zDHE1^Fd@LeSqG(r(~IMjnC*8=bMKZ=U%<2x@$QzeUVPiSbwimH=I$D#7dUN5ri)wd
z8dMiFWkw2{m%f?c2)7kdvW)Ak0F(=a-h!u0N@4TT3#8hHlq{h@cLy1g1)ef5+nyTS
z4@pFlt4vOctM-#YTCEVIWxTxGurM^wLlS|cDpM3w<x*+8h?HR7nx|Cn4g@5*&WZAV
zLWV?Q?E^KIduKzbz#)yxbZxi6<UZ1WLnN2La<`8H95T9eoLn63H-vVXY2MpDID}~<
z;=Jtx1D5&oyWP=a^mG@OZyO*HXWwDQ51U~~<gU?rL!g(jYlXY8ULMsqBzk#(a)EST
z{Bn?a@0|RmpL^Be@V6CWy@Ud%10xqm0}jz%)x5Q18{Rt}tv6(O8D3xh@oBjsyleK$
zyBk8Y4T)VQf?fhUo;$8RvbJ_;ZgqKM-P;HHYe?s+E~7+|OQEyk;BfY8*@Ad`0t7i3
zvxqoqIUo>|Sp1`}5S#vz-M;mW<s;MIvh%tjK%k}0TMab1G>RxTwzyNXVeXF6L2)ge
zJh?ooXNcn}PW!C#=ACPcX<VR*a}sEJX&jT(q3i6@+?}I!hBz+eQU;z}9<4J3a;0$i
z@ZxFD_aKkZ#C|fqrWUANC>=CJb2$oix05;=Lc07|xma3hi0d*^w|!}d@6il6FjFV8
zl}#-(<uYlfA-GE~)7?@!YDn;MQ}t3gN2z1l{MB{eA?d0i&r2;dFy%6-uOZY+N4h&o
zs}1R1ZmL`=MK{EJS)q6)uJfHm#DSS|-|8tUG37ET!XfZWFVo#onsG?}a#Ph(836{U
za$DQ><s-gbmpNcyR=QB)R$6B2Wnzdz7?^r(r>BhaNdz;&K&>?B*@cEVl1rsVhiEWM
zsl_k7b9R(29n!&^Ou0-7bqEO4k=`}Rps-d53UgECQmNM=Ff0!poi5e>9%=FbO?|c@
zwe6;%a-nqW5FO?y)ZJ8ScE}L(WA$P=N~uHN$_rd}Pklq8V~0R7W2I%LTqYGego^1%
zcS~v6AzaK&l}n{$hmbK9>h)dVz_dE{<nu>$Qc6f8(s)-ZBi$XPWrxr){o)3y+)+w)
z2p_XTOYW(!Q95=AAagSHGC4$6#U6Kbt4w!CY1tu+OfS^iQA&0QB~zi^HOgQzFpZ8~
zyu8k5QihWOX;kbo(%VsHl(j-enSODd9c5G*pvLV8W9#}<1^}Q`A5IrlK2X$kh}S2?
zI4INi^W-gtab(~ar(?>|%i%awA46}yW{JcOXQL>~A&o3IgH=$uP`Yx6By%v`%Z<7m
zvdR2dxmenA2q`mC%MPgzEKj6EP??h{mq}F)VP!hf-BQ|e2rY9{<x(ljA-qh5y6=f}
z<q%-zWa?#dh-|7nN#H<vA<)m=QQC3{Gt&$8-V-UyA=JzY-4Z(I#!FWY!DddT+)=7>
z2shJ_?vC<6Iz*nisdA~5<q(9XLfvbWt{kG#oJ_S$hNS^%GN8#CoeZx3oR*nspqhNB
zlB$=AF$&RXmaBU@&S{yR2ByiCS!jYOmq~FB0cv`g-Zjb~HBilSR4mlF>w-t9ag0U}
zUG>jt!2{GP*=eZUQ#y8tR?GJ^V!E43%?=rBeym(9Ejxs*8L72>7dSGNNh&k-GC4=K
z<LF5W!NiXIzu5sGsm2Xyd`=Nj<x(lxA%sossP|q99GS)qsaR-W%4JfqLl~Q0rhAd{
zKsto7xv6rgl<W}BW`$0)!I|k!9GI!!)r;T1nJkzRQ!bN=9YWglGTj}eWrwgfH&ri{
zV`Mv$E?-^u8l_{0;5M~T?;53IhwwHX>Fp>p+(4B!qsiz_EnDp<quc;B-qV0{ozpVE
ztrhay@)tBv<Wi{7A+*f`arcSt9OBvpj$969IRv%kFs<=73Dx&j@CY@HLlsmnl*3_j
z<pk5+M(T3NZgX7UfR&4-Er$>{Bem?1@;NPVWSTaXaDmHAxlF2Z2y@fRbhnfT(jnB%
zO_fWfEQfG673#ew0tcq;6Vm05I`4^~5$PJD>D@^mUD|R8d($tjca2h(L+G0oI=K;h
zy_LalV45DeCLJv?<-w&YhwwMOO!pe4Er&=rH&rf`vK)fqRH%E6@<2L7$2pmDnN;Ou
zm|XnuQy!ukYm~N}OqLU<YN?Ew1Jt}76^{~q7R$^zFwJt5VCrRJh>`(x>1$g=de<n^
z=s-1ZMoFq%Dz!QpO_wS)I=*zal#-o{sIx=mLh0Das5%qV-Bcb)CnM_=tXwQDI~iS9
zNiDmlEDcPsYb6ux%uKmVDt0o+&VqEel$M>$vlFO#shlI*k#uVsnlyRC-nmTa*vVwO
ze4$R?1r126VwaDG-5sT6Cll^0g$AnJQA&0)=B^rAa!-AY(y^0icV?zsCKWpwdS^kp
zJ4(w=X5R@^xl~GaG6F9T_4+PwWE$;gykI6%<Xxjw>|`LGrA&86Y1zqKJb@~Al#-o{
z#~VQptgJ2`+Q)C`;df!y`pr(};(=6U*3~4nBnILER5dV3khf>d!fPe7@TxDLb2kNr
zs496*l;vG<%)o0UGw`a4a;`Hl2ogrZ|1`<;JFtYQn6l&nq!A~h@5<ND-67g<GViW}
zAeTVZO@`dX0QcIV+a{CkDj9Ma6xn2)-3W1bWog@2hhUeh|3(L3AXZvzo+y_{g-r(9
z<x6yTjt-m5v@<j1GAXdhSi3Z`_QM|Y>;SCVFQP`4fqH=)oXXED&*Hh;NO4W3*k!M&
z_Z~)RO-9&di0+M&!kWykvjOEUQdX0pby=Rfi<H%5be$C`7fDe~2G~^)CpzwZ69cS9
zD;Z&D1F8iw#16bFFNxRUR)O9wGRh7}`*sr{^&&Ah$w0fzMfI*lM%n>bC9u_3)oGUu
zwZk{bRDHwTES|fI6xw90UG|zfcT~Vgw9L(EI@$&d#LACVY`?5&qCBXS-DI>~bWtOq
zyOH$YWWt@9DVIqdP6pkjk!9hKkIEQz2cUU$%0Rt94vtBU$EP@0_u<6B-A2lBGV(5a
zO}*=pYMczd%MjhWD#bV%eP;v8U8EN$!|$>@cNcl4n~cA+BIP1!#>oJ@3gXt?*!M;R
zSO8Xjyg1rar{9njlu9m;UYrcT%NFSFBGovVeP>1LMRIN`KUG|PQQd2iVw_C8%arIf
z9C$Il(#>Kx;~jM}rqiP>^`=1@krp+sK#E-?jKr$j)UhMmj&Me~m$eUZNCUF+i<b>Z
zxk$J%h*fSpo9JAT<N;bK$R?;<C=40Qs`-hTW$Vx<7KRL>eSXSBy+lq=r6A8N(cMZo
zGMM)H*;Zs=%4I^4@!gJ9o6n3go!cv6Bsyltx4)7_pks7uY_K>$o=;}tjqV{cFRQnc
zuw)>G;hEe=oCi9Dfmrz=h(wVkqCBjyWB`TH*#f%PBup7hVR|B_UM7d9@(_8(b#<>v
zm@<ft+3~JPm@>XiR@inXSGC5Pge`;Vn4T$9?j&p(NaOGnKPc&X1R)H>N?R_?PJ$?x
z2wMiwI66h3dri`qlNoxc-8IQn%Vdlmkfz~@&aZr}WP%=$rr}8=^&&Am$?&}Rx<*8A
zGnt(SsA+(zpmL#f=45nUMWNe7%h^mCb22o~NR&&YFel^k@<4Yhai~0>OwBVh<ua+v
z$?&`!(tEE23`CnTkLGncJI;G0SOnU)o^3#PC+W<|B)$A)4NSR{ROV!$UKMhp1YUn-
zrXGm)ou}q`9a#X{gn2wVTP3<XNoP(b>*X)2x06)nWVl`q>0XmKTb@s*>lumiuu_<l
z@p^fnyOVV0WaOTiDVIrQPR8)XNQ|oIxBvO_WSBH3F!VAw9yTeSesijPL<ImTrtM@(
zkxQZaCZqVu9SL80R{tE8seC{h?mtyXwMYi^foK?s1kt;^nA103ozu7N$*CGD7fSc7
zj_Zq=(_6duo(C9+meb-;Wtk|KNbyA=T17(O%o5$Lr28U-s0LFnlhb24E-qhJ?>!GN
z5)Io;1yL@M;)^dR*RSgI{6yYUfg{r~JyWJ!Ce^n(z%MbdTaxCx0|EwO<wq*=M7czY
zFTQ5X-pQ>J-D{HWTOHdsv&g`d%cS~N$M(%Yy4NJlw>q*fPt;4~=va=5C!4^elhgH~
zuSvRZbzI-fBE6lY`c}vFm5{{)Kj)gH`Bq2u<%x19sloWrv29HjOLVVE9u`-}^vx{N
z+eu1tbxdCY8J+eyTS;ZcrvWTiHjPw!2^fi{yH-T>HWNHTO#@T~)eFS{;gr#QY0EY#
zZ$5yLXu5Tcsw@-b5-H5p(R}g1ZUeeoNoTH(=Cd&6GO5hifAmgz@0EaoXfx(wCj(I~
zk-}UZ$v5Mw`tOxY<eRTf<g+m4GO5hffqV(%L<zjllz@>a%ub1@m&nnv3G;Y#wt((6
zNoTH(;hT9`o%c%6fV3PJpQdk9C5%M-jca-$?p>1<=IZ#p8CSK&n&e?|b^M-%DGw}_
zxjKHYfW)Y}&QbC5U|2p<(VPnmxeQ8ib+lfo4EIW-30H^dr73zT9D?zY@ao#p!;6Pj
zUb?(*W%bChjrF<x#||7=KDuvh|N8RL>lZgx)>ike99}%QJV#yGr#lOsH-zSiEV(SY
zZivdWqd?KCVnQBJqrE~ysYNju52Vrk#X?%&SlmFb7GLZs3S$w{@HCWM6h$^f;0dG^
z$NVt|Z$9MUX(+iUifahN+ahf&9zD3cwX7De<%)9lhngBv@GJnm0QP?@rP<ea^@0)_
zqV6mJxc~}hh_ySmx<yz!wy}L{tJMPvW{9w}A><+`lp(IpKydYeXRaZg&W4bSpfHZY
z=x(_9;?LfEPXDd9{^?-H^glO#|MY*3{zv+spLgPgC(ro1>xNIf^xIDT#$`u`TYo40
z@9%uW`KMYBUN^k(8UNjXa{Ea3-~YuQ`qfjPY&~W8zR!J6|B=6W%J47W^vcbHAKGzg
zj#K{gbIyEe@4fduW%#B`|GIbfA3Y_@{pZJ?-+k8eA09mMGjBifeJ^=<aP+6&*n4R0
zJ=wok<~Cn&dGFw54-a1ct{?CG^y?nZUf;ccJb$8p;?BYMUhoh1yyV_H2k(9TCH>3K
zyiw-1e(jkjUUDa1=f@w@e$3nM9Nc-?f8YD^SNzKGvb!$qUGi7A4R3t#-@Bh~{@38l
z{_FZ{PkUhSiBtc(dCfDf8eZ1>?&HhNtA~%A^`m>g@jw22@Ua`3{pVfsz~GjJ=WqVb
z6^|Q!;R9cH&wsve=Wyp&cJKYl7kiif>&g=bYk##T?nnE&dw1Ob;lVRHJNj?y-y7Ba
zHTPWczVG<RpI-X3pZ}5fKJv*2FTL-U&YtG8zG3)HPy5@GuYA|e;me+R*1daff6w5<
zKX~KGi@r1d`yV^|TTdP8{LbL^EC1=9m9LBc{Gzu!wE6D49v<9r;^%s2U;Kr^YoGMH
z{Wm|db9n35e(L0ZzGCO_)nEUdQ-iO6&*0Oy+|YacBb(9B_XnSP+=;nY{=(pkH@)k`
z<1V^6`tx(2``UZ2UV8R$>9Y5py6M`hhVTFNjXmccKX>^4x6JjQ@XnpXFP!<B&3$iv
z`Jnq3{oZ-)`0vmC(FLbo`n^9txcEa~>|DPm{`0F2p0nqZAA0-XPv(FB#NXU==iv5d
ze(UB{w_h~8@-uJj{qQgB9NuyMQ}?`L<>tYw-t)}kzx(9)?;kzqFZ)mao+}2MFFyD9
z`Z;@}pMU3lgHs>8|NP;tcYf^Tm;dd&;m4kJ<DToUJa^c?^`Y)Z-?Vdh##skXt^Z4R
z@a<<jbo^&OwI}-f&pF=Pe8(FHgVhiHS?B5h9RK<Km+aUx_x^VczUvR)**pKDxIc}L
zKREca{#}D#dCif%-+29%!yi8O)XhsidhYPctsm&^zkBEKw?F>*{>E9C4ZeEACCAVI
zZ2bDYrO%#v;y-i-H~qwabdG-}{`3EO)<;i0bob+jC;NAGZ(83u{Kz|g?9`9_=3sF4
zH(%ZP$RqK;f9oR`?fJkJgTbwT{L$_WpO621>5dEczV)4tAO65~_n-X4%XbdXyZP?^
zxl3O&_@$TL(fO9|KQp>M|8~J|ZGP!BHx2gx(vNlryW;*l_k$N~?)vlB5ANUHJ9*Z_
zw+%kN{OrAR-*w*boU0$(|HbF;9DdV%PrCO@KlQS~#$6YkxH!?@H(Yhi=1+e6ErZYf
z-9MgafBB|}fAr9a&2wLL>)@wecXe;=?{6Dic-`%LZ~5HW!&iUg*}Z)`&l{e#c;%j_
zzx&d`cf8}C?pX`*d|dacAKLqmuRSq1<NU9l-1o)J==GN_{`BSt9{=INr4N2>|6i|p
zVDREwcHev3+b<a2{Ms+|?!Is5aPYb(ZhroZcMmQ(^9h~bI~w=@%=R6pzH;lI430nj
zw%$#jh+qG}lP=tR>I3%={^F{;dpCT`1A|BJPm?>#_`dq+{pr#B)1&vNNAFKNc8vFf
zNAFL?{qg_%{sj8ud1>_DZ~x%`T=DeVo_u-pmrq@O%gXaF-*MwtuNeP?zkK;w=bsn-
G^Zx-TuGq@}

literal 0
HcmV?d00001

diff --git a/tests/scripts/save_policy_to_safetensor.py b/tests/scripts/save_policy_to_safetensor.py
new file mode 100644
index 00000000..70337c17
--- /dev/null
+++ b/tests/scripts/save_policy_to_safetensor.py
@@ -0,0 +1,101 @@
+import shutil
+from pathlib import Path
+
+import torch
+from safetensors.torch import save_file
+
+from lerobot.common.datasets.factory import make_dataset
+from lerobot.common.policies.factory import make_policy
+from lerobot.common.utils.utils import init_hydra_config, set_global_seed
+from lerobot.scripts.train import make_optimizer_and_scheduler
+from tests.utils import DEFAULT_CONFIG_PATH
+
+
+def get_policy_stats(env_name, policy_name, extra_overrides=None):
+    cfg = init_hydra_config(
+        DEFAULT_CONFIG_PATH,
+        overrides=[
+            f"env={env_name}",
+            f"policy={policy_name}",
+            "device=cpu",
+        ]
+        + extra_overrides,
+    )
+    set_global_seed(1337)
+    dataset = make_dataset(cfg)
+    policy = make_policy(cfg, dataset_stats=dataset.stats)
+    policy.train()
+    optimizer, _ = make_optimizer_and_scheduler(cfg, policy)
+
+    dataloader = torch.utils.data.DataLoader(
+        dataset,
+        num_workers=0,
+        batch_size=cfg.training.batch_size,
+        shuffle=False,
+    )
+
+    batch = next(iter(dataloader))
+    output_dict = policy.forward(batch)
+    output_dict = {k: v for k, v in output_dict.items() if isinstance(v, torch.Tensor)}
+    loss = output_dict["loss"]
+
+    loss.backward()
+    grad_stats = {}
+    for key, param in policy.named_parameters():
+        if param.requires_grad:
+            grad_stats[f"{key}_mean"] = param.grad.mean()
+            grad_stats[f"{key}_std"] = (
+                param.grad.std() if param.grad.numel() > 1 else torch.tensor(float(0.0))
+            )
+
+    optimizer.step()
+    param_stats = {}
+    for key, param in policy.named_parameters():
+        param_stats[f"{key}_mean"] = param.mean()
+        param_stats[f"{key}_std"] = param.std() if param.numel() > 1 else torch.tensor(float(0.0))
+
+    optimizer.zero_grad()
+    policy.reset()
+
+    # HACK: We reload a batch with no delta_timestamps as `select_action` won't expect a timestamps dimension
+    dataset.delta_timestamps = None
+    batch = next(iter(dataloader))
+    obs = {
+        k: batch[k]
+        for k in batch
+        if k in ["observation.image", "observation.images.top", "observation.state"]
+    }
+
+    actions_queue = (
+        cfg.policy.n_action_steps if "n_action_steps" in cfg.policy else cfg.policy.n_action_repeats
+    )
+    actions = {str(i): policy.select_action(obs).contiguous() for i in range(actions_queue)}
+    return output_dict, grad_stats, param_stats, actions
+
+
+def save_policy_to_safetensors(output_dir, env_name, policy_name, extra_overrides):
+    env_policy_dir = Path(output_dir) / f"{env_name}_{policy_name}"
+
+    if env_policy_dir.exists():
+        shutil.rmtree(env_policy_dir)
+
+    env_policy_dir.mkdir(parents=True, exist_ok=True)
+    output_dict, grad_stats, param_stats, actions = get_policy_stats(env_name, policy_name, extra_overrides)
+    save_file(output_dict, env_policy_dir / "output_dict.safetensors")
+    save_file(grad_stats, env_policy_dir / "grad_stats.safetensors")
+    save_file(param_stats, env_policy_dir / "param_stats.safetensors")
+    save_file(actions, env_policy_dir / "actions.safetensors")
+
+
+if __name__ == "__main__":
+    env_policies = [
+        # ("xarm", "tdmpc", ["policy.n_action_repeats=2"]),
+        (
+            "pusht",
+            "diffusion",
+            ["policy.n_action_steps=8", "policy.num_inference_steps=10", "policy.down_dims=[128, 256, 512]"],
+        ),
+        ("aloha", "act", ["policy.n_action_steps=10"]),
+    ]
+    for env, policy, extra_overrides in env_policies:
+        save_policy_to_safetensors("tests/data/save_policy_to_safetensors", env, policy, extra_overrides)
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index e50d4108..22b271be 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -265,7 +265,7 @@ def test_backward_compatibility(repo_id):
 
         for key in new_frame:
             assert torch.isclose(
-                new_frame[key], old_frame[key], rtol=1e-05, atol=1e-08
+                new_frame[key], old_frame[key]
             ).all(), f"{key=} for index={i} does not contain the same value"
 
     # test2 first frames of first episode
diff --git a/tests/test_policies.py b/tests/test_policies.py
index 51cdb93e..7d2f19ba 100644
--- a/tests/test_policies.py
+++ b/tests/test_policies.py
@@ -1,8 +1,10 @@
 import inspect
+from pathlib import Path
 
 import pytest
 import torch
 from huggingface_hub import PyTorchModelHubMixin
+from safetensors.torch import load_file
 
 from lerobot import available_policies
 from lerobot.common.datasets.factory import make_dataset
@@ -13,7 +15,8 @@ from lerobot.common.policies.factory import get_policy_and_config_classes, make_
 from lerobot.common.policies.normalize import Normalize, Unnormalize
 from lerobot.common.policies.policy_protocol import Policy
 from lerobot.common.utils.utils import init_hydra_config
-from tests.utils import DEFAULT_CONFIG_PATH, DEVICE, require_env
+from tests.scripts.save_policy_to_safetensor import get_policy_stats
+from tests.utils import DEFAULT_CONFIG_PATH, DEVICE, require_env, require_x86_64_kernel
 
 
 @pytest.mark.parametrize("policy_name", available_policies)
@@ -228,3 +231,37 @@ def test_normalize(insert_temporal_dim):
     new_unnormalize = Unnormalize(output_shapes, unnormalize_output_modes, stats=None)
     new_unnormalize.load_state_dict(unnormalize.state_dict())
     unnormalize(output_batch)
+
+
+@pytest.mark.parametrize(
+    "env_name, policy_name, extra_overrides",
+    [
+        # ("xarm", "tdmpc", ["policy.n_action_repeats=2"]),
+        (
+            "pusht",
+            "diffusion",
+            ["policy.n_action_steps=8", "policy.num_inference_steps=10", "policy.down_dims=[128, 256, 512]"],
+        ),
+        ("aloha", "act", ["policy.n_action_steps=10"]),
+    ],
+)
+# As artifacts have been generated on an x86_64 kernel, this test won't
+# pass if it's run on another platform due to floating point errors
+@require_x86_64_kernel
+def test_backward_compatibility(env_name, policy_name, extra_overrides):
+    env_policy_dir = Path("tests/data/save_policy_to_safetensors") / f"{env_name}_{policy_name}"
+    saved_output_dict = load_file(env_policy_dir / "output_dict.safetensors")
+    saved_grad_stats = load_file(env_policy_dir / "grad_stats.safetensors")
+    saved_param_stats = load_file(env_policy_dir / "param_stats.safetensors")
+    saved_actions = load_file(env_policy_dir / "actions.safetensors")
+
+    output_dict, grad_stats, param_stats, actions = get_policy_stats(env_name, policy_name, extra_overrides)
+
+    for key in saved_output_dict:
+        assert torch.isclose(output_dict[key], saved_output_dict[key], rtol=0.1, atol=1e-7).all()
+    for key in saved_grad_stats:
+        assert torch.isclose(grad_stats[key], saved_grad_stats[key], rtol=0.1, atol=1e-7).all()
+    for key in saved_param_stats:
+        assert torch.isclose(param_stats[key], saved_param_stats[key], rtol=50, atol=1e-7).all()
+    for key in saved_actions:
+        assert torch.isclose(actions[key], saved_actions[key], rtol=0.1, atol=1e-7).all()
diff --git a/tests/utils.py b/tests/utils.py
index f3fe5790..6a706694 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,3 +1,5 @@
+import platform
+
 import pytest
 import torch
 
@@ -9,6 +11,51 @@ DEFAULT_CONFIG_PATH = "lerobot/configs/default.yaml"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 
+def require_x86_64_kernel(func):
+    """
+    Decorator that skips the test if plateform device is not an x86_64 cpu.
+    """
+    from functools import wraps
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        if platform.machine() != "x86_64":
+            pytest.skip("requires x86_64 plateform")
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+def require_cpu(func):
+    """
+    Decorator that skips the test if device is not cpu.
+    """
+    from functools import wraps
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        if DEVICE != "cpu":
+            pytest.skip("requires cpu")
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+def require_cuda(func):
+    """
+    Decorator that skips the test if cuda is not available.
+    """
+    from functools import wraps
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        if not torch.cuda.is_available():
+            pytest.skip("requires cuda")
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
 def require_env(func):
     """
     Decorator that skips the test if the required environment package is not installed.