mirror of
https://github.com/Farama-Foundation/Gymnasium.git
synced 2025-08-22 07:02:19 +00:00
Add testing for step api compatibility functions and wrapper (#3028)
* Initial commit * Fixed tests and forced TimeLimit.truncated to always exist when truncated or terminated * Fix CI issues * pre-commit * Revert back to old language * Revert changes to step api wrapper
This commit is contained in:
Mark Towers
GitHub
@@ -36,66 +36,41 @@ def step_to_new_api(
|
||||
assert len(step_returns) == 4
|
||||
observations, rewards, dones, infos = step_returns
|
||||
|
||||
terminateds = []
|
||||
truncateds = []
|
||||
if not is_vector_env:
|
||||
dones = [dones]
|
||||
|
||||
for i in range(len(dones)):
|
||||
# For every condition, handling - info single env / info vector env (list) / info vector env (dict)
|
||||
|
||||
# TimeLimit.truncated attribute not present - implies either terminated or episode still ongoing based on `done`
|
||||
if (not is_vector_env and "TimeLimit.truncated" not in infos) or (
|
||||
is_vector_env
|
||||
and (
|
||||
(
|
||||
isinstance(infos, list)
|
||||
and "TimeLimit.truncated" not in infos[i]
|
||||
) # vector env, list info api
|
||||
or (
|
||||
"TimeLimit.truncated" not in infos
|
||||
or (
|
||||
"TimeLimit.truncated" in infos
|
||||
and not infos["TimeLimit.truncated"][i]
|
||||
)
|
||||
)
|
||||
# vector env, dict info api, for env i, vector mask `_TimeLimit.truncated` is not considered, to be compatible with envpool
|
||||
# For env i, `TimeLimit.truncated` not being present is treated same as being present and set to False.
|
||||
# therefore, terminated=True, truncated=True simultaneously is not allowed while using compatibility functions
|
||||
# with vector info
|
||||
)
|
||||
):
|
||||
terminateds.append(dones[i])
|
||||
truncateds.append(False)
|
||||
|
||||
# This means info["TimeLimit.truncated"] exists and this elif checks if it is True, which means the truncation has occurred but termination has not.
|
||||
elif (
|
||||
infos["TimeLimit.truncated"]
|
||||
if not is_vector_env
|
||||
else (
|
||||
infos["TimeLimit.truncated"][i]
|
||||
if isinstance(infos, dict)
|
||||
else infos[i]["TimeLimit.truncated"]
|
||||
)
|
||||
):
|
||||
assert dones[i]
|
||||
terminateds.append(False)
|
||||
truncateds.append(True)
|
||||
else:
|
||||
# This means info["TimeLimit.truncated"] exists but is False, which means the core environment had already terminated,
|
||||
# but it also exceeded maximum timesteps at the same step. However to be compatible with envpool, and to be backward compatible
|
||||
# truncated is set to False here.
|
||||
assert dones[i]
|
||||
terminateds.append(True)
|
||||
truncateds.append(False)
|
||||
|
||||
return (
|
||||
observations,
|
||||
rewards,
|
||||
np.array(terminateds, dtype=np.bool_) if is_vector_env else terminateds[0],
|
||||
np.array(truncateds, dtype=np.bool_) if is_vector_env else truncateds[0],
|
||||
infos,
|
||||
)
|
||||
# Cases to handle - info single env / info vector env (list) / info vector env (dict)
|
||||
if is_vector_env is False:
|
||||
truncated = infos.pop("TimeLimit.truncated", False)
|
||||
return (
|
||||
observations,
|
||||
rewards,
|
||||
dones and not truncated,
|
||||
dones and truncated,
|
||||
infos,
|
||||
)
|
||||
elif isinstance(infos, list):
|
||||
truncated = np.array(
|
||||
[info.pop("TimeLimit.truncated", False) for info in infos]
|
||||
)
|
||||
return (
|
||||
observations,
|
||||
rewards,
|
||||
np.logical_and(dones, np.logical_not(truncated)),
|
||||
np.logical_and(dones, truncated),
|
||||
infos,
|
||||
)
|
||||
elif isinstance(infos, dict):
|
||||
num_envs = len(dones)
|
||||
truncated = infos.pop("TimeLimit.truncated", np.zeros(num_envs, dtype=bool))
|
||||
return (
|
||||
observations,
|
||||
rewards,
|
||||
np.logical_and(dones, np.logical_not(truncated)),
|
||||
np.logical_and(dones, truncated),
|
||||
infos,
|
||||
)
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Unexpected value of infos, as is_vector_envs=False, expects `info` to be a list or dict, actual type: {type(infos)}"
|
||||
)
|
||||
|
||||
|
||||
def step_to_old_api(
|
||||
@@ -111,44 +86,45 @@ def step_to_old_api(
|
||||
return step_returns
|
||||
else:
|
||||
assert len(step_returns) == 5
|
||||
observations, rewards, terminateds, truncateds, infos = step_returns
|
||||
dones = []
|
||||
if not is_vector_env:
|
||||
terminateds = [terminateds]
|
||||
truncateds = [truncateds]
|
||||
observations, rewards, terminated, truncated, infos = step_returns
|
||||
|
||||
n_envs = len(terminateds)
|
||||
|
||||
for i in range(n_envs):
|
||||
dones.append(terminateds[i] or truncateds[i])
|
||||
if truncateds[i]:
|
||||
if is_vector_env:
|
||||
# handle vector infos for dict and list
|
||||
if isinstance(infos, dict):
|
||||
if "TimeLimit.truncated" not in infos:
|
||||
# TODO: This should ideally not be done manually and should use vector_env's _add_info()
|
||||
infos["TimeLimit.truncated"] = np.zeros(n_envs, dtype=bool)
|
||||
infos["_TimeLimit.truncated"] = np.zeros(n_envs, dtype=bool)
|
||||
|
||||
infos["TimeLimit.truncated"][i] = (
|
||||
not terminateds[i] or infos["TimeLimit.truncated"][i]
|
||||
)
|
||||
infos["_TimeLimit.truncated"][i] = True
|
||||
else:
|
||||
# if vector info is a list
|
||||
infos[i]["TimeLimit.truncated"] = not terminateds[i] or infos[
|
||||
i
|
||||
].get("TimeLimit.truncated", False)
|
||||
else:
|
||||
infos["TimeLimit.truncated"] = not terminateds[i] or infos.get(
|
||||
"TimeLimit.truncated", False
|
||||
)
|
||||
return (
|
||||
observations,
|
||||
rewards,
|
||||
np.array(dones, dtype=np.bool_) if is_vector_env else dones[0],
|
||||
infos,
|
||||
)
|
||||
# Cases to handle - info single env / info vector env (list) / info vector env (dict)
|
||||
if is_vector_env is False:
|
||||
if truncated or terminated:
|
||||
infos["TimeLimit.truncated"] = truncated and not terminated
|
||||
return (
|
||||
observations,
|
||||
rewards,
|
||||
terminated or truncated,
|
||||
infos,
|
||||
)
|
||||
elif isinstance(infos, list):
|
||||
for info, env_truncated, env_terminated in zip(
|
||||
infos, truncated, terminated
|
||||
):
|
||||
if env_truncated or env_terminated:
|
||||
info["TimeLimit.truncated"] = env_truncated and not env_terminated
|
||||
return (
|
||||
observations,
|
||||
rewards,
|
||||
np.logical_or(terminated, truncated),
|
||||
infos,
|
||||
)
|
||||
elif isinstance(infos, dict):
|
||||
if np.logical_or(np.any(truncated), np.any(terminated)):
|
||||
infos["TimeLimit.truncated"] = np.logical_and(
|
||||
truncated, np.logical_not(terminated)
|
||||
)
|
||||
return (
|
||||
observations,
|
||||
rewards,
|
||||
np.logical_or(terminated, truncated),
|
||||
infos,
|
||||
)
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Unexpected value of infos, as is_vector_envs=False, expects `info` to be a list or dict, actual type: {type(infos)}"
|
||||
)
|
||||
|
||||
|
||||
def step_api_compatibility(
|
||||
|
||||
Reference in New Issue
Block a user