# Path: rl_teacher/envs.py
# def get_timesteps_per_episode(env):
#     if hasattr(env, "_max_episode_steps"):
#         return env._max_episode_steps
#     if hasattr(env, "spec"):
#         return env.spec.tags.get("wrapper_config.TimeLimit.max_episode_steps")
#     if hasattr(env, "env"):
#         return get_timesteps_per_episode(env.env)
#     return None


import math
import numpy as np
import gym.spaces.prng as space_prng
from multiprocessing import Pool
from rl_teacher.envs import get_timesteps_per_episode
        for k, v in path.items()
        if k in ['obs', "actions", 'original_rewards', 'human_obs']}

def create_segment_q_states(segment):
    obs_Ds = segment["obs"]
    act_Ds = segment["actions"]
    return np.concatenate([obs_Ds, act_Ds], axis=1)

def sample_segment_from_path(path, segment_length):
    """Returns a segment sampled from a random place in a path. Returns None if the path is too short"""
    path_length = len(path["obs"])
    if path_length < segment_length:
        return None

    start_pos = np.random.randint(0, path_length - segment_length + 1)

    # Build segment
    segment = _slice_path(path, segment_length, start_pos)

    # Add q_states
    segment["q_states"] = create_segment_q_states(segment)
    return segment

def random_action(env, ob):
    """ Pick an action by uniformly sampling the environment's action space. """
    return env.action_space.sample()

def do_rollout(env, action_function):
    """ Builds a path by running through an environment using a provided function to select actions. """
    obs, rewards, actions, human_obs = [], [], [], []

# Path: rl_teacher/envs.py
# def get_timesteps_per_episode(env):
#     if hasattr(env, "_max_episode_steps"):
#         return env._max_episode_steps
#     if hasattr(env, "spec"):
#         return env.spec.tags.get("wrapper_config.TimeLimit.max_episode_steps")
#     if hasattr(env, "env"):
#         return get_timesteps_per_episode(env.env)
#     return None


# Path: rl_teacher/segment_sampling.py
import math
import numpy as np
import gym.spaces.prng as space_prng
from multiprocessing import Pool
from rl_teacher.envs import get_timesteps_per_episode

        for k, v in path.items()
        if k in ['obs', "actions", 'original_rewards', 'human_obs']}

def create_segment_q_states(segment):
    obs_Ds = segment["obs"]
    act_Ds = segment["actions"]
    return np.concatenate([obs_Ds, act_Ds], axis=1)

def sample_segment_from_path(path, segment_length):
    """Returns a segment sampled from a random place in a path. Returns None if the path is too short"""
    path_length = len(path["obs"])
    if path_length < segment_length:
        return None

    start_pos = np.random.randint(0, path_length - segment_length + 1)

    # Build segment
    segment = _slice_path(path, segment_length, start_pos)

    # Add q_states
    segment["q_states"] = create_segment_q_states(segment)
    return segment

def random_action(env, ob):
    """ Pick an action by uniformly sampling the environment's action space. """
    return env.action_space.sample()

def do_rollout(env, action_function):
    """ Builds a path by running through an environment using a provided function to select actions. """
    obs, rewards, actions, human_obs = [], [], [], []

    max_timesteps_per_episode = get_timesteps_per_episode(env)