code: def __init__(self, dataset_path, pct_traj, max_ep_len, scale):
      self.dataset_path = dataset_path
      self.pct_traj = pct_traj
      self.scale = scale
      self.device = torch.device(
          "cuda" if torch.cuda.is_available() else "cpu"
      )
      self.max_ep_len = max_ep_len
      with open(dataset_path, "rb") as f:
          trajectories = pickle.load(f)
  
      states, traj_lens, returns = [], [], []
      for path in trajectories:
          states.append(path["observations"])
          traj_lens.append(len(path["observations"]))
          returns.append(path["rewards"].sum())
      traj_lens, returns = np.array(traj_lens), np.array(returns)
  
      # used for input normalization
      states = np.concatenate(states, axis=0)
      self.state_mean, self.state_std = (
          np.mean(states, axis=0),
          np.std(states, axis=0) + 1e-6,
      )
  
      num_timesteps = sum(traj_lens)
  
      # only train on top pct_traj trajectories (for %BC experiment)
      num_timesteps = max(int(pct_traj * num_timesteps), 1)
      sorted_inds = np.argsort(returns)  # lowest to highest
      num_trajectories = 1
      timesteps = traj_lens[sorted_inds[-1]]
      ind = len(trajectories) - 2
      while (
          ind >= 0
          and timesteps + traj_lens[sorted_inds[ind]] <= num_timesteps
      ):
          timesteps += traj_lens[sorted_inds[ind]]
          num_trajectories += 1
          ind -= 1
      sorted_inds = sorted_inds[-num_trajectories:]
  
      # used to reweight sampling so we sample according to timesteps instead of trajectories
      p_sample = traj_lens[sorted_inds] / sum(traj_lens[sorted_inds])
  
      logger.info(
          f"{len(traj_lens)} trajectories, {num_timesteps} timesteps found"
      )
      logger.info(
          f"Average return: {np.mean(returns):.2f}, std: {np.std(returns):.2f}"
      )
      logger.info(
          f"Max return: {np.max(returns):.2f}, min: {np.min(returns):.2f}"
      )
      self.trajectories = trajectories
      self.num_trajectories = num_trajectories
      self.p_sample = p_sample
      self.sorted_inds = sorted_inds
      self.state_dim = trajectories[0]["observations"].shape[1]
      self.act_dim = trajectories[0]["actions"].shape[1]
  
file_path: PARL/benchmark/torch/DT/data_loader.py
imports: import numpy
  import pickle
  import random
  import torch
  from parl.utils import logger
label_desc: Use image node_id 1 to create a new DataLoader object from inherited base classes: object with example: obj = DataLoader(dataset_path, pct_traj, max_ep_len, scale)