File size: 3,739 Bytes
2631354 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
diff --git a/data/envs/metaworld/generate_dataset.py b/data/envs/metaworld/generate_dataset.py
index dc51b4f..ed99a4a 100644
--- a/data/envs/metaworld/generate_dataset.py
+++ b/data/envs/metaworld/generate_dataset.py
@@ -4,7 +4,9 @@ from typing import Dict, Optional
import gym
import metaworld
import numpy as np
+import pandas as pd
import torch
+from datasets import Dataset
from huggingface_hub import HfApi, repocard, upload_folder
from sample_factory.algo.learning.learner import Learner
from sample_factory.algo.sampling.batched_sampling import preprocess_actions
@@ -12,11 +14,7 @@ from sample_factory.algo.utils.action_distributions import argmax_actions
from sample_factory.algo.utils.env_info import extract_env_info
from sample_factory.algo.utils.make_env import make_env_func_batched
from sample_factory.algo.utils.rl_utils import make_dones, prepare_and_normalize_obs
-from sample_factory.cfg.arguments import (
- load_from_checkpoint,
- parse_full_cfg,
- parse_sf_args,
-)
+from sample_factory.cfg.arguments import load_from_checkpoint, parse_full_cfg, parse_sf_args
from sample_factory.envs.env_utils import register_env
from sample_factory.model.actor_critic import create_actor_critic
from sample_factory.model.model_utils import get_rnn_size
@@ -165,10 +163,8 @@ def create_dataset(cfg: Config):
# Create dataset
dataset_size = 100_000
dataset = {
- "observations": np.zeros(
- (dataset_size, *env.observation_space["obs"].shape), dtype=env.observation_space["obs"].dtype
- ),
- "actions": np.zeros((dataset_size, *env.action_space.shape), env.action_space.dtype),
+ "observations": np.zeros((dataset_size, *env.observation_space["obs"].shape), dtype=np.float32),
+ "actions": np.zeros((dataset_size, *env.action_space.shape), np.float32),
"dones": np.zeros((dataset_size,), bool),
"rewards": np.zeros((dataset_size,), np.float32),
}
@@ -206,6 +202,13 @@ def create_dataset(cfg: Config):
env.close()
+ # Convert dict of numpy array to pandas dataframe
+# dataset = Dataset.from_dict(dataset)
+# dataset.create_config_id
+ # Set the card of the dataset
+# dataset.card = f""""""
+# dataset.push_to_hub("qgallouedec/prj_gia_dataset_metaworld_assembly_v2_1111_demo")
+
# Save dataset
repo_path = f"{cfg.train_dir}/datasets/{cfg.experiment}"
os.makedirs(repo_path, exist_ok=True)
diff --git a/data/envs/metaworld/generate_dataset_all.sh b/data/envs/metaworld/generate_dataset_all.sh
index 802bf5c..3cc4f97 100755
--- a/data/envs/metaworld/generate_dataset_all.sh
+++ b/data/envs/metaworld/generate_dataset_all.sh
@@ -1,34 +1,6 @@
#!/bin/bash
ENVS=(
- assembly
- basketball
- bin-picking
- box-close
- button-press-topdown
- button-press-topdown-wall
- button-press
- button-press-wall
- coffee-button
- coffee-pull
- coffee-push
- dial-turn
- disassemble
- door-close
- door-lock
- door-open
- door-unlock
- drawer-close
- drawer-open
- faucet-close
- faucet-open
- hammer
- hand-insert
- handle-press-side
- handle-press
- handle-pull-side
- handle-pull
- lever-pull
peg-insert-side
peg-unplug-side
pick-out-of-hole
@@ -40,19 +12,8 @@ ENVS=(
plate-slide
push-back
push
- push-wall
- reach
- reach-wall
- shelf-place
- soccer
- stick-pull
- stick-push
- sweep-into
- sweep
- window-close
- window-open
)
for ENV in "${ENVS[@]}"; do
- python generate_dataset.py --env $ENV-v2 --experiment $ENV-v2 --train_dir=./train_dir
+ python generate_dataset.py --env $ENV-v2 --experiment $ENV-v2 --train_dir=./train_dir
done
|