def __init__(
self,
rollouts_dir: str,
embedding: Embedding,
num_segments: int = None,
len_segments: int = None,
rand_init_rollout: bool = True,
):
"""
Constructor
:param rollouts_dir: directory where to find the of pre-recorded rollouts
:param num_segments: number of segments in which the rollouts are split into. For every segment, the initial
state of the simulation is reset, and thus for every set the features of the trajectories
are computed separately. Either specify `num_segments` or `len_segments`.
:param embedding: embedding used for pre-processing the data before (later) passing it to the posterior
:param len_segments: length of the segments in which the rollouts are split into. For every segment, the initial
state of the simulation is reset, and thus for every set the features of the trajectories
are computed separately. Either specify `num_segments` or `len_segments`.
:param rand_init_rollout: if `True`, chose the first rollout at random, and then cycle through the list
"""
if not os.path.isdir(rollouts_dir):
raise pyrado.PathErr(given=rollouts_dir)
Serializable._init(self, locals())
super().__init__(None, None, embedding, num_segments, len_segments)
# Crawl through the directory and load every file that starts with the word rollout
rollouts_rec = []
for root, dirs, files in os.walk(rollouts_dir):
dirs.clear() # prevents walk() from going into subdirectories
rollouts_rec = [pyrado.load(name=f, load_dir=root) for f in files if f.startswith("rollout")]
check_all_lengths_equal(rollouts_rec)
if not rollouts_rec:
raise pyrado.ValueErr(msg="No rollouts have been found!")
self.rollouts_dir = rollouts_dir
self.rollouts_rec = rollouts_rec
self._ring_idx = np.random.randint(0, len(rollouts_rec)) if rand_init_rollout else 0
self._set_action_field(self.rollouts_rec)
def plot_rollouts_segment_wise(
plot_type: str,
segments_ground_truth: List[List[StepSequence]],
segments_multiple_envs: List[List[List[StepSequence]]],
segments_nominal: List[List[StepSequence]],
use_rec_str: bool,
idx_iter: int,
idx_round: int,
state_labels: Optional[Iterable[str]] = None,
act_labels: Optional[Iterable[str]] = None,
x_limits: Optional[Tuple[int]] = None,
plot_act: bool = False,
data_field: str = "states",
cmap_samples: Optional[colors.Colormap] = None,
save_dir: Optional[pyrado.PathLike] = None,
file_format: Iterable[str] = ("pdf", "pgf", "png"),
) -> List[plt.Figure]:
r"""
Plot the different rollouts in separate figures and the different state dimensions along the columns.
:param plot_type: type of plot, pass "samples" to plot the rollouts of the most likely domain parameters as
individual lines, or pass "confidence" to plot the most likely one, and the mean $\pm$ 1 std
:param segments_ground_truth: list of lists containing rollout segments from the ground truth environment
:param segments_multiple_envs: list of lists of lists containing rollout segments from different environment
instances, e.g. samples from a posterior coming from `NDPR`
:param segments_nominal: list of lists containing rollout segments from the nominal environment
:param use_rec_str: `True` if pre-recorded actions have been used to generate the rollouts
:param idx_iter: selected iteration
:param idx_round: selected round
:param state_labels: y-axes labels to for the state trajectories, no label by default
:param act_labels: y-axes labels to for the action trajectories, no label by default
:param x_limits: tuple containing the lower and upper limits for the x-axis
:param plot_act: if `True`, also plot the actions
:param data_field: data field of the rollout, e.g. "states" or "observations"
:param cmap_samples: color map for the trajectories resulting from different domain parameter samples
:param save_dir: if not `None` create a subfolder plots in `save_dir` and save the plots in there
:param file_format: select the file format to store the plots
:return: list of handles to the created figures
"""
if plot_type not in ["samples", "confidence"]:
raise pyrado.ValueErr(given=plot_type,
eq_constraint="samples or confidence")
if data_field not in ["states", "observations"]:
raise pyrado.ValueErr(given=data_field,
eq_constraint="states or observations")
# Extract the state dimension, and the number of most likely samples from the data
dim_state = segments_ground_truth[0][0].get_data_values(data_field)[
0, :].size
dim_act = segments_ground_truth[0][0].get_data_values("actions")[0, :].size
num_samples = len(segments_multiple_envs[0][0])
# Validate the labels
if state_labels is None:
state_labels = [""] * dim_state
else:
if len(state_labels) != dim_state:
raise pyrado.ShapeErr(given=state_labels,
expected_match=(dim_state, ))
if act_labels is None:
act_labels = [""] * dim_act
else:
if len(act_labels) != dim_act:
raise pyrado.ShapeErr(given=act_labels, expected_match=(dim_act, ))
if cmap_samples is None:
cmap_samples = plt.get_cmap("Reds")(np.linspace(0.6, 0.8, num_samples))
fig_list = []
label_samples = "ml" if plot_type == "confidence" else "samples"
# Plot
for idx_r in range(len(segments_ground_truth)):
num_rows = dim_state + dim_act if plot_act else dim_state
fig, axs = plt.subplots(nrows=num_rows,
figsize=(16, 9),
tight_layout=True,
sharex="col")
axs = np.atleast_1d(axs)
# Plot the states
for idx_state in range(dim_state):
# Plot the real segments
cnt_step = [0]
for segment_gt in segments_ground_truth[idx_r]:
axs[idx_state].plot(
np.arange(cnt_step[-1], cnt_step[-1] + segment_gt.length),
segment_gt.get_data_values(data_field,
truncate_last=True)[:,
idx_state],
zorder=0,
c="black",
lw=1.0,
label="target" if cnt_step[-1] == 0 else "", # print once
)
cnt_step.append(cnt_step[-1] + segment_gt.length)
# Plot the maximum likely simulated segments
for idx_seg, sml in enumerate(segments_multiple_envs[idx_r]):
for idx_dp, sdp in enumerate(sml):
axs[idx_state].plot(
np.arange(cnt_step[idx_seg],
cnt_step[idx_seg] + sdp.length),
sdp.get_data_values(data_field,
truncate_last=True)[:, idx_state],
zorder=2 if idx_dp == 0 else 0, # most likely on top
c=cmap_samples[idx_dp],
ls="-",
lw=1.5
if plot_type == "confidence" or idx_dp == 0 else 0.5,
alpha=1.0
if plot_type == "confidence" or idx_dp == 0 else 0.1,
label=label_samples
if cnt_step[idx_seg] == idx_seg == idx_dp == 0 else
"", # print once
)
if plot_type != "samples":
# Stop here, unless the rollouts of most likely all domain parameters should be plotted
break
if plot_type == "confidence":
assert check_all_lengths_equal(
segments_multiple_envs[idx_r]
[0]) # all segments need to be equally long
states_all = []
for idx_dp in range(num_samples):
# Reconstruct the step sequences for all domain parameters
ss_dp = StepSequence.concat([
seg_dp[idx_dp] for seg_dp in
[segs_ro for segs_ro in segments_multiple_envs[idx_r]]
])
states = ss_dp.get_data_values(data_field,
truncate_last=True)
states_all.append(states)
states_all = np.stack(states_all, axis=0)
states_mean = np.mean(states_all, axis=0)
states_std = np.std(states_all, axis=0)
for idx_seg in range(len(segments_multiple_envs[idx_r])):
len_segs = min([
len(seg)
for seg in segments_multiple_envs[idx_r][idx_seg]
]) # use shortest
m_i = states_mean[cnt_step[idx_seg]:cnt_step[idx_seg] +
len_segs, idx_state]
s_i = states_std[cnt_step[idx_seg]:cnt_step[idx_seg] +
len_segs, idx_state]
draw_curve(
"std",
axs[idx_state],
pd.DataFrame(dict(mean=m_i, std=s_i)),
x_grid=np.arange(cnt_step[idx_seg],
cnt_step[idx_seg] + len_segs),
show_legend=False,
area_label="2 std" if idx_seg == 0 else None,
plot_kwargs=dict(color=cmap_samples[0]),
)
# Plot the nominal simulation's segments
for idx_seg, sn in enumerate(segments_nominal[idx_r]):
axs[idx_state].plot(
np.arange(cnt_step[idx_seg],
cnt_step[idx_seg] + sn.length),
sn.get_data_values(data_field,
truncate_last=True)[:, idx_state],
zorder=2,
c="green", # former: steelblue"
ls="--",
lw=1.0,
label="nom sim"
if cnt_step[idx_seg] == 0 else "", # print once
)
axs[idx_state].set_ylabel(state_labels[idx_state])
if plot_act:
# Plot the actions
for idx_act in range(dim_act):
# Plot the real segments
cnt_step = [0]
for segment_gt in segments_ground_truth[idx_r]:
axs[dim_state + idx_act].plot(
np.arange(cnt_step[-1],
cnt_step[-1] + segment_gt.length),
segment_gt.get_data_values(
"actions", truncate_last=False)[:, idx_act],
zorder=0,
c="black",
label="target"
if cnt_step[-1] == 0 else "", # print once
)
cnt_step.append(cnt_step[-1] + segment_gt.length)
# Plot the maximum likely simulated segments
for idx_seg, sml in enumerate(segments_multiple_envs[idx_r]):
for idx_dp, sdp in enumerate(sml):
axs[dim_state + idx_act].plot(
np.arange(cnt_step[idx_seg],
cnt_step[idx_seg] + sdp.length),
sdp.get_data_values("actions",
truncate_last=False)[:,
idx_act],
zorder=2
if idx_dp == 0 else 0, # most likely on top
c=cmap_samples[idx_dp],
ls="--",
lw=1.5 if plot_type == "confidence" or idx_dp == 0
else 0.5,
alpha=1.0 if plot_type == "confidence"
or idx_dp == 0 else 0.4,
label=label_samples
if cnt_step[idx_seg] == idx_seg == idx_dp == 0 else
"", # print once
)
if plot_type != "samples":
# Stop here, unless the rollouts of most likely all domain parameters should be plotted
break
if plot_type == "confidence":
len_segs = len(segments_multiple_envs[idx_r][0][0])
assert check_all_lengths_equal(
segments_multiple_envs[idx_r]
[0]) # all segments need to be equally long
acts_all = []
for idx_dp in range(num_samples):
# Reconstruct the step sequences for all domain parameters
ss_dp = StepSequence.concat([
seg_dp[idx_dp] for seg_dp in [
segs_ro
for segs_ro in segments_multiple_envs[idx_r]
]
])
acts = ss_dp.get_data_values("actions",
truncate_last=False)
acts_all.append(acts)
acts_all = np.stack(acts_all, axis=0)
acts_mean = np.mean(acts_all, axis=0)
acts_std = np.std(acts_all, axis=0)
for idx_seg in range(len(segments_multiple_envs[idx_r])):
m_i = acts_mean[cnt_step[idx_seg]:cnt_step[idx_seg] +
len_segs, idx_act]
s_i = acts_std[cnt_step[idx_seg]:cnt_step[idx_seg] +
len_segs, idx_act]
draw_curve(
"std",
axs[dim_state + idx_act],
pd.DataFrame(dict(mean=m_i, std=s_i)),
x_grid=np.arange(cnt_step[idx_seg],
cnt_step[idx_seg] + len_segs),
show_legend=False,
area_label="2 std" if idx_seg == 0 else None,
plot_kwargs=dict(color=cmap_samples[0]),
)
# Plot the nominal simulation's segments
for idx_seg, sn in enumerate(segments_nominal[idx_r]):
axs[dim_state + idx_act].plot(
np.arange(cnt_step[idx_seg],
cnt_step[idx_seg] + sn.length),
sn.get_data_values("actions",
truncate_last=False)[:, idx_act],
zorder=2,
c="steelblue",
ls="-.",
label="nom sim"
if cnt_step[idx_seg] == 0 else "", # print once
)
axs[dim_state + idx_act].set_ylabel(act_labels[idx_act])
# Settings for all subplots
for idx_axs in range(num_rows):
if x_limits is not None:
axs[idx_axs].set_xlim(x_limits[0], x_limits[1])
# Set window title and the legend, placing the latter above the plot expanding and expanding it fully
use_rec_str = ", using rec actions" if use_rec_str else ""
round_str = f"round {idx_round}, " if idx_round != -1 else ""
fig.canvas.manager.set_window_title(
f"Target Domain and Simulated Rollouts (iteration {idx_iter}, {round_str}rollout {idx_r}{use_rec_str})"
)
lg = axs[0].legend(
ncol=2 + num_samples,
bbox_to_anchor=(0.0, 1.02, 1.0, 0.102),
loc="lower left",
mode="expand",
borderaxespad=0.0,
)
# Save if desired
if save_dir is not None:
for fmt in file_format:
os.makedirs(os.path.join(save_dir, "plots"), exist_ok=True)
len_seg_str = f"seglen_{segments_ground_truth[0][0].length}"
use_rec_str = "_use_rec" if use_rec_str else ""
round_str = f"_round_{idx_round}" if idx_round != -1 else ""
fig.savefig(
os.path.join(
save_dir,
"plots",
f"posterior_iter_{idx_iter}{round_str}_rollout_{idx_r}_{len_seg_str}{use_rec_str}.{fmt}",
),
bbox_extra_artists=(lg, ),
dpi=150,
)
# Append current figure
fig_list.append(fig)
return fig_list
# Get the experiments' directories to load from
prefixes = [
osp.join(pyrado.EXP_DIR, 'ENV_NAME', 'ALGO_NAME'),
]
ex_names = [
'',
]
ex_labels = [
'',
]
else:
raise pyrado.ValueErr(given=args.env_name, eq_constraint=f'{QCartPoleSwingUpSim.name}, {QCartPoleStabSim.name},'
f' or {QCartPoleSwingUpSim.name}')
if not check_all_lengths_equal([prefixes, ex_names, ex_labels]):
raise pyrado.ShapeErr(msg=f'The lengths of prefixes, ex_names, and ex_labels must be equal, '
f'but they are {len(prefixes)}, {len(ex_names)}, and {len(ex_labels)}!')
# Create Randomizer
pert = create_conservative_randomizer(env)
pert.add_domain_params(UniformDomainParam(name='act_delay', mean=20, halfspan=20, clip_lo=0, roundint=True))
# Loading the policies
ex_dirs = [osp.join(p, e) for p, e in zip(prefixes, ex_names)]
env_sim_list = []
policy_list = []
for ex_dir in ex_dirs:
env_sim, policy, _ = load_experiment(ex_dir, args)
policy_list.append(policy)
def test_check_all_lengths_equal(x, b):
assert check_all_lengths_equal(x) == b
