generate dataset

.gitattributes

@@ -1 +1,2 @@
 *.blend filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text

rlcube/.gitignore

@@ -217,3 +217,6 @@ __marimo__/
 
 # Blender
 *.blend1
+
+# Dataset
+dataset.pt

rlcube/pyproject.toml

@@ -12,4 +12,5 @@ dependencies = [
     "numpy>=2.3.2",
     "tensordict>=0.10.0",
     "torch>=2.8.0",
+    "tqdm>=4.67.1",
 ]

rlcube/rlcube/models/dataset.py

@@ -1,27 +1,43 @@
 from torch.utils.data import Dataset
 from rlcube.envs.cube2 import Cube2
 import numpy as np
+import torch
+from tqdm import tqdm
+
+
+def create_dataset(num_envs: int = 10000, num_steps: int = 50):
+    states = []
+    neighbors = []
+    D = []
+    for _ in tqdm(range(num_envs)):
+        env = Cube2()
+        obs, _ = env.reset()
+        for _ in range(num_steps):
+            action = env.action_space.sample()
+            obs, _, _, _, _ = env.step(action)
+            states.append(obs)
+            neighbors.append(env.neighbors())
+            D.append(env.step_count)
+    states = np.array(states)
+    neighbors = np.array(neighbors)
+    D = np.array(D)
+    dataseet = {
+        "states": torch.tensor(states),
+        "neighbors": torch.tensor(neighbors),
+        "D": torch.tensor(D),
+    }
+    torch.save(dataseet, "dataset.pt")
 
 
 class Cube2Dataset(Dataset):
-    def __init__(self, num_envs: int = 1000, num_steps: int = 20):
-        self.num_envs = num_envs
-        self.num_steps = num_steps
-        self.states = []
-        self.D = []
-        for _ in range(num_envs):
-            env = Cube2()
-            obs, _ = env.reset()
-            for _ in range(num_steps):
-                action = env.action_space.sample()
-                obs, _, _, _, _ = env.step(action)
-                self.states.append(obs)
-                self.D.append(env.step_count)
-        self.states = np.array(self.states)
-        self.D = np.array(self.D)
+    def __init__(self, filepath: str = "dataset.pt"):
+        self.dataset = torch.load(filepath)
+        self.states = self.dataset["states"]
+        self.neighbors = self.dataset["neighbors"]
+        self.D = self.dataset["D"]
 
     def __len__(self):
         return len(self.states)
 
     def __getitem__(self, idx):
-        return self.states[idx], self.D[idx]
+        return self.states[idx], self.neighbors[idx], self.D[idx]

rlcube/rlcube/train/train.py

@@ -1,6 +1,14 @@
 from rlcube.models.dataset import Cube2Dataset
+from rlcube.envs.cube2 import Cube2
+import numpy as np
+from tqdm import tqdm
+
+
+def train(epochs: int = 100):
+    dataset = Cube2Dataset()
+    for _ in tqdm(range(epochs)):
+        pass
 
 
 if __name__ == "__main__":
-    dataset = Cube2Dataset(num_envs=10, num_steps=20)
-    print(dataset[10])
+    train()

rlcube/uv.lock

@@ -1536,6 +1536,7 @@ dependencies = [
     { name = "numpy" },
     { name = "tensordict" },
     { name = "torch" },
+    { name = "tqdm" },
 ]
 
 [package.metadata]
@@ -1547,6 +1548,7 @@ requires-dist = [
     { name = "numpy", specifier = ">=2.3.2" },
     { name = "tensordict", specifier = ">=0.10.0" },
     { name = "torch", specifier = ">=2.8.0" },
+    { name = "tqdm", specifier = ">=4.67.1" },
 ]
 
 [[package]]