Spaces:

crossentropy-ai
/

rlcube

Sleeping

imwithye commited on Sep 20

Commit

e71fda3

1 Parent(s): ee0d368

train

Files changed (1) hide show

rlcube/rlcube/train/train.py CHANGED Viewed

@@ -17,7 +17,7 @@ print(f"Using device: {device}")
 def train(epochs: int = 100):
     if not os.path.exists("dataset.pt"):
-        create_dataset(num_envs=1000, num_steps=20, filepath="dataset.pt")
     dataset = Cube2Dataset("dataset.pt")
     print("Number of samples:", len(dataset))
     print("Number of epochs:", epochs)

 def train(epochs: int = 100):
     if not os.path.exists("dataset.pt"):
+        create_dataset(num_envs=10000, num_steps=30, filepath="dataset.pt")
     dataset = Cube2Dataset("dataset.pt")
     print("Number of samples:", len(dataset))
     print("Number of epochs:", epochs)