File size: 1,218 Bytes
987ed1b | 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 | import torch
import numpy as np
def slice_episode(episode, horizon, stride, start=0):
is_torch = hasattr(episode, 'cuda')
shape = episode.shape
N, T = shape[:2]
sliced_fragments = []
for current_start in range(start, T, stride):
end = current_start + horizon
if end > T:
current_start = max(0, T - horizon)
end = T
fragment = episode[:, current_start:end, ...]
sliced_fragments.append(fragment)
if end == T:
break
if is_torch:
return torch.stack(sliced_fragments)
else:
return np.stack(sliced_fragments)
def slice_episode_time(episode, horizon, stride):
is_torch = hasattr(episode, 'cuda')
shape = episode.shape
T = shape[0]
sliced_fragments = []
for start in range(0, T, stride):
end = start + horizon
if end > T:
start = max(0, T - horizon)
end = T
fragment = episode[start:end, ...]
sliced_fragments.append(fragment)
if end == T:
break
if is_torch:
return torch.stack(sliced_fragments)
else:
return np.stack(sliced_fragments) |