fix dataratio bug

xmy86 · Dec 9, 2023 · 261c602 · 261c602
1 parent 6164990
commit 261c602
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jaxrl5/data/dsrl_datasets.py b/jaxrl5/data/dsrl_datasets.py
@@ -36,7 +36,7 @@ def __init__(self, env: gym.Env, clip_to_eps: bool = True, eps: float = 1e-5, cr
                 file_list = dataset_name.split('-')
                 ratio_num = int(float(file_list[-1].split('.')[0]) * ratio)
                 dataset_ratio = '-'.join(file_list[:-1]) + '-' + str(ratio_num) + '-' + str(ratio) + '.hdf5'
-                dataset_dict = env.get_dataset(dataset_ratio)
+                dataset_dict = env.get_dataset(os.path.join('data', dataset_ratio))
             print('max_episode_reward', env.max_episode_reward, 
                 'min_episode_reward', env.min_episode_reward,
                 'mean_episode_reward', env._max_episode_steps * np.mean(dataset_dict['rewards']))