From 058ac991ebb00599a20ba6d8918d90f6928f1c60 Mon Sep 17 00:00:00 2001
From: Simon Alibert <alibert.sim@gmail.com>
Date: Mon, 25 Mar 2024 16:35:46 +0100
Subject: [PATCH] Add simxarm back into tests

---
 lerobot/common/envs/simxarm/env.py             |   2 +-
 .../replay_buffer/action.memmap                |   3 +++
 .../replay_buffer/episode.memmap               |   3 +++
 .../replay_buffer/frame_id.memmap              |   3 +++
 .../xarm_lift_medium/replay_buffer/meta.json   |   1 +
 .../replay_buffer/next/done.memmap             |   3 +++
 .../replay_buffer/next/meta.json               |   1 +
 .../next/observation/image.memmap              |   3 +++
 .../replay_buffer/next/observation/meta.json   |   1 +
 .../next/observation/state.memmap              |   3 +++
 .../replay_buffer/next/reward.memmap           |   3 +++
 .../replay_buffer/observation/image.memmap     |   3 +++
 .../replay_buffer/observation/meta.json        |   1 +
 .../replay_buffer/observation/state.memmap     |   3 +++
 tests/data/xarm_lift_medium/stats.pth          | Bin 0 -> 4306 bytes
 tests/test_datasets.py                         |   3 +--
 tests/test_envs.py                             |  17 +++++++++--------
 tests/test_policies.py                         |  10 ++--------
 18 files changed, 44 insertions(+), 19 deletions(-)
 create mode 100644 tests/data/xarm_lift_medium/replay_buffer/action.memmap
 create mode 100644 tests/data/xarm_lift_medium/replay_buffer/episode.memmap
 create mode 100644 tests/data/xarm_lift_medium/replay_buffer/frame_id.memmap
 create mode 100644 tests/data/xarm_lift_medium/replay_buffer/meta.json
 create mode 100644 tests/data/xarm_lift_medium/replay_buffer/next/done.memmap
 create mode 100644 tests/data/xarm_lift_medium/replay_buffer/next/meta.json
 create mode 100644 tests/data/xarm_lift_medium/replay_buffer/next/observation/image.memmap
 create mode 100644 tests/data/xarm_lift_medium/replay_buffer/next/observation/meta.json
 create mode 100644 tests/data/xarm_lift_medium/replay_buffer/next/observation/state.memmap
 create mode 100644 tests/data/xarm_lift_medium/replay_buffer/next/reward.memmap
 create mode 100644 tests/data/xarm_lift_medium/replay_buffer/observation/image.memmap
 create mode 100644 tests/data/xarm_lift_medium/replay_buffer/observation/meta.json
 create mode 100644 tests/data/xarm_lift_medium/replay_buffer/observation/state.memmap
 create mode 100644 tests/data/xarm_lift_medium/stats.pth

diff --git a/lerobot/common/envs/simxarm/env.py b/lerobot/common/envs/simxarm/env.py
index 9c996139..9b08be6a 100644
--- a/lerobot/common/envs/simxarm/env.py
+++ b/lerobot/common/envs/simxarm/env.py
@@ -84,7 +84,7 @@ class SimxarmEnv(AbstractEnv):
         else:
             obs = {"state": torch.tensor(raw_obs["observation"], dtype=torch.float32)}
 
-        obs = TensorDict(obs, batch_size=[])
+        # obs = TensorDict(obs, batch_size=[])
         return obs
 
     def _reset(self, tensordict: Optional[TensorDict] = None):
diff --git a/tests/data/xarm_lift_medium/replay_buffer/action.memmap b/tests/data/xarm_lift_medium/replay_buffer/action.memmap
new file mode 100644
index 00000000..c90afbe9
--- /dev/null
+++ b/tests/data/xarm_lift_medium/replay_buffer/action.memmap
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:10ec2f944de18f1a2aa3fc2f9a4185c03e3a5afc31148c85c98b58602ac4186e
+size 800
diff --git a/tests/data/xarm_lift_medium/replay_buffer/episode.memmap b/tests/data/xarm_lift_medium/replay_buffer/episode.memmap
new file mode 100644
index 00000000..7924f028
--- /dev/null
+++ b/tests/data/xarm_lift_medium/replay_buffer/episode.memmap
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a589cba6bf0dfce138110864b6509508a804d7ea5c519d0a3cd67c4a87fa2d0
+size 200
diff --git a/tests/data/xarm_lift_medium/replay_buffer/frame_id.memmap b/tests/data/xarm_lift_medium/replay_buffer/frame_id.memmap
new file mode 100644
index 00000000..a633d346
--- /dev/null
+++ b/tests/data/xarm_lift_medium/replay_buffer/frame_id.memmap
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6afe7098f30bdc8564526517c085e62613f6cb67194153840567974cfa6f3815
+size 400
diff --git a/tests/data/xarm_lift_medium/replay_buffer/meta.json b/tests/data/xarm_lift_medium/replay_buffer/meta.json
new file mode 100644
index 00000000..33dc932c
--- /dev/null
+++ b/tests/data/xarm_lift_medium/replay_buffer/meta.json
@@ -0,0 +1 @@
+{"action": {"device": "cpu", "shape": [50, 4], "dtype": "torch.float32"}, "episode": {"device": "cpu", "shape": [50], "dtype": "torch.int32"}, "frame_id": {"device": "cpu", "shape": [50], "dtype": "torch.int64"}, "shape": [50], "device": "cpu", "_type": "<class 'tensordict._td.TensorDict'>"}
\ No newline at end of file
diff --git a/tests/data/xarm_lift_medium/replay_buffer/next/done.memmap b/tests/data/xarm_lift_medium/replay_buffer/next/done.memmap
new file mode 100644
index 00000000..cf5e9cca
--- /dev/null
+++ b/tests/data/xarm_lift_medium/replay_buffer/next/done.memmap
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dab3a9712c413c4bfcd91c645752ab981306b23d25bcd4da4c422911574673f7
+size 50
diff --git a/tests/data/xarm_lift_medium/replay_buffer/next/meta.json b/tests/data/xarm_lift_medium/replay_buffer/next/meta.json
new file mode 100644
index 00000000..d69cadad
--- /dev/null
+++ b/tests/data/xarm_lift_medium/replay_buffer/next/meta.json
@@ -0,0 +1 @@
+{"reward": {"device": "cpu", "shape": [50], "dtype": "torch.float32"}, "done": {"device": "cpu", "shape": [50], "dtype": "torch.bool"}, "shape": [50], "device": "cpu", "_type": "<class 'tensordict._td.TensorDict'>"}
\ No newline at end of file
diff --git a/tests/data/xarm_lift_medium/replay_buffer/next/observation/image.memmap b/tests/data/xarm_lift_medium/replay_buffer/next/observation/image.memmap
new file mode 100644
index 00000000..462d0117
--- /dev/null
+++ b/tests/data/xarm_lift_medium/replay_buffer/next/observation/image.memmap
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d6f9d1422ce3764e7253f70ed4da278f0c07fafef0d5386479f09d6b6b9d8259
+size 1058400
diff --git a/tests/data/xarm_lift_medium/replay_buffer/next/observation/meta.json b/tests/data/xarm_lift_medium/replay_buffer/next/observation/meta.json
new file mode 100644
index 00000000..b13b8ec9
--- /dev/null
+++ b/tests/data/xarm_lift_medium/replay_buffer/next/observation/meta.json
@@ -0,0 +1 @@
+{"image": {"device": "cpu", "shape": [50, 3, 84, 84], "dtype": "torch.uint8"}, "state": {"device": "cpu", "shape": [50, 4], "dtype": "torch.float32"}, "shape": [50], "device": "cpu", "_type": "<class 'tensordict._td.TensorDict'>"}
\ No newline at end of file
diff --git a/tests/data/xarm_lift_medium/replay_buffer/next/observation/state.memmap b/tests/data/xarm_lift_medium/replay_buffer/next/observation/state.memmap
new file mode 100644
index 00000000..1dbe6024
--- /dev/null
+++ b/tests/data/xarm_lift_medium/replay_buffer/next/observation/state.memmap
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52e7c1a3c4fb2423b195e66ffee2c9e23b3ea0ad7c8bfc4dec30a35c65cadcbb
+size 800
diff --git a/tests/data/xarm_lift_medium/replay_buffer/next/reward.memmap b/tests/data/xarm_lift_medium/replay_buffer/next/reward.memmap
new file mode 100644
index 00000000..9ff5d5a1
--- /dev/null
+++ b/tests/data/xarm_lift_medium/replay_buffer/next/reward.memmap
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c4dbe8ea1966e5cc6da6daf5704805b9b5810f4575de7016b8f6cb1da1d7bb8a
+size 200
diff --git a/tests/data/xarm_lift_medium/replay_buffer/observation/image.memmap b/tests/data/xarm_lift_medium/replay_buffer/observation/image.memmap
new file mode 100644
index 00000000..c9416940
--- /dev/null
+++ b/tests/data/xarm_lift_medium/replay_buffer/observation/image.memmap
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8fca8ddbda3f7bb2f6e7553297c18f3ab8f8b73d64b5c9f81a3695ad9379d403
+size 1058400
diff --git a/tests/data/xarm_lift_medium/replay_buffer/observation/meta.json b/tests/data/xarm_lift_medium/replay_buffer/observation/meta.json
new file mode 100644
index 00000000..b13b8ec9
--- /dev/null
+++ b/tests/data/xarm_lift_medium/replay_buffer/observation/meta.json
@@ -0,0 +1 @@
+{"image": {"device": "cpu", "shape": [50, 3, 84, 84], "dtype": "torch.uint8"}, "state": {"device": "cpu", "shape": [50, 4], "dtype": "torch.float32"}, "shape": [50], "device": "cpu", "_type": "<class 'tensordict._td.TensorDict'>"}
\ No newline at end of file
diff --git a/tests/data/xarm_lift_medium/replay_buffer/observation/state.memmap b/tests/data/xarm_lift_medium/replay_buffer/observation/state.memmap
new file mode 100644
index 00000000..3bae16df
--- /dev/null
+++ b/tests/data/xarm_lift_medium/replay_buffer/observation/state.memmap
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b3e3e12896d553c208ee152f6d447c877c435e15d010c4a6171966d5b8a0c0b
+size 800
diff --git a/tests/data/xarm_lift_medium/stats.pth b/tests/data/xarm_lift_medium/stats.pth
new file mode 100644
index 0000000000000000000000000000000000000000..0accffb0c063b11524c5cdf723837d5456d86d2f
GIT binary patch
literal 4306
zcmbuDUu;uV9LI0hF-95NfPtcAn}5KLKU@FY{+nxy1*=dU<&HmP3iYnN>|W{G{o2#P
zHY<=g2_(K4<BMqaU`Tv0(KmvCF-Ay;3DNi$f<}Gu!3QK5j1R_h?zwBvxf4jsNlxy)
z-TC~!zw<l4`#bk65w|u9LUXe)ueAytLJ`TR*q4!!+&g|S8;iCJ!=`q$A&rz=F<;1N
zY1ErS8QV_%aSML5fzWddEM;hLBo|`)sf;qAr4_IaCKr>(DhQ3Nl@ijjMO$hJ6|~&w
z_FQ(78e4#kH92%NKT=c*6Ef2BIasI?mF6-ma;VlpwNWMGY>wmtEJ6HNV6i@;r3_1$
z5v`6Sp13fk$T?_BqkJK)_NENwf-+LlvKhm{)P&mx?G8GJA|5G^DzKCywpcbVqapgE
z4VEP>*rpSfC$0FJ9xp*hyb-~Q9dRMvI0Gx=mKk{9L^_?%W|cJYzG&NC$S4IRL#Oy)
zj|8g_JS4%xx`iUjfL+Df>Qa^3cEM^@x7lgf>>@V1PpAZa1i_;cJf@>#@?m(KL*1~3
zqH7KG34%H(+Jj)71nViP<zNGcdSD|(dkwUYpiYXq5V$4ql<1<4$SAF*irSO{q7kDZ
z{9prb&kXn^@F$nx(BMQ#WwfzWP99T=5Ev|t=<B78k=O?H_3AIflgxSzIMg=IClNx}
zCo6(5@yQyu#4X1mLN4TH6ZP#W1W^e#t9%f>u!W*eGiYoMeTJgXBG@XyHkA*e58@Pk
zjzOQFLkB6k9YI2Z9V#D0KS&fEV$hv)=q`%xMzBYMWJwiOze?g#z0>nGryNGGR}au|
z8D3xl^rA}Ngneb$&-|QnsEgGWHK_JiEfp))lTCOrWIYw(CEcoAh7kr&o5LP}44wr}
zYK6=KC-ni63Zv9}6@ezf%cgJ#;UGn_1~NvF97Xa7#w7qI%njHjmlFz<DH;qSOo=iR
z2o6bbn2E2>xL}epQwB3lm{%xMM)0ZxN2)=n-t{Kp3!Bd_y^g-X6@m1`9gAX{5H^b9
z$Rtv5)l&+wX#c#bXeyF-dCY1EWfkJUOz5{&WHm^?vjZ){CZ;O)Ib+d*`%oVqzP9>u
z1eDJ3@f#hHWW(Mt@ucY||K)b<1?hM5#K(H#ibdUZc(V1x*#U9Qy93kvmk$&!9S9Nq
zk3+XC_tsw`csYF>3Rmq82*S|_+ll_}4)yoi^>F~|)BhCfS31!@*RZH?Z~Yd5i@&?R
z`0x7Y*T}ig{tUl%Wq+g}lnBxPqXE5Mv%aT3{kM+26ZzzJ@cif7{te%nJb9kzuio<=
zs##yGPyg?ASEDy~UyF8~{xM3l_kOrC)q?Y$^t<JLKbw2rSkza0>giR%E_X+6ywMqX
z-`l0f{?OFb%Qfr!>em;9Bb(<sng7}=*Ml|d`|n2|2ij;vKY#Jt>6-Ne_3MAL>16n3
zkNv{ged{l5`9Qrud`SHA{pp(XFIZpx-M)AuB>wzY=$|idhgt`I(Cgpy*Kh8*w|=$$
zO(=yT{+6t({&;L;#$UZRx{GX~q~AH~?Sg1zeJ?o@^x<!CP0ngl`geFr%fzD2n%@GK
z=yZu*mrwKry&kX28SsjJm(S_-_+0^~7!>_}cfjfOdx)2p4hE+k*hA72aSL&_a&lp}
z?^f~$uJPv>Gt2J0Z5+1(A0jAQ+3w9SyIGRNt+ux-upTX^nPj&)p6tfo09BX!TxOEp
zpvc#BwauEyd7|vz!lT{zh{lS}Q)RaW-rrdsb({TV3p!7B<AWz7ikF#RR<<1T<Sb8$
zW|A$WJZWFVX?o2hTLO9V6i@ohBwM_A(r)84{brIaIXroaCj(}ZEfzf4jStyOY=dT!
zed+mBn5~dTX)&Yhn`p0QAHG|LeW0TMUn3Inrlt)B(P`aSVM)CsZSH!89H4PKW6TM2
a<X;HFxG{j7vlH<J<YO5#?E)E3$NmR@`sT0z

literal 0
HcmV?d00001

diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index b7d1e6f8..c3fcfccd 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -10,9 +10,8 @@ from .utils import DEVICE, init_config
     "env_name,dataset_id",
     [
         # TODO(rcadene): simxarm is depreciated for now
-        # ("simxarm", "lift"),
+        ("simxarm", "lift"),
         ("pusht", "pusht"),
-        # TODO(aliberts): add aloha when dataset is available on hub
         ("aloha", "sim_insertion_human"),
         ("aloha", "sim_insertion_scripted"),
         ("aloha", "sim_transfer_cube_human"),
diff --git a/tests/test_envs.py b/tests/test_envs.py
index 8931cf52..6535535e 100644
--- a/tests/test_envs.py
+++ b/tests/test_envs.py
@@ -39,19 +39,20 @@ def print_spec_rollout(env):
     print("data from rollout:", simple_rollout(100))
 
 
-@pytest.mark.skip(reason="Simxarm is deprecated")
+# @pytest.mark.skip(reason="Simxarm is deprecated")
 @pytest.mark.parametrize(
     "task,from_pixels,pixels_only",
     [
         ("lift", False, False),
         ("lift", True, False),
         ("lift", True, True),
-        ("reach", False, False),
-        ("reach", True, False),
-        ("push", False, False),
-        ("push", True, False),
-        ("peg_in_box", False, False),
-        ("peg_in_box", True, False),
+        # TODO(aliberts): Add simxarm other task or remove them completely from repo
+        # ("reach", False, False),
+        # ("reach", True, False),
+        # ("push", False, False),
+        # ("push", True, False),
+        # ("peg_in_box", False, False),
+        # ("peg_in_box", True, False),
     ],
 )
 def test_simxarm(task, from_pixels, pixels_only):
@@ -84,7 +85,7 @@ def test_pusht(from_pixels, pixels_only):
 @pytest.mark.parametrize(
     "env_name",
     [
-        # "simxarm",
+        "simxarm",
         "pusht",
         "aloha",
     ],
diff --git a/tests/test_policies.py b/tests/test_policies.py
index 92508dac..cd08fc4e 100644
--- a/tests/test_policies.py
+++ b/tests/test_policies.py
@@ -19,12 +19,13 @@ from .utils import DEVICE, init_config
     [
         ("simxarm", "tdmpc", ["policy.mpc=true"]),
         ("pusht", "tdmpc", ["policy.mpc=false"]),
-        ("simxarm", "diffusion", []),
         ("pusht", "diffusion", []),
         ("aloha", "act", ["env.task=sim_insertion_scripted"]),
         ("aloha", "act", ["env.task=sim_insertion_human"]),
         ("aloha", "act", ["env.task=sim_transfer_cube_scripted"]),
         ("aloha", "act", ["env.task=sim_transfer_cube_human"]),
+        # TODO(aliberts): simxarm not working with diffusion
+        # ("simxarm", "diffusion", []),
     ],
 )
 def test_concrete_policy(env_name, policy_name, extra_overrides):
@@ -45,13 +46,6 @@ def test_concrete_policy(env_name, policy_name, extra_overrides):
     # Check that we can make the policy object.
     policy = make_policy(cfg)
     # Check that we run select_actions and get the appropriate output.
-    if env_name == "simxarm":
-        # TODO(rcadene): Not implemented
-        return
-    if policy_name == "tdmpc":
-        # TODO(alexander-soare): TDMPC does not use n_obs_steps but the environment requires this.
-        with open_dict(cfg):
-            cfg["n_obs_steps"] = 1
     offline_buffer = make_offline_buffer(cfg)
     env = make_env(cfg, transform=offline_buffer.transform)