From 35bd577deba0309aec7015374a0807e669fde634 Mon Sep 17 00:00:00 2001
From: Cadene <re.cadene@gmail.com>
Date: Sat, 9 Mar 2024 14:13:14 +0000
Subject: [PATCH] Add mock_dataset.py

---
 tests/data/pusht/next/done.memmap |   2 +-
 tests/data/pusht/stats.pth        | Bin 4498 -> 4498 bytes
 tests/scripts/create_dataset.py   |  41 ++++++++++++++++++++++++++++++
 tests/scripts/mock_dataset.py     |  41 ++++++++++++++++++++++++++++++
 4 files changed, 83 insertions(+), 1 deletion(-)
 create mode 100644 tests/scripts/create_dataset.py
 create mode 100644 tests/scripts/mock_dataset.py

diff --git a/tests/data/pusht/next/done.memmap b/tests/data/pusht/next/done.memmap
index 3c77d33..44fd709 100644
--- a/tests/data/pusht/next/done.memmap
+++ b/tests/data/pusht/next/done.memmap
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d1a8f20ab8c1dead0f61b7b38de300b0ebd0df1d870babfbbe03ce9d2b81e36a
+oid sha256:cc2786e1f9910a9d811400edcddaf7075195f7a16b216dcbefba3bc7c4f2ae51
 size 50
diff --git a/tests/data/pusht/stats.pth b/tests/data/pusht/stats.pth
index 329cb35f799f828ca8d9dff52ad54fb2538166f8..334a01bef4da3170d970dc0eb6255301f90a9d78 100644
GIT binary patch
delta 434
zcmbQFJV|-O0XBh<3yIE}pC>v8c(Zd@MLOzDzQ9%>ux8dmM}rj|jv&D>g{+N}1K5lB
z{yQ@?{0FNO^GffYd;uuL00tmcnfH}aCkJp82rycIcD9)R7pzX>_&TP^7dSB0aTZAO
zl$ARBA1inE`&kB78@(=e9uLqsp2_EVWG8>)^pd<ayV^njQ>{b0eysz@icLo<OAso`
zxXy?}tOt=G`TS#7tC8fny#)IGQk-LICxgvrSzRDM`2u%=Kz+gj$DPkx!G4>yYI4rz
z03KFGX-iW}OG8Tw6C*=2a}x_wQ&UTGkeM3}72TX%z&nuz95nWPVw0u$6j;DPV$WyB
z0uG2uK64hZ_qPBUU@yN1N`Xb}`NStX@XND+JU3ZTKx}d)zXA)$A3%=y<SqP>EFeb%
fIc6+i_kk=(6i{FRxoa{=^BjRF7O)$B3CsWh6hMl1

delta 434
zcmbQFJV|-O0X6}dl<CgWAKIJ)yxBPnr1j@bzQ9%>aQ+dC<C#ah9YKO>v3)j^1K5lB
zb~`gPJOHcv#Pj9S<O@KdRA&Z;J|~c>gm1ItCkJp82z1}w<$Oa7qHf0QR{zNtIEwh<
zoW1&2fYqH#vs^PdfU`jI!1{7$U5hH`8k<V6nbI3>zvf|JVBnd2o=0}_H%>3fv)}Wq
z8_%0tM=!c&39>@VVq*|OMH$x_@g^q*hD$&!05)Ms!{m)f^4wkmhm~eH&-%~@mUXzc
zm}~L{?gD`gvpF33mh1rgZ7KhzBbx(wSQ(`aObyK~%*+f83=EA8jVw*f%)xfv*e<R(
zxqx>f3pi-(`NSqm^C_@^gT$WCj0GGJm3-zbVDE1MGQeJb50nCX)Sgd#vID<73&?Ym
z1qH+=SMn>cfcydEh)>?aAISo8G>~J)0(KwBf<yrY7LdCpgEY?(h++Y|;g`S+0OKf=
A-2eap

diff --git a/tests/scripts/create_dataset.py b/tests/scripts/create_dataset.py
new file mode 100644
index 0000000..c58280d
--- /dev/null
+++ b/tests/scripts/create_dataset.py
@@ -0,0 +1,41 @@
+"""
+    usage: `python tests/scripts/mock_dataset.py --in-data-dir data/pusht --out-data-dir tests/data/pusht`
+"""
+
+import argparse
+import shutil
+
+from tensordict import TensorDict
+from pathlib import Path
+
+
+def mock_dataset(in_data_dir, out_data_dir, num_frames=50):
+    # load full dataset as a tensor dict
+    in_td_data = TensorDict.load_memmap(in_data_dir)
+
+    # use 1 frame to know the specification of the dataset
+    # and copy it over `n` frames in the test artifact directory
+    out_td_data = in_td_data[0].expand(num_frames).memmap_like(out_data_dir)
+
+    # copy the first `n` frames so that we have real data
+    out_td_data[:num_frames] = in_td_data[:num_frames].clone()
+
+    # make sure everything has been properly written
+    out_td_data.lock_()
+
+    # copy the full statistics of dataset since it's pretty small
+    in_stats_path = Path(in_data_dir) / "stats.pth"
+    out_stats_path = Path(out_data_dir) / "stats.pth"
+    shutil.copy(in_stats_path, out_stats_path)
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(description="Create dataset")
+
+    parser.add_argument("--in-data-dir", type=str, help="Path to input data")
+    parser.add_argument("--out-data-dir", type=str, help="Path to save the output data")
+
+    args = parser.parse_args()
+
+    mock_dataset(args.in_data_dir, args.out_data_dir)
\ No newline at end of file
diff --git a/tests/scripts/mock_dataset.py b/tests/scripts/mock_dataset.py
new file mode 100644
index 0000000..c58280d
--- /dev/null
+++ b/tests/scripts/mock_dataset.py
@@ -0,0 +1,41 @@
+"""
+    usage: `python tests/scripts/mock_dataset.py --in-data-dir data/pusht --out-data-dir tests/data/pusht`
+"""
+
+import argparse
+import shutil
+
+from tensordict import TensorDict
+from pathlib import Path
+
+
+def mock_dataset(in_data_dir, out_data_dir, num_frames=50):
+    # load full dataset as a tensor dict
+    in_td_data = TensorDict.load_memmap(in_data_dir)
+
+    # use 1 frame to know the specification of the dataset
+    # and copy it over `n` frames in the test artifact directory
+    out_td_data = in_td_data[0].expand(num_frames).memmap_like(out_data_dir)
+
+    # copy the first `n` frames so that we have real data
+    out_td_data[:num_frames] = in_td_data[:num_frames].clone()
+
+    # make sure everything has been properly written
+    out_td_data.lock_()
+
+    # copy the full statistics of dataset since it's pretty small
+    in_stats_path = Path(in_data_dir) / "stats.pth"
+    out_stats_path = Path(out_data_dir) / "stats.pth"
+    shutil.copy(in_stats_path, out_stats_path)
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(description="Create dataset")
+
+    parser.add_argument("--in-data-dir", type=str, help="Path to input data")
+    parser.add_argument("--out-data-dir", type=str, help="Path to save the output data")
+
+    args = parser.parse_args()
+
+    mock_dataset(args.in_data_dir, args.out_data_dir)
\ No newline at end of file