(env_id, make_env, n_desired_segments, clip_length_in_seconds, workers)
return [segment for sublist in results for segment in sublist]
Generate a list of path segments by doing random rollouts. Can use multiple processes.
Generate a list of path segments by doing random rollouts. Can use multiple processes.
Generate a list of path segments by doing random rollouts . Can use multiple processes .
def segments_from_rand_rollout(env_id, make_env, n_desired_segments, clip_length_in_seconds, workers):
    """ Generate a list of path segments by doing random rollouts. Can use multiple processes. """
    if workers < 2:  # Default to basic segment collection
        return basic_segments_from_rand_rollout(env_id, make_env, n_desired_segments, clip_length_in_seconds)

    pool = Pool(processes=workers)
    segments_per_worker = int(math.ceil(n_desired_segments / workers))
    # One job per worker. Only the first worker is verbose.
    jobs = [
        (env_id, make_env, segments_per_worker, clip_length_in_seconds, i, i == 0, workers)
        for i in range(workers)]
    results = pool.starmap(basic_segments_from_rand_rollout, jobs)
    pool.close()
    return [segment for sublist in results for segment in sublist]
def segments_from_rand_rollout ( env_id , make_env , n_desired_segments , clip_length_in_seconds , workers ) : if workers < 2 : # Default to basic segment collection return basic_segments_from_rand_rollout ( env_id , make_env , n_desired_segments , clip_length_in_seconds ) pool = Pool ( processes = workers ) segments_per_worker = int ( math . ceil ( n_desired_segments / workers ) ) # One job per worker. Only the first worker is verbose. jobs = [ ( env_id , make_env , segments_per_worker , clip_length_in_seconds , i , i == 0 , workers ) for i in range ( workers ) ] results = pool . starmap ( basic_segments_from_rand_rollout , jobs ) pool . close ( ) return [ segment for sublist in results for segment in sublist ]