# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
from __future__ import annotations
import gymnasium as gym
import torch
from typing import Dict, Literal, TypeVar
from omni.isaac.lab.utils import configclass
##
# Configuration.
##
[docs]@configclass
class ViewerCfg:
"""Configuration of the scene viewport camera."""
eye: tuple[float, float, float] = (7.5, 7.5, 7.5)
"""Initial camera position (in m). Default is (7.5, 7.5, 7.5)."""
lookat: tuple[float, float, float] = (0.0, 0.0, 0.0)
"""Initial camera target position (in m). Default is (0.0, 0.0, 0.0)."""
cam_prim_path: str = "/OmniverseKit_Persp"
"""The camera prim path to record images from. Default is "/OmniverseKit_Persp",
which is the default camera in the viewport.
"""
resolution: tuple[int, int] = (1280, 720)
"""The resolution (width, height) of the camera specified using :attr:`cam_prim_path`.
Default is (1280, 720).
"""
origin_type: Literal["world", "env", "asset_root"] = "world"
"""The frame in which the camera position (eye) and target (lookat) are defined in. Default is "world".
Available options are:
* ``"world"``: The origin of the world.
* ``"env"``: The origin of the environment defined by :attr:`env_index`.
* ``"asset_root"``: The center of the asset defined by :attr:`asset_name` in environment :attr:`env_index`.
"""
env_index: int = 0
"""The environment index for frame origin. Default is 0.
This quantity is only effective if :attr:`origin` is set to "env" or "asset_root".
"""
asset_name: str | None = None
"""The asset name in the interactive scene for the frame origin. Default is None.
This quantity is only effective if :attr:`origin` is set to "asset_root".
"""
##
# Types.
##
SpaceType = TypeVar("SpaceType", gym.spaces.Space, int, set, tuple, list, dict)
"""A sentinel object to indicate a valid space type to specify states, observations and actions."""
VecEnvObs = Dict[str, torch.Tensor | Dict[str, torch.Tensor]]
"""Observation returned by the environment.
The observations are stored in a dictionary. The keys are the group to which the observations belong.
This is useful for various setups such as reinforcement learning with asymmetric actor-critic or
multi-agent learning. For non-learning paradigms, this may include observations for different components
of a system.
Within each group, the observations can be stored either as a dictionary with keys as the names of each
observation term in the group, or a single tensor obtained from concatenating all the observation terms.
For example, for asymmetric actor-critic, the observation for the actor and the critic can be accessed
using the keys ``"policy"`` and ``"critic"`` respectively.
Note:
By default, most learning frameworks deal with default and privileged observations in different ways.
This handling must be taken care of by the wrapper around the :class:`ManagerBasedRLEnv` instance.
For included frameworks (RSL-RL, RL-Games, skrl), the observations must have the key "policy". In case,
the key "critic" is also present, then the critic observations are taken from the "critic" group.
Otherwise, they are the same as the "policy" group.
"""
VecEnvStepReturn = tuple[VecEnvObs, torch.Tensor, torch.Tensor, torch.Tensor, dict]
"""The environment signals processed at the end of each step.
The tuple contains batched information for each sub-environment. The information is stored in the following order:
1. **Observations**: The observations from the environment.
2. **Rewards**: The rewards from the environment.
3. **Terminated Dones**: Whether the environment reached a terminal state, such as task success or robot falling etc.
4. **Timeout Dones**: Whether the environment reached a timeout state, such as end of max episode length.
5. **Extras**: A dictionary containing additional information from the environment.
"""
AgentID = TypeVar("AgentID")
"""Unique identifier for an agent within a multi-agent environment.
The identifier has to be an immutable object, typically a string (e.g.: ``"agent_0"``).
"""
ObsType = TypeVar("ObsType", torch.Tensor, Dict[str, torch.Tensor])
"""A sentinel object to indicate the data type of the observation.
"""
ActionType = TypeVar("ActionType", torch.Tensor, Dict[str, torch.Tensor])
"""A sentinel object to indicate the data type of the action.
"""
StateType = TypeVar("StateType", torch.Tensor, dict)
"""A sentinel object to indicate the data type of the state.
"""
EnvStepReturn = tuple[
Dict[AgentID, ObsType],
Dict[AgentID, torch.Tensor],
Dict[AgentID, torch.Tensor],
Dict[AgentID, torch.Tensor],
Dict[AgentID, dict],
]
"""The environment signals processed at the end of each step.
The tuple contains batched information for each sub-environment (keyed by the agent ID).
The information is stored in the following order:
1. **Observations**: The observations from the environment.
2. **Rewards**: The rewards from the environment.
3. **Terminated Dones**: Whether the environment reached a terminal state, such as task success or robot falling etc.
4. **Timeout Dones**: Whether the environment reached a timeout state, such as end of max episode length.
5. **Extras**: A dictionary containing additional information from the environment.
"""
