Description

%load_ext autoreload
%autoreload 2

# System libraries.
import logging

# Third-party libraries.
import numpy as np
import pandas as pd

import helpers.hdbg as hdbg

_LOG = logging.getLogger(__name__)

hdbg.init_logger(verbosity=logging.INFO)
# hnotebook.config_notebook()

try:
    from IPython.display import display
except ImportError:
    display = print  # type: ignore

import gymnasium as gym
import helpers.hintrospection as hintros

# Inspect what Env exposes as a markdown list.
hintros.print_obj_info(gym.Env)

# Create the simplest classic-control environment.
env = gym.make("CartPole-v1", render_mode=None)
print("type(env):", type(env))
print("env:", env)

# Print type and content of env attributes.
print("observation_space type:", type(env.observation_space))
print("observation_space:", env.observation_space)

print("\naction_space type:", type(env.action_space))
print("action_space:", env.action_space)

print("\nspec:", env.spec)

print("\nunwrapped type:", type(env.unwrapped))
print("unwrapped:", env.unwrapped)

# Reset and inspect the return values.
observation, info = env.reset(seed=42)
print("observation:", observation)
print("observation.shape:", observation.shape)
print("observation.dtype:", observation.dtype)
print("info:", info)

# Show that reset is reproducible with the same seed.
obs1, _ = env.reset(seed=0)
obs2, _ = env.reset(seed=0)
print("Same seed → same obs:", np.array_equal(obs1, obs2))

import tutorials.gymnasium.gymnasium_utils as gymutils

env.reset(seed=42)
# Take one random action and inspect all five return values.
action = env.action_space.sample()
obs, reward, terminated, truncated, info = env.step(action)
gymutils.print_step(
    obs=obs, reward=reward, terminated=terminated, truncated=truncated, info=info
)

# Reset the env with a fixed seed for reproducibility.
env.reset(seed=42)
total_reward = 0.0
step_count = 0
done = False
# Loop until the episode ends (either terminated or truncated).
while not done:
    # Sample a random action from the action space.
    action = env.action_space.sample()
    # Advance the environment by one timestep.
    _obs, reward, terminated, truncated, _info = env.step(action)
    total_reward += reward
    step_count += 1
    # An episode ends when terminated (goal/failure) or truncated (time limit).
    done = terminated or truncated
print(f"Steps: {step_count}  Total reward: {total_reward}")

from gymnasium import spaces

# `Discrete(n)` represents a set of n integer values 0, 1, ..., n-1.
# It is used for action spaces with a finite number of discrete choices
# (e.g., move left/right in CartPole has n=2).
d = spaces.Discrete(3)
print("type:", type(d))
print("n:", d.n)
print("dtype:", d.dtype)

# Sample several times to see the range.
samples = [int(d.sample()) for _ in range(10)]
print("samples:", samples)

# Membership check.
print("2 in d:", d.contains(2))
print("5 in d:", d.contains(5))

# `Box(low, high, shape)` represents a continuous tensor where each element is
# bounded in [low, high]. It is used for observation spaces with continuous
# values (e.g., position, velocity in CartPole has shape (4,) with bounds).
b = spaces.Box(low=-1.0, high=1.0, shape=(3,), dtype=np.float32)

print("type:", type(b))
print("low:", b.low)
print("high:", b.high)
print("shape:", b.shape)
print("dtype:", b.dtype)
print("sample:", b.sample())
print(
    "contains [0.5, -0.5, 0.0]:",
    b.contains(np.array([0.5, -0.5, 0.0], dtype=np.float32)),
)
print(
    "contains [2.0, 0.0, 0.0]:",
    b.contains(np.array([2.0, 0.0, 0.0], dtype=np.float32)),
)

cart = gym.make("CartPole-v1")
mountain = gym.make("MountainCar-v0")
# Compare observation and action spaces of two classic-control envs.
for name, e in [("CartPole-v1", cart), ("MountainCar-v0", mountain)]:
    print(f"{name}:")
    print(f"  obs_space: {e.observation_space}")
    print(f"  obs_shape: {e.observation_space.shape}")
    print(f"  act_space: {e.action_space}")
    print()
cart.close()
mountain.close()

# MultiDiscrete: vector of independent discrete dims.
md = spaces.MultiDiscrete([3, 4, 2])
print("MultiDiscrete sample:", md.sample())

# MultiBinary: vector of n binary flags.
mb = spaces.MultiBinary(5)
print("MultiBinary sample:", mb.sample())

# Dict: named collection of spaces.
dict_space = spaces.Dict(
    {
        "pos": spaces.Box(low=0.0, high=1.0, shape=(2,), dtype=np.float32),
        "speed": spaces.Discrete(5),
    }
)
print("Dict sample:", dict_space.sample())

# Tuple: ordered collection of spaces.
tup = spaces.Tuple((spaces.Discrete(3), spaces.Box(0.0, 1.0, shape=(2,))))
print("Tuple sample:", tup.sample())

# Use rgb_array mode to capture a frame without opening a window.
# Note: requires `pip install "gymnasium[classic-control]"` (pygame).
try:
    env_rgb = gym.make("CartPole-v1", render_mode="rgb_array")
    env_rgb.reset(seed=0)
    frame = env_rgb.render()
    print("frame type:", type(frame))
    print("frame shape:", frame.shape)
    print("frame dtype:", frame.dtype)
    env_rgb.close()
except Exception as exc:
    print(f"render skipped ({exc})")

# Always close when finished.
env.close()
print("env closed")

# Show a sample of registered environments.
all_ids = sorted(gym.envs.registry.keys())
print(f"Total registered environments: {len(all_ids)}")
sample_ids = all_ids[:10]
display(pd.DataFrame({"id": sample_ids}))

# Inspect an EnvSpec.
spec = gym.spec("CartPole-v1")
data = {
    "field": ["id", "entry_point", "max_episode_steps", "reward_threshold"],
    "value": [
        spec.id,
        spec.entry_point,
        spec.max_episode_steps,
        spec.reward_threshold,
    ],
}
display(pd.DataFrame(data))

class OneDGridEnv(gym.Env):
    """
    Walk a 1-D grid to reach position (size-1).
    """

    def __init__(self, size: int = 5) -> None:
        super().__init__()
        # Grid length and starting position (leftmost cell).
        self._size = size
        self._pos = 0
        # Observation: current integer position (0 to size-1).
        self.observation_space = spaces.Discrete(size)
        # Actions: 0 = left, 1 = right.
        self.action_space = spaces.Discrete(2)

    def reset(self, *, seed=None, options=None):
        # Reset the RNG and place the agent at the leftmost cell.
        super().reset(seed=seed)
        self._pos = 0
        return self._pos, {}

    def step(self, action):
        # Move right (action=1) or left (action=0), clamped to grid bounds.
        self._pos += 1 if action == 1 else -1
        self._pos = int(np.clip(self._pos, 0, self._size - 1))
        # Episode terminates when the agent reaches the rightmost cell.
        terminated = self._pos == self._size - 1
        # Reward is 1 only on reaching the goal, 0 otherwise.
        reward = 1.0 if terminated else 0.0
        return self._pos, reward, terminated, False, {}


# Verify the custom env behaves like any other Gym env.
custom_env = OneDGridEnv(size=5)
obs, info = custom_env.reset(seed=7)
print("initial obs:", obs)
# Walk right until done.
step_count = 0
done = False
while not done:
    obs, reward, terminated, truncated, info = custom_env.step(1)
    print(
        f"Step {step_count}: obs={obs}, reward={reward}, terminated={terminated}"
    )
    done = terminated or truncated
    step_count += 1

# Register the custom env so it can be created via gym.make().
gym.register(id="OneDGrid-v0", entry_point=OneDGridEnv, max_episode_steps=20)
# Create the registered env with a custom size (overriding the default).
env2 = gym.make("OneDGrid-v0", size=4)
# Reset returns the starting position (0 for leftmost cell).
obs, _ = env2.reset(seed=0)
print("obs after reset:", obs)
# Step right: action=1 moves one cell to the right.
obs, reward, terminated, truncated, info = env2.step(1)
print("obs after step right:", obs, "reward:", reward)
# Clean up the env.
env2.close()

from gymnasium.wrappers import RecordEpisodeStatistics, TimeLimit

base_env = gym.make("CartPole-v1")
# Stack two wrappers.
env_w = TimeLimit(base_env, max_episode_steps=50)
env_w = RecordEpisodeStatistics(env_w)
print("type of wrapped env:", type(env_w))
print("is still Env?", isinstance(env_w, gym.Env))
# unwrapped reaches through all wrapper layers.
print("unwrapped type:", type(env_w.unwrapped))

# Run an episode and inspect the statistics injected into `info`.
obs, info = env_w.reset(seed=0)
done = False
while not done:
    obs, reward, terminated, truncated, info = env_w.step(
        env_w.action_space.sample()
    )
    # Use compact=True to print all five values on a single line.
    gymutils.print_step(
        obs=obs,
        reward=reward,
        terminated=terminated,
        truncated=truncated,
        info=info,
        compact=True,
    )
    done = terminated or truncated
# RecordEpisodeStatistics adds 'episode' key at episode end.
print("episode stats:", info.get("episode"))
env_w.close()

# Run 3 CartPole envs in parallel.
vec_env = gym.make_vec("CartPole-v1", num_envs=3, vectorization_mode="sync")
print("type:", type(vec_env))
obs_batch, info_batch = vec_env.reset(seed=0)
print("obs_batch shape:", obs_batch.shape)  # (3, 4)
# Step all 3 envs simultaneously.
actions = vec_env.action_space.sample()
print("actions batch:", actions)
obs_batch, rewards, terminations, truncations, info_batch = vec_env.step(actions)
data = {
    "field": ["obs_batch.shape", "rewards", "terminations", "truncations"],
    "value": [
        str(obs_batch.shape),
        str(rewards),
        str(terminations),
        str(truncations),
    ],
}
display(pd.DataFrame(data))
vec_env.close()

# Explore any env's interface.
env6 = gym.make("MountainCar-v0")
print("Public methods:", [m for m in dir(env6) if not m.startswith("_")])

# What does reset return?
obs, info = env6.reset()
print("obs:", obs, "  info:", info)

# What does a random step look like?
obs, reward, terminated, truncated, info = env6.step(env6.action_space.sample())
print("obs:", obs, "reward:", reward, "terminated:", terminated)

# What is the observation space bound?
print("obs low:", env6.observation_space.low)
print("obs high:", env6.observation_space.high)
env6.close()