diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json new file mode 100644 index 0000000..ca3af0e --- /dev/null +++ b/.devcontainer/devcontainer.json @@ -0,0 +1,57 @@ +{ + "name": "${localWorkspaceFolderBasename}", + "build": { + "context": "${localWorkspaceFolder}", + "dockerfile": "${localWorkspaceFolder}/Dockerfile", + "cacheFrom": "andrejorsula/space_robotics_bench" + }, + "workspaceFolder": "/root/ws", + "workspaceMount": "type=bind,source=${localWorkspaceFolder},target=/root/ws", + "runArgs": [ + // Network mode + "--network=host", + "--ipc=host", + // NVIDIA GPU + "--gpus=all", + // Other GPUs + "--device=/dev/dri:/dev/dri", + "--group-add=video" + ], + "mounts": [ + /// Common + // Time + "type=bind,source=/etc/localtime,target=/etc/localtime,readonly", + "type=bind,source=/etc/timezone,target=/etc/timezone,readonly", + // GUI (X11) + "type=bind,source=/tmp/.X11-unix,target=/tmp/.X11-unix", + "type=bind,source=${localEnv:TMPDIR:/tmp}/xauth_docker_vsc_${localWorkspaceFolderBasename},target=${localEnv:TMPDIR:/tmp}/xauth_docker_vsc_${localWorkspaceFolderBasename}", + /// Isaac Sim + // Data + "type=bind,source=${localEnv:HOME}/.nvidia-omniverse/data/isaac-sim,target=/root/isaac-sim/kit/data", + // Cache + "type=bind,source=${localEnv:HOME}/.cache/isaac-sim,target=/root/isaac-sim/kit/cache", + "type=bind,source=${localEnv:HOME}/.cache/nvidia/GLCache,target=/root/.cache/nvidia/GLCache", + "type=bind,source=${localEnv:HOME}/.cache/ov,target=/root/.cache/ov", + "type=bind,source=${localEnv:HOME}/.nv/ComputeCache,target=/root/.nv/ComputeCache", + // Logs + "type=bind,source=${localEnv:HOME}/.nvidia-omniverse/logs,target=/root/.nvidia-omniverse/logs", + "type=bind,source=${localEnv:HOME}/.nvidia-omniverse/logs/isaac-sim,target=/root/isaac-sim/kit/logs", + /// Project + // Cache + "type=bind,source=${localEnv:HOME}/.cache/srb,target=/root/.cache/srb" + ], + "containerEnv": { + // GUI (X11) + "DISPLAY": "${localEnv:DISPLAY}", + "XAUTHORITY": "${localEnv:TMPDIR:/tmp}/xauth_docker_vsc_${localWorkspaceFolderBasename}", + // NVIDIA GPU + "NVIDIA_VISIBLE_DEVICES": "all", + "NVIDIA_DRIVER_CAPABILITIES": "all" + }, + "initializeCommand": "XAUTH=\"${localEnv:TMPDIR:/tmp}/xauth_docker_vsc_${localWorkspaceFolderBasename}\"; touch \"${XAUTH}\"; chmod a+r \"${XAUTH}\"; XAUTH_LIST=$(xauth nlist \"${localEnv:DISPLAY}\"); if [ -n \"${XAUTH_LIST}\" ]; then echo \"${XAUTH_LIST}\" | sed -e 's/^..../ffff/' | xauth -f \"${XAUTH}\" nmerge -; fi", + "customizations": { + "vscode": { + "extensions": [] + } + } +} diff --git a/.devcontainer/open.bash b/.devcontainer/open.bash new file mode 100755 index 0000000..e84bab9 --- /dev/null +++ b/.devcontainer/open.bash @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +### Open the Dev Container in VS Code +### Usage: open.bash +set -e + +SCRIPT_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" &>/dev/null && pwd)" +REPOSITORY_DIR="$(dirname "${SCRIPT_DIR}")" + +## Determine the workspace folder +if [[ -n "$1" ]]; then + # Use the first argument as the workspace folder if provided + WORKSPACE_FOLDER="$1" +else + # Otherwise, try to extract the workspace folder from `./devcontainer.json` + WORKSPACE_FOLDER="$(grep -Po '"workspaceFolder":.*?[^\\]",' "${SCRIPT_DIR}/devcontainer.json" | cut -d'"' -f4 || true)" + if [[ -z "${WORKSPACE_FOLDER}" ]]; then + # If `./devcontainer.json` does not contain the workspace folder, default to the root + WORKSPACE_FOLDER="/" + fi +fi + +## Open the Dev Container in VS Code +CODE_REMOTE_CMD=( + code --remote + "dev-container+$(printf "%s" "${REPOSITORY_DIR}" | xxd -p | tr -d "[:space:]")" + "${WORKSPACE_FOLDER}" +) +echo -e "\033[1;90m${CODE_REMOTE_CMD[*]}\033[0m" | xargs +# shellcheck disable=SC2048 +exec ${CODE_REMOTE_CMD[*]} diff --git a/.docker/build.bash b/.docker/build.bash new file mode 100755 index 0000000..42e4b6a --- /dev/null +++ b/.docker/build.bash @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +### Build the Docker image +### Usage: build.bash [TAG] [BUILD_ARGS...] +set -e + +SCRIPT_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" &>/dev/null && pwd)" +REPOSITORY_DIR="$(dirname "${SCRIPT_DIR}")" + +## If the current user is not in the docker group, all docker commands will be run as root +if ! grep -qi /etc/group -e "docker.*${USER}"; then + echo "[INFO] The current user '${USER}' is not detected in the docker group. All docker commands will be run as root." + WITH_SUDO="sudo" +fi + +## Determine the name of the image to build +DOCKERHUB_USER="$(${WITH_SUDO} docker info 2>/dev/null | sed '/Username:/!d;s/.* //')" +PROJECT_NAME="$(basename "${REPOSITORY_DIR}")" +IMAGE_NAME="${DOCKERHUB_USER:+${DOCKERHUB_USER}/}${PROJECT_NAME}" + +## Parse TAG and forward additional build arguments +if [ "${#}" -gt "0" ]; then + if [[ "${1}" != "-"* ]]; then + IMAGE_NAME+=":${1}" + BUILD_ARGS=${*:2} + else + BUILD_ARGS=${*:1} + fi +fi + +## Build the image +DOCKER_BUILD_CMD=( + "${WITH_SUDO}" docker build + "${REPOSITORY_DIR}" + --file "${REPOSITORY_DIR}/Dockerfile" + --tag "${IMAGE_NAME}" + "${BUILD_ARGS}" +) +echo -e "\033[1;90m[TRACE] ${DOCKER_BUILD_CMD[*]}\033[0m" | xargs +# shellcheck disable=SC2048 +exec ${DOCKER_BUILD_CMD[*]} diff --git a/.docker/demo.bash b/.docker/demo.bash new file mode 100755 index 0000000..1ac7eac --- /dev/null +++ b/.docker/demo.bash @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +set -e + +SCRIPT_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" &>/dev/null && pwd)" + +## Config +# Additional volumes to mount inside the container +EXTRA_DOCKER_VOLUMES=( + "${HOME}/Videos:/root/Videos" +) +# Additional environment variables to set inside the container +EXTRA_DOCKER_ENVIRON=( + SRB_SKIP_EXT_MOD_UPDATE="1" +) + +## Parse arguments +DEFAULT_CMD="cargo run --release --package space_robotics_bench_gui" +if [ "${#}" -gt "0" ]; then + CMD=${*:1} +fi + +## Run the container +DOCKER_RUN_CMD=( + "${SCRIPT_DIR}/run.bash" + "${EXTRA_DOCKER_VOLUMES[@]/#/"-v "}" + "${EXTRA_DOCKER_ENVIRON[@]/#/"-e "}" + "${CMD:-${DEFAULT_CMD}}" +) +echo -e "\033[1;90m[TRACE] ${DOCKER_RUN_CMD[*]}\033[0m" | xargs +# shellcheck disable=SC2048 +exec ${DOCKER_RUN_CMD[*]} diff --git a/.docker/dev.bash b/.docker/dev.bash new file mode 100755 index 0000000..37bdbab --- /dev/null +++ b/.docker/dev.bash @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +### Run a Docker container with additional development volumes mounted +### Usage: dev.bash [-v HOST_DIR:DOCKER_DIR:OPTIONS] [-e ENV=VALUE] [TAG] [CMD] +set -e + +SCRIPT_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" &>/dev/null && pwd)" +REPOSITORY_DIR="$(dirname "${SCRIPT_DIR}")" +WS_DIR="$(dirname "${REPOSITORY_DIR}")" + +## Config +# Development volumes to mount inside the container +DOCKER_DEV_VOLUMES=( + "${WS_DIR}/isaaclab:/root/isaaclab:rw" + # "${WS_DIR}/dreamerv3:/root/dreamerv3:rw" +) +# Development environment variables to set inside the container +DOCKER_DEV_ENVIRON=( + SRB_WITH_TRACEBACK="${SRB_WITH_TRACEBACK:-true}" +) + +## Run the container with development volumes +DOCKER_DEV_CMD=( + WITH_DEV_VOLUME=true "${SCRIPT_DIR}/run.bash" + "${DOCKER_DEV_VOLUMES[@]/#/"-v "}" + "${DOCKER_DEV_ENVIRON[@]/#/"-e "}" + "${*:1}" +) +echo -e "\033[1;90m[TRACE] ${DOCKER_DEV_CMD[*]}\033[0m" | xargs +# shellcheck disable=SC2048 +exec ${DOCKER_DEV_CMD[*]} diff --git a/.docker/host/install_docker.bash b/.docker/host/install_docker.bash new file mode 100755 index 0000000..38f7357 --- /dev/null +++ b/.docker/host/install_docker.bash @@ -0,0 +1,130 @@ +#!/usr/bin/env bash +### Install Docker and NVIDIA Container Toolkit +### Usage: install_docker.bash +set -e + +## Install curl if missing +if ! command -v curl >/dev/null 2>&1; then + if command -v apt-get >/dev/null 2>&1; then + sudo apt-get install -y curl + elif command -v dnf >/dev/null 2>&1; then + sudo dnf install -y curl + elif command -v yum >/dev/null 2>&1; then + sudo yum install -y curl + fi +fi + +## Install Docker via the convenience script +curl -fsSL https://get.docker.com | sh +sudo systemctl enable --now docker + +## (Optional) Install support for NVIDIA if an NVIDIA GPU is detected and the installation is requested +check_nvidia_gpu() { + if ! lshw -C display 2>/dev/null | grep -qi "vendor.*nvidia"; then + return 1 # NVIDIA GPU is not present + elif ! command -v nvidia-smi >/dev/null 2>&1; then + return 1 # NVIDIA GPU is present but nvidia-utils not installed + elif ! nvidia-smi -L &>/dev/null; then + return 1 # NVIDIA GPU is present but is not working properly + else + return 0 # NVIDIA GPU is present and appears to be working + fi +} +configure_nvidia_apt_repository() { + curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | + sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg && + curl -sL https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | + sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | + sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list + sudo apt-get update +} +if check_nvidia_gpu; then + echo -e "[INFO] NVIDIA GPU detected." + DOCKER_VERSION="$(sudo docker version --format '{{.Server.Version}}' 2>/dev/null)" + MIN_VERSION_FOR_TOOLKIT="19.3" + if [ "$(printf '%s\n' "${MIN_VERSION_FOR_TOOLKIT}" "${DOCKER_VERSION}" | sort -V | head -n1)" = "$MIN_VERSION_FOR_TOOLKIT" ]; then + if ! command -v nvidia-container-toolkit >/dev/null 2>&1; then + while true; do + read -erp "Do you want to install NVIDIA Container Toolkit? [Y/n]: " INSTALL_NVIDIA_CONTAINER_TOOLKIT + case "${INSTALL_NVIDIA_CONTAINER_TOOLKIT,,}" in + "" | y | yes) + INSTALL_NVIDIA_CONTAINER_TOOLKIT=true + break + ;; + n | no) + INSTALL_NVIDIA_CONTAINER_TOOLKIT=false + break + ;; + esac + done + if [[ "${INSTALL_NVIDIA_CONTAINER_TOOLKIT}" = true ]]; then + if command -v apt-get >/dev/null 2>&1; then + configure_nvidia_apt_repository + sudo apt-get install -y nvidia-container-toolkit + elif command -v yum >/dev/null 2>&1; then + curl -s -L https://nvidia.github.io/libnvidia-container/stable/rpm/nvidia-container-toolkit.repo | + sudo tee /etc/yum.repos.d/nvidia-container-toolkit.repo + sudo yum install -y nvidia-container-toolkit + else + echo >&2 -e "\033[1;31m[ERROR] Supported package manager not found. Please install nvidia-container-toolkit manually.\033[0m" + fi + sudo systemctl restart --now docker + fi + else + echo -e "[INFO] NVIDIA Container Toolkit is already installed." + fi + else + if ! command -v nvidia-docker >/dev/null 2>&1; then + while true; do + read -erp "Do you want to install NVIDIA Docker [Y/n]: " INSTALL_NVIDIA_DOCKER + case "${INSTALL_NVIDIA_DOCKER,,}" in + "" | y | yes) + INSTALL_NVIDIA_DOCKER=true + break + ;; + n | no) + INSTALL_NVIDIA_DOCKER=false + break + ;; + esac + done + if [[ "${INSTALL_NVIDIA_DOCKER}" = true ]]; then + if command -v apt-get >/dev/null 2>&1; then + configure_nvidia_apt_repository + sudo apt-get install -y nvidia-docker2 + else + echo >&2 -e "\033[1;31m[ERROR] Supported package manager not found. Please install nvidia-docker2 manually.\033[0m" + fi + sudo systemctl restart --now docker + fi + else + echo -e "[INFO] NVIDIA Docker is already installed." + fi + fi +fi + +if [[ $(grep /etc/group -e "docker") != *"${USER}"* ]]; then + [ -z "${PS1}" ] + ## (Optional) Add user to docker group + while true; do + read -erp "Do you want to add the current user ${USER} to the docker group? [Y/n]: " ADD_USER_TO_DOCKER_GROUP + case "${ADD_USER_TO_DOCKER_GROUP,,}" in + "" | y | yes) + ADD_USER_TO_DOCKER_GROUP=true + break + ;; + n | no) + ADD_USER_TO_DOCKER_GROUP=false + break + ;; + esac + done + if [[ "${ADD_USER_TO_DOCKER_GROUP}" = true ]]; then + sudo groupadd -f docker + sudo usermod -aG docker "${USER}" + newgrp docker + echo -e "[INFO] The current user ${USER} was added to the docker group." + fi +else + echo -e "[INFO] The current user ${USER} is already in the docker group." +fi diff --git a/.docker/hpc/build.bash b/.docker/hpc/build.bash new file mode 100755 index 0000000..1bd604a --- /dev/null +++ b/.docker/hpc/build.bash @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +### Build a Singularity image from the Docker image +### Usage: build.bash [TAG] [BUILD_ARGS...] +set -e + +SCRIPT_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" &>/dev/null && pwd)" +DOT_DOCKER_DIR="$(dirname "${SCRIPT_DIR}")" +REPOSITORY_DIR="$(dirname "${DOT_DOCKER_DIR}")" +IMAGES_DIR="${SCRIPT_DIR}/images" +DOCKER_BUILD_SCRIPT="${DOT_DOCKER_DIR}/build.bash" +export APPTAINER_TMPDIR="${APPTAINER_TMPDIR:-"${HOME}/.apptainer/tmp"}" + +## If the current user is not in the docker group, all docker commands will be run as root +if ! grep -qi /etc/group -e "docker.*${USER}"; then + echo "[INFO] The current user '${USER}' is not detected in the docker group. All docker commands will be run as root." + WITH_SUDO="sudo" +fi + +## Determine the name of the image to build and the output path +DOCKERHUB_USER="$(${WITH_SUDO} docker info 2>/dev/null | sed '/Username:/!d;s/.* //')" +PROJECT_NAME="$(basename "${REPOSITORY_DIR}")" +IMAGE_NAME="${DOCKERHUB_USER:+${DOCKERHUB_USER}/}${PROJECT_NAME}" +OUTPUT_PATH="${IMAGES_DIR}/${PROJECT_NAME}.sif" + +## Parse TAG and forward additional build arguments +if [ "${#}" -gt "0" ]; then + if [[ "${1}" != "-"* ]]; then + TAG="${1}" + BUILD_ARGS=${*:2} + else + BUILD_ARGS=${*:1} + fi +fi +TAG="${TAG:-"latest"}" +IMAGE_NAME+=":${TAG}" + +## Create the temporary directory for the Singularity image +mkdir -p "${APPTAINER_TMPDIR}" + +## Build the Docker image +"${DOCKER_BUILD_SCRIPT}" "${TAG}" "${BUILD_ARGS}" + +## Convert the Docker image to a Singularity image +APPTAINER_BUILD_CMD=( + apptainer build + "${OUTPUT_PATH}" + "docker-daemon:${IMAGE_NAME}" +) +echo -e "\033[1;90m[TRACE] ${APPTAINER_BUILD_CMD[*]}\033[0m" | xargs +# shellcheck disable=SC2048 +exec ${APPTAINER_BUILD_CMD[*]} diff --git a/.docker/hpc/run.bash b/.docker/hpc/run.bash new file mode 100755 index 0000000..d226076 --- /dev/null +++ b/.docker/hpc/run.bash @@ -0,0 +1,94 @@ +#!/usr/bin/env bash +### Run a command inside the Singularity container +### Usage: run.bash [CMD] +set -e + +SCRIPT_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" &>/dev/null && pwd)" +DOT_DOCKER_DIR="$(dirname "${SCRIPT_DIR}")" +REPOSITORY_DIR="$(dirname "${DOT_DOCKER_DIR}")" +IMAGES_DIR="${SCRIPT_DIR}/images" +PROJECT_NAME="$(basename "${REPOSITORY_DIR}")" + +## Configuration +# Path to the image to run +IMAGE_PATH="${IMAGE_PATH:-"${IMAGES_DIR}/${PROJECT_NAME}.sif"}" +# Options for running the container +SINGULARITY_EXEC_OPTS="${SINGULARITY_EXEC_OPTS:- + --no-home +}" +# Flag to enable GPU +WITH_GPU="${WITH_GPU:-true}" +# Volumes to mount inside the container +VOLUMES_ROOT="${SCRATCH:-${HOME}}/volumes/${PROJECT_NAME}" +SINGULARITY_VOLUMES=( + ## Isaac Sim + # Data + "${VOLUMES_ROOT}/.nvidia-omniverse/data/isaac-sim:/root/isaac-sim/kit/data:rw" + # Cache + "${VOLUMES_ROOT}/.cache/isaac-sim:/root/isaac-sim/kit/cache:rw" + "${VOLUMES_ROOT}/.cache/nvidia/GLCache:/root/.cache/nvidia/GLCache:rw" + "${VOLUMES_ROOT}/.cache/ov:/root/.cache/ov:rw" + "${VOLUMES_ROOT}/.nv/ComputeCache:/root/.nv/ComputeCache:rw" + # Logs + "${VOLUMES_ROOT}/.nvidia-omniverse/logs:/root/.nvidia-omniverse/logs:rw" + "${VOLUMES_ROOT}/.nvidia-omniverse/logs/isaac-sim:/root/isaac-sim/kit/logs:rw" + ## Project + # Source + "${REPOSITORY_DIR}:/root/ws:rw" + # Cache + "${VOLUMES_ROOT}/.cache/srb:/root/.cache/srb:rw" + + + "${VOLUMES_ROOT}/home/users:/home/users:rw" +) + +## Ensure the image exists +if [ ! -f "${IMAGE_PATH}" ]; then + echo >&2 -e "\033[1;31m[ERROR] Singularity image not found at ${IMAGE_PATH}\033[0m" + exit 1 +fi + +## Parse CMD +if [ "${#}" -gt "0" ]; then + CMD=${*:1} +else + echo >&2 -e "\033[1;31m[ERROR] No command provided.\033[0m" + exit 1 +fi + +## Ensure the host directories exist +for volume in "${SINGULARITY_VOLUMES[@]}"; do + if [[ "${volume}" =~ ^([^:]+):([^:]+).*$ ]]; then + host_dir="${BASH_REMATCH[1]}" + if [ ! -d "${host_dir}" ]; then + mkdir -p "${host_dir}" + echo -e "\033[1;90m[INFO] Created directory ${host_dir}\033[0m" + fi + fi +done + +## GPU +if [[ "${WITH_GPU,,}" = true ]]; then + SINGULARITY_EXEC_OPTS+=" --nv" +fi + +## Environment +if ! command -v module >/dev/null 2>&1; then + echo >&2 -e "\033[1;31m[ERROR] The 'module' command is not available. Please run this script on a compute node.\033[0m" + exit 1 +fi +# Load the Singularity module +module purge +module load tools/Singularity + +## Run the container +SINGULARITY_EXEC_CMD=( + singularity exec + "${SINGULARITY_EXEC_OPTS}" + "${SINGULARITY_VOLUMES[@]/#/"--bind "}" + "${IMAGE_PATH}" + "${CMD}" +) +echo -e "\033[1;90m[TRACE] ${SINGULARITY_EXEC_CMD[*]}\033[0m" | xargs +# shellcheck disable=SC2048 +exec ${SINGULARITY_EXEC_CMD[*]} diff --git a/.docker/hpc/submit.bash b/.docker/hpc/submit.bash new file mode 100755 index 0000000..82d12b8 --- /dev/null +++ b/.docker/hpc/submit.bash @@ -0,0 +1,65 @@ +#!/usr/bin/env bash +### Submit a job using Slurm +### Usage: submit.bash [CMD] +set -e + +SCRIPT_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" &>/dev/null && pwd)" +DOT_DOCKER_DIR="$(dirname "${SCRIPT_DIR}")" +REPOSITORY_DIR="$(dirname "${DOT_DOCKER_DIR}")" +IMAGES_DIR="${SCRIPT_DIR}/images" +PROJECT_NAME="$(basename "${REPOSITORY_DIR}")" + +## Configuration +# Path to the image to run +IMAGE_PATH="${IMAGE_PATH:-"${IMAGES_DIR}/${PROJECT_NAME}.sif"}" +JOBS_DIR="${JOBS_DIR:-"${HOME}/jobs"}" + +## Ensure the image exists +if [ ! -f "${IMAGE_PATH}" ]; then + echo >&2 -e "\033[1;31m[ERROR] Singularity image not found at ${IMAGE_PATH}\033[0m" + exit 1 +fi + +## Parse CMD +if [ "${#}" -gt "0" ]; then + CMD=${*:1} +else + echo >&2 -e "\033[1;31m[ERROR] No command provided.\033[0m" + exit 1 +fi + +# Create a job file directory if it does not exist +if [ ! -d "${JOBS_DIR}" ]; then + mkdir -p "${JOBS_DIR}" + echo -e "\033[1;90m[INFO] Created directory ${JOBS_DIR}\033[0m" +fi + +# Create a job file +TIMESTAMP="$(date +%Y-%m-%d_%H-%M-%S)" +JOB_PATH="${JOBS_DIR}/job_${TIMESTAMP}.bash" +cat <"${JOB_PATH}" +#!/bin/bash -l +#SBATCH --job-name="job_${TIMESTAMP}" +#SBATCH --time=1-12:00:00 +#SBATCH --qos=normal +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --cpus-per-task=3 +#SBATCH --mem=168GB +#SBATCH --partition=gpu +#SBATCH --constraint volta32 +#SBATCH --gpus=1 +#SBATCH --gpus-per-node=1 +#SBATCH --gpus-per-task=1 +#SBATCH --signal=B:SIGKILL@60 + +bash "${SCRIPT_DIR}/run.bash" "${CMD}" +EOT + +## Submit the job +SBATCH_CMD=( + sbatch "${JOB_PATH}" +) +echo -e "\033[1;90m[TRACE] ${SBATCH_CMD[*]}\033[0m" | xargs +# shellcheck disable=SC2048 +exec ${SBATCH_CMD[*]} diff --git a/.docker/join.bash b/.docker/join.bash new file mode 100755 index 0000000..eb9fb7a --- /dev/null +++ b/.docker/join.bash @@ -0,0 +1,96 @@ +#!/usr/bin/env bash +### Join a running Docker container +### Usage: join.bash [ID] [CMD] +set -e + +SCRIPT_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" &>/dev/null && pwd)" +REPOSITORY_DIR="$(dirname "${SCRIPT_DIR}")" + +## If the current user is not in the docker group, all docker commands will be run as root +if ! grep -qi /etc/group -e "docker.*${USER}"; then + echo "[INFO] The current user '${USER}' is not detected in the docker group. All docker commands will be run as root." + WITH_SUDO="sudo" +fi + +## Config +# Name of the Docker image to join if an image with locally-defined name does not exist +DOCKERHUB_IMAGE_NAME="${DOCKERHUB_IMAGE_NAME:-"andrejorsula/space_robotics_bench"}" +# Options for executing a command inside the container +DOCKER_EXEC_OPTS="${DOCKER_EXEC_OPTS:- + --interactive + --tty +}" +# Default command to execute inside the container +DEFAULT_CMD="${DEFAULT_CMD:-"bash"}" + +## Parse ID and CMD +if [ "${#}" -gt "0" ]; then + if [[ "${1}" =~ ^[0-9]+$ ]]; then + ID="${1}" + if [ "${#}" -gt "1" ]; then + CMD=${*:2} + else + CMD="${DEFAULT_CMD}" + fi + else + CMD=${*:1} + fi +else + CMD="${DEFAULT_CMD}" +fi + +## Determine the name of the container to join +DOCKERHUB_USER="$(${WITH_SUDO} docker info 2>/dev/null | sed '/Username:/!d;s/.* //')" +PROJECT_NAME="$(basename "${REPOSITORY_DIR}")" +IMAGE_NAME="${DOCKERHUB_USER:+${DOCKERHUB_USER}/}${PROJECT_NAME}" +if [[ -z "$(${WITH_SUDO} docker images -q "${IMAGE_NAME}" 2>/dev/null)" ]] && [[ -n "$(curl -fsSL "https://registry.hub.docker.com/v2/repositories/${DOCKERHUB_IMAGE_NAME}" 2>/dev/null)" ]]; then + IMAGE_NAME="${DOCKERHUB_IMAGE_NAME}" +fi +CONTAINER_NAME="${IMAGE_NAME##*/}" +CONTAINER_NAME="${CONTAINER_NAME//[^a-zA-Z0-9]/_}" + +## Verify/select the appropriate container to join +RELEVANT_CONTAINERS=$(${WITH_SUDO} docker container list --all --format "{{.Names}}" | grep -i "${CONTAINER_NAME}" || :) +RELEVANT_CONTAINERS_COUNT=$(echo "${RELEVANT_CONTAINERS}" | wc -w) +if [ "${RELEVANT_CONTAINERS_COUNT}" -eq "0" ]; then + echo >&2 -e "\033[1;31m[ERROR] No containers with the name '${CONTAINER_NAME}' found. Run the container first.\033[0m" + exit 1 +elif [ "${RELEVANT_CONTAINERS_COUNT}" -eq "1" ]; then + CONTAINER_NAME="${RELEVANT_CONTAINERS}" +else + print_usage_with_relevant_containers() { + echo >&2 "Usage: ${0} [ID] [CMD]" + echo "${RELEVANT_CONTAINERS}" | sort --version-sort | while read -r container; do + id=$(echo "${container}" | grep -oE '[0-9]+$' || :) + if [ -z "${id}" ]; then + id=0 + fi + echo >&2 -e " ${container}\t(ID=${id})" + done + } + if [[ -n "${ID}" ]]; then + if [ "${ID}" -gt "0" ]; then + CONTAINER_NAME="${CONTAINER_NAME}${ID}" + fi + if ! echo "${RELEVANT_CONTAINERS}" | grep -qi "${CONTAINER_NAME}"; then + echo >&2 -e "\033[1;31m[ERROR] Container with 'ID=${ID}' does not exist. Specify the correct ID as the first argument.\033[0m" + print_usage_with_relevant_containers + exit 2 + fi + else + echo >&2 -e "\033[1;31m[ERROR] Multiple containers with the name '${CONTAINER_NAME}' found. ID of the container must be specified as the first argument.\033[0m" + print_usage_with_relevant_containers + exit 2 + fi +fi + +## Execute command inside the container +DOCKER_EXEC_CMD=( + "${WITH_SUDO}" docker exec + "${DOCKER_EXEC_OPTS}" + "${CONTAINER_NAME}" + "${CMD}" +) +echo -e "\033[1;90m[TRACE] ${DOCKER_EXEC_CMD[*]}\033[0m" | xargs +# shellcheck disable=SC2048 +exec ${DOCKER_EXEC_CMD[*]} diff --git a/.docker/run.bash b/.docker/run.bash new file mode 100755 index 0000000..0bae2b0 --- /dev/null +++ b/.docker/run.bash @@ -0,0 +1,201 @@ +#!/usr/bin/env bash +### Run a Docker container +### Usage: run.bash [-v HOST_DIR:DOCKER_DIR:OPTIONS] [-e ENV=VALUE] [TAG] [CMD] +set -e + +SCRIPT_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" &>/dev/null && pwd)" +REPOSITORY_DIR="$(dirname "${SCRIPT_DIR}")" + +## If the current user is not in the docker group, all docker commands will be run as root +if ! grep -qi /etc/group -e "docker.*${USER}"; then + echo "[INFO] The current user '${USER}' is not detected in the docker group. All docker commands will be run as root." + WITH_SUDO="sudo" +fi + +## Config +# Name of the Docker image to run if an image with locally-defined name does not exist +DOCKERHUB_IMAGE_NAME="${DOCKERHUB_IMAGE_NAME:-"andrejorsula/space_robotics_bench"}" +# Options for running the container +DOCKER_RUN_OPTS="${DOCKER_RUN_OPTS:- + --interactive + --tty + --rm + --network host + --ipc host + --privileged +}" +# Flag to enable GPU +WITH_GPU="${WITH_GPU:-true}" +# Flag to enable GUI (X11) +WITH_GUI="${WITH_GUI:-true}" +# Flag to enable mounting the source code as a volume +WITH_DEV_VOLUME="${WITH_DEV_VOLUME:-true}" +# Volumes to mount inside the container +DOCKER_VOLUMES=( + ## Common + # Time + "/etc/localtime:/etc/localtime:ro" + "/etc/timezone:/etc/timezone:ro" + # Hot-pluggable IO devices + "/dev:/dev" + ## Isaac Sim + # Data + "${HOME}/.nvidia-omniverse/data/isaac-sim:/root/isaac-sim/kit/data:rw" + # Cache + "${HOME}/.cache/isaac-sim:/root/isaac-sim/kit/cache:rw" + "${HOME}/.cache/nvidia/GLCache:/root/.cache/nvidia/GLCache:rw" + "${HOME}/.cache/ov:/root/.cache/ov:rw" + "${HOME}/.nv/ComputeCache:/root/.nv/ComputeCache:rw" + # Logs + "${HOME}/.nvidia-omniverse/logs:/root/.nvidia-omniverse/logs:rw" + "${HOME}/.nvidia-omniverse/logs/isaac-sim:/root/isaac-sim/kit/logs:rw" + ## Project + # Cache + "${HOME}/.cache/srb:/root/.cache/srb:rw" +) +if [[ "${WITH_DEV_VOLUME,,}" = true ]]; then + DOCKER_VOLUMES+=( + "${REPOSITORY_DIR}:/root/ws:rw" + ) +fi +# Environment variables to set inside the container +DOCKER_ENVIRON=( + ROS_DOMAIN_ID="${ROS_DOMAIN_ID:-"0"}" + ROS_LOCALHOST_ONLY="${ROS_LOCALHOST_ONLY:-"1"}" + RMW_IMPLEMENTATION="${RMW_IMPLEMENTATION:-"rmw_cyclonedds_cpp"}" +) + +## DDS config +if [ -n "${CYCLONEDDS_URI}" ]; then + DOCKER_VOLUMES+=("${CYCLONEDDS_URI//file:\/\//}:/root/.ros/cyclonedds.xml:ro") + DOCKER_ENVIRON+=("CYCLONEDDS_URI=file:///root/.ros/cyclonedds.xml") +fi +if [ -n "${FASTRTPS_DEFAULT_PROFILES_FILE}" ]; then + DOCKER_VOLUMES+=("${FASTRTPS_DEFAULT_PROFILES_FILE}:/root/.ros/fastrtps.xml:ro") + DOCKER_ENVIRON+=("FASTRTPS_DEFAULT_PROFILES_FILE=/root/.ros/fastrtps.xml") +fi + +## Determine the name of the image to run +DOCKERHUB_USER="$(${WITH_SUDO} docker info 2>/dev/null | sed '/Username:/!d;s/.* //')" +PROJECT_NAME="$(basename "${REPOSITORY_DIR}")" +IMAGE_NAME="${DOCKERHUB_USER:+${DOCKERHUB_USER}/}${PROJECT_NAME}" +if [[ -z "$(${WITH_SUDO} docker images -q "${IMAGE_NAME}" 2>/dev/null)" ]] && [[ -n "$(curl -fsSL "https://registry.hub.docker.com/v2/repositories/${DOCKERHUB_IMAGE_NAME}" 2>/dev/null)" ]]; then + IMAGE_NAME="${DOCKERHUB_IMAGE_NAME}" +fi + +## Generate a unique container name +CONTAINER_NAME="${IMAGE_NAME##*/}" +CONTAINER_NAME="${CONTAINER_NAME//[^a-zA-Z0-9]/_}" +ALL_CONTAINER_NAMES=$(${WITH_SUDO} docker container list --all --format "{{.Names}}") +if echo "${ALL_CONTAINER_NAMES}" | grep -qi "${CONTAINER_NAME}"; then + ID=1 + while echo "${ALL_CONTAINER_NAMES}" | grep -qi "${CONTAINER_NAME}${ID}"; do + ID=$((ID + 1)) + done + CONTAINER_NAME="${CONTAINER_NAME}${ID}" +fi +DOCKER_RUN_OPTS+=" --name ${CONTAINER_NAME}" + +## Parse volumes and environment variables +while getopts ":v:e:" opt; do + case "${opt}" in + v) DOCKER_VOLUMES+=("${OPTARG}") ;; + e) DOCKER_ENVIRON+=("${OPTARG}") ;; + *) + echo >&2 "Usage: ${0} [-v HOST_DIR:DOCKER_DIR:OPTIONS] [-e ENV=VALUE] [TAG] [CMD]" + exit 2 + ;; + esac +done +shift "$((OPTIND - 1))" + +## Parse TAG and forward CMD arguments +if [ "${#}" -gt "0" ]; then + if [[ $(${WITH_SUDO} docker images --format "{{.Tag}}" "${IMAGE_NAME}") =~ (^|[[:space:]])${1}($|[[:space:]]) || $(curl -fsSL "https://registry.hub.docker.com/v2/repositories/${IMAGE_NAME}/tags" 2>/dev/null | grep -Poe '(?<=(\"name\":\")).*?(?=\")') =~ (^|[[:space:]])${1}($|[[:space:]]) ]]; then + IMAGE_NAME+=":${1}" + CMD=${*:2} + else + CMD=${*:1} + fi +fi + +## GPU +if [[ "${WITH_GPU,,}" = true ]]; then + check_nvidia_gpu() { + if [[ -n "${WITH_GPU_FORCE_NVIDIA}" ]]; then + if [[ "${WITH_GPU_FORCE_NVIDIA,,}" = true ]]; then + echo "[INFO] NVIDIA GPU is force-enabled via 'WITH_GPU_FORCE_NVIDIA=true'." + return 0 # NVIDIA GPU is force-enabled + else + echo "[INFO] NVIDIA GPU is force-disabled via 'WITH_GPU_FORCE_NVIDIA=false'." + return 1 # NVIDIA GPU is force-disabled + fi + elif ! lshw -C display 2>/dev/null | grep -qi "vendor.*nvidia"; then + return 1 # NVIDIA GPU is not present + elif ! command -v nvidia-smi >/dev/null 2>&1; then + echo >&2 -e "\e[33m[WARNING] NVIDIA GPU is detected, but its functionality cannot be verified. This container will not be able to use the GPU. Please install nvidia-utils on the host system or force-enable NVIDIA GPU via 'WITH_GPU_FORCE_NVIDIA=true'.\e[0m" + return 1 # NVIDIA GPU is present but nvidia-utils not installed + elif ! nvidia-smi -L &>/dev/null; then + echo >&2 -e "\e[33m[WARNING] NVIDIA GPU is detected, but it does not seem to be working properly. This container will not be able to use the GPU. Please ensure the NVIDIA drivers are properly installed on the host system.\e[0m" + return 1 # NVIDIA GPU is present but is not working properly + else + return 0 # NVIDIA GPU is present and appears to be working + fi + } + if check_nvidia_gpu; then + # Enable GPU either via NVIDIA Container Toolkit or NVIDIA Docker (depending on Docker version) + DOCKER_VERSION="$(${WITH_SUDO} docker version --format '{{.Server.Version}}')" + MIN_VERSION_FOR_TOOLKIT="19.3" + if [ "$(printf '%s\n' "${MIN_VERSION_FOR_TOOLKIT}" "${DOCKER_VERSION}" | sort -V | head -n1)" = "$MIN_VERSION_FOR_TOOLKIT" ]; then + DOCKER_RUN_OPTS+=" --gpus all" + else + DOCKER_RUN_OPTS+=" --runtime nvidia" + fi + DOCKER_ENVIRON+=( + NVIDIA_VISIBLE_DEVICES="all" + NVIDIA_DRIVER_CAPABILITIES="all" + ) + fi + if [[ -e /dev/dri ]]; then + DOCKER_RUN_OPTS+=" --device=/dev/dri:/dev/dri" + if [[ $(getent group video) ]]; then + DOCKER_RUN_OPTS+=" --group-add video" + fi + fi +fi + +## GUI +if [[ "${WITH_GUI,,}" = true ]]; then + # To enable GUI, make sure processes in the container can connect to the x server + XAUTH="${TMPDIR:-"/tmp"}/xauth_docker_${PROJECT_NAME}" + touch "${XAUTH}" + chmod a+r "${XAUTH}" + XAUTH_LIST=$(xauth nlist "${DISPLAY}") + if [ -n "${XAUTH_LIST}" ]; then + echo "${XAUTH_LIST}" | sed -e 's/^..../ffff/' | xauth -f "${XAUTH}" nmerge - + fi + # GUI-enabling volumes + DOCKER_VOLUMES+=( + "${XAUTH}:${XAUTH}" + "/tmp/.X11-unix:/tmp/.X11-unix" + "/dev/input:/dev/input" + ) + # GUI-enabling environment variables + DOCKER_ENVIRON+=( + DISPLAY="${DISPLAY}" + XAUTHORITY="${XAUTH}" + ) +fi + +## Run the container +DOCKER_RUN_CMD=( + "${WITH_SUDO}" docker run + "${DOCKER_RUN_OPTS}" + "${DOCKER_VOLUMES[@]/#/"--volume "}" + "${DOCKER_ENVIRON[@]/#/"--env "}" + "${IMAGE_NAME}" + "${CMD}" +) +echo -e "\033[1;90m[TRACE] ${DOCKER_RUN_CMD[*]}\033[0m" | xargs +# shellcheck disable=SC2048 +exec ${DOCKER_RUN_CMD[*]} diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..d86a7e0 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,36 @@ +## Python +**/__pycache__/ +**/.pytest_cache/ +**/*.pcd +**/*.py[cod] +**/*.pyc +**/*$py.class + +## Rust +target/ +**/*.rs.bk +**/*.pdb +**/.cargo/config.toml +!.cargo/config.toml + +## Shared Objects +**/*.so + +## Dump files +**/core.* + +## Logs +logs/ +*.log + +## Apptainer +**/*.sif + +## Docker +Dockerfile* +.docker/ +.dockerignore + +## Git +.git/ +.gitignore diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000..a1cd7d2 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,11 @@ +version: 2 +updates: + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "weekly" + + - package-ecosystem: "cargo" + directory: "/" + schedule: + interval: "weekly" diff --git a/.github/workflows/dependabot.yml b/.github/workflows/dependabot.yml new file mode 100644 index 0000000..17b7161 --- /dev/null +++ b/.github/workflows/dependabot.yml @@ -0,0 +1,25 @@ +name: Dependabot automation +on: pull_request + +permissions: + contents: write + pull-requests: write + +jobs: + dependabot_automation: + runs-on: ubuntu-latest + if: ${{ github.actor == 'dependabot[bot]' }} + env: + GH_TOKEN: ${{secrets.GITHUB_TOKEN}} + PR_URL: ${{github.event.pull_request.html_url}} + steps: + - name: Fetch metadata + id: metadata + uses: dependabot/fetch-metadata@v2 + with: + github-token: "${{ secrets.GITHUB_TOKEN }}" + - name: Approve the PR + run: gh pr review --approve "$PR_URL" + continue-on-error: true + - name: Enable auto-merge for the PR + run: gh pr merge --auto --squash "$PR_URL" diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml new file mode 100644 index 0000000..65eb15b --- /dev/null +++ b/.github/workflows/docker.yml @@ -0,0 +1,90 @@ +name: Docker + +on: + push: + branches: + - main + paths-ignore: + - "docs/**" + pull_request: + release: + types: [published] + workflow_dispatch: + +concurrency: + group: "${{ github.workflow }} @ ${{ github.event.pull_request.head.label || github.head_ref || github.ref }}" + cancel-in-progress: true + +jobs: + docker: + permissions: + packages: write + contents: read + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + dockerfile: + - Dockerfile + steps: + ## Free up space by removing unnecessary files + - name: Maximize build space + uses: AdityaGarg8/remove-unwanted-software@v4 + with: + remove-android: "true" + remove-cached-tools: "true" + remove-codeql: "true" + remove-docker-images: "true" + remove-dotnet: "true" + remove-haskell: "true" + remove-large-packages: "true" + + - uses: actions/checkout@v4 + + ## Login to Docker Hub and/or GitHub container registry + - name: Login to GitHub container registry + if: ${{ github.event_name == 'release' }} + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + - name: Login to Docker Hub + env: + DOCKERHUB_USERNAME: ${{ secrets.DOCKERHUB_USERNAME }} + DOCKERHUB_PASSWORD: ${{ secrets.DOCKERHUB_PASSWORD }} + if: ${{ github.event_name == 'release' && env.DOCKERHUB_USERNAME != '' && env.DOCKERHUB_PASSWORD != '' }} + id: login-dockerhub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_PASSWORD }} + - name: Login to NGC + env: + NVCR_USERNAME: ${{ secrets.NVCR_USERNAME }} + NVCR_PASSWORD: ${{ secrets.NVCR_PASSWORD }} + if: ${{ env.NVCR_USERNAME != '' && env.NVCR_PASSWORD != '' }} + uses: docker/login-action@v3 + with: + registry: nvcr.io + username: ${{ secrets.NVCR_USERNAME }} + password: ${{ secrets.NVCR_PASSWORD }} + + ## Extract metadata (tags, labels) from Git reference and GitHub events for the Docker image + - name: Extract metadata + id: metadata + uses: docker/metadata-action@v5 + with: + images: | + name=ghcr.io/${{ github.repository }},enable=true + name=${{ github.repository }},enable=${{ steps.login-dockerhub.outcome == 'success' }} + + ## Build and push if the workflow was triggered by a release + - name: Build (and push on release) + uses: docker/build-push-action@v6 + with: + context: . + file: ${{ matrix.dockerfile }} + push: ${{ github.event_name == 'release' && matrix.dockerfile == 'Dockerfile' }} + tags: ${{ steps.metadata.outputs.tags }} + labels: ${{ steps.metadata.outputs.labels }} diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 0000000..868b988 --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,50 @@ +name: Docs + +on: + push: + branches: + - main + paths: + - ".github/workflows/docs.yml" + - "docs/**" + pull_request: + release: + types: [published] + workflow_dispatch: + +permissions: + contents: read + pages: write + id-token: write + +concurrency: + group: "${{ github.workflow }} @ ${{ github.event.pull_request.head.label || github.head_ref || github.ref }}" + cancel-in-progress: true + +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Setup mdBook + uses: peaceiris/actions-mdbook@v2 + with: + mdbook-version: latest + - name: Build + run: mdbook build docs + - name: Upload artifact + uses: actions/upload-pages-artifact@v3 + with: + path: docs/book + + deploy: + if: ${{ github.event_name != 'pull_request' }} + needs: build + runs-on: ubuntu-latest + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + steps: + - name: Deploy + id: deployment + uses: actions/deploy-pages@v4 diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml new file mode 100644 index 0000000..ac43c8e --- /dev/null +++ b/.github/workflows/pre-commit.yml @@ -0,0 +1,32 @@ +name: pre-commit + +on: + push: + branches: + - main + pull_request: + workflow_dispatch: + +concurrency: + group: "${{ github.workflow }} @ ${{ github.event.pull_request.head.label || github.head_ref || github.ref }}" + cancel-in-progress: true + +env: + PYTHON_VERSION: "3.10" + SKIP: cargo-fmt,cargo-update,cargo-clippy,cargo-check,cargo-test,cargo-test-doc,cargo-doc,cargo-miri-test,cargo-miri-run,cargo-deny-check + +jobs: + pre-commit: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ env.PYTHON_VERSION }} + + ## Run pre-commit and try to apply fixes + - name: Run pre-commit + uses: pre-commit/action@v3.0.1 + - name: Apply fixes from pre-commit + uses: pre-commit-ci/lite-action@v1.1.0 + if: always() diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml new file mode 100644 index 0000000..aa42b81 --- /dev/null +++ b/.github/workflows/python.yml @@ -0,0 +1,100 @@ +name: Python + +on: + push: + branches: + - main + paths-ignore: + - "docs/**" + pull_request: + release: + types: [published] + workflow_dispatch: + +concurrency: + group: "${{ github.workflow }} @ ${{ github.event.pull_request.head.label || github.head_ref || github.ref }}" + cancel-in-progress: true + +env: + CARGO_TERM_COLOR: always + PYTHON_VERSION: "3.10" + +jobs: + pytest: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: Swatinem/rust-cache@v2 + with: + save-if: ${{ github.event_name == 'push'}} + - uses: actions/setup-python@v5 + with: + python-version: ${{ env.PYTHON_VERSION }} + + - name: pip install + run: pip install .[tests] + - name: pytest + run: pytest + + wheels_sdist: + if: ${{ github.event_name == 'release' }} + needs: pytest + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Build sdist + uses: PyO3/maturin-action@v1 + with: + command: sdist + args: --out dist + - name: Upload sdist + uses: actions/upload-artifact@v4 + with: + name: wheels + path: dist + + wheels_linux: + if: ${{ github.event_name == 'release' }} + needs: pytest + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ env.PYTHON_VERSION }} + + - name: Build wheels + uses: PyO3/maturin-action@v1 + with: + args: --release --out dist --find-interpreter + sccache: "true" + manylinux: auto + - name: Upload wheels + uses: actions/upload-artifact@v4 + with: + name: wheels + path: dist + + pypi_release: + permissions: + contents: read + if: ${{ github.event_name == 'release' }} + needs: + - wheels_sdist + - wheels_linux + runs-on: ubuntu-latest + env: + MATURIN_PYPI_TOKEN: ${{ secrets.PYPI_TOKEN }} + steps: + - name: Download wheels + if: ${{ env.MATURIN_PYPI_TOKEN != '' }} + uses: actions/download-artifact@v4 + with: + name: wheels + - name: Publish to PyPI + if: ${{ env.MATURIN_PYPI_TOKEN != '' }} + uses: PyO3/maturin-action@v1 + with: + command: upload + args: --skip-existing * diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml new file mode 100644 index 0000000..579af8d --- /dev/null +++ b/.github/workflows/rust.yml @@ -0,0 +1,123 @@ +name: Rust + +on: + push: + branches: + - main + paths-ignore: + - "docs/**" + pull_request: + release: + types: [published] + workflow_dispatch: + +concurrency: + group: "${{ github.workflow }} @ ${{ github.event.pull_request.head.label || github.head_ref || github.ref }}" + cancel-in-progress: true + +env: + MSRV: "1.80" + CARGO_TERM_COLOR: always + PYTHON_PACKAGE_NAME: space_robotics_bench_py + GUI_PACKAGE_NAME: space_robotics_bench_gui + +jobs: + rustfmt: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + with: + components: rustfmt + + ## cargo fmt + - name: cargo fmt + run: cargo fmt --all --check --verbose + + cargo: + needs: rustfmt + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + toolchain: + - MSRV + - stable + - beta + steps: + - uses: actions/checkout@v4 + - uses: Swatinem/rust-cache@v2 + with: + save-if: ${{ github.event_name == 'push'}} + - uses: dtolnay/rust-toolchain@master + if: ${{ matrix.toolchain != 'MSRV' && matrix.toolchain != 'stable' }} + with: + toolchain: ${{ matrix.toolchain }} + - uses: dtolnay/rust-toolchain@master + if: ${{ matrix.toolchain == 'MSRV' }} + with: + toolchain: ${{ env.MSRV }} + - uses: dtolnay/rust-toolchain@master + if: ${{ matrix.toolchain == 'stable' }} + with: + toolchain: ${{ matrix.toolchain }} + components: clippy + + ## cargo check + - name: cargo check + run: cargo check --workspace --exclude ${{ env.PYTHON_PACKAGE_NAME }} --exclude ${{ env.GUI_PACKAGE_NAME }} --all-targets --verbose + - name: cargo check --no-default-features + run: cargo check --workspace --exclude ${{ env.PYTHON_PACKAGE_NAME }} --exclude ${{ env.GUI_PACKAGE_NAME }} --all-targets --no-default-features --verbose + - name: cargo check --all-features + run: cargo check --workspace --exclude ${{ env.PYTHON_PACKAGE_NAME }} --exclude ${{ env.GUI_PACKAGE_NAME }} --all-targets --all-features --verbose + + ## cargo test + - name: cargo test + run: cargo test --workspace --exclude ${{ env.PYTHON_PACKAGE_NAME }} --exclude ${{ env.GUI_PACKAGE_NAME }} --all-targets --verbose + - name: cargo test --no-default-features + run: cargo test --workspace --exclude ${{ env.PYTHON_PACKAGE_NAME }} --exclude ${{ env.GUI_PACKAGE_NAME }} --all-targets --no-default-features --verbose + - name: cargo test --all-features + run: cargo test --workspace --exclude ${{ env.PYTHON_PACKAGE_NAME }} --exclude ${{ env.GUI_PACKAGE_NAME }} --all-targets --all-features --verbose + + ## cargo test --doc + - name: cargo test --doc + run: cargo test --workspace --exclude ${{ env.PYTHON_PACKAGE_NAME }} --exclude ${{ env.GUI_PACKAGE_NAME }} --doc --verbose + - name: cargo test --doc --no-default-features + run: cargo test --workspace --exclude ${{ env.PYTHON_PACKAGE_NAME }} --exclude ${{ env.GUI_PACKAGE_NAME }} --doc --no-default-features --verbose + - name: cargo test --doc --all-features + run: cargo test --workspace --exclude ${{ env.PYTHON_PACKAGE_NAME }} --exclude ${{ env.GUI_PACKAGE_NAME }} --doc --all-features --verbose + + ## [stable] cargo clippy + - name: stable | cargo clippy + if: ${{ matrix.toolchain == 'stable' }} + run: cargo clippy --workspace --exclude ${{ env.PYTHON_PACKAGE_NAME }} --exclude ${{ env.GUI_PACKAGE_NAME }} --all-targets --all-features --no-deps --verbose -- --deny warnings + + ## [stable] cargo doc + - name: stable | cargo doc --document-private-items + if: ${{ matrix.toolchain == 'stable' }} + run: cargo doc --workspace --exclude ${{ env.PYTHON_PACKAGE_NAME }} --exclude ${{ env.GUI_PACKAGE_NAME }} --all-features --no-deps --document-private-items --verbose + + ## [stable] Code coverage + - name: stable | Install cargo llvm-cov for code coverage + uses: taiki-e/install-action@cargo-llvm-cov + if: ${{ matrix.toolchain == 'stable' }} + ## [stable] Generate coverage with cargo llvm-cov + - name: stable | Generate coverage + if: ${{ matrix.toolchain == 'stable' }} + run: cargo llvm-cov --workspace --exclude ${{ env.PYTHON_PACKAGE_NAME }} --exclude ${{ env.GUI_PACKAGE_NAME }} --all-features --lcov --output-path lcov.info + ## [stable] Upload coverage to codecov.io + - name: stable | Upload coverage + if: ${{ matrix.toolchain == 'stable' }} + uses: codecov/codecov-action@v4 + with: + token: ${{ secrets.CODECOV_TOKEN }} + files: lcov.info + fail_ci_if_error: false + + deny: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: EmbarkStudios/cargo-deny-action@v1 + with: + command: check bans licenses sources diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..fdf9f8a --- /dev/null +++ b/.gitignore @@ -0,0 +1,27 @@ +## Python +**/__pycache__/ +**/.pytest_cache/ +**/*.pcd +**/*.py[cod] +**/*.pyc +**/*$py.class + +## Rust +target/ +**/*.rs.bk +**/*.pdb +**/.cargo/config.toml +!.cargo/config.toml + +## Shared Objects +**/*.so + +## Dump files +**/core.* + +## Logs +logs/ +*.log + +## Apptainer +**/*.sif diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..1a423f6 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,4 @@ +[submodule "assets/srb_assets"] + path = assets/srb_assets + url = https://github.com/AndrejOrsula/srb_assets + branch = main diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..38f359d --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,80 @@ +# See https://pre-commit.com for more information +# See https://pre-commit.com/hooks.html for more hooks + +exclude: Cargo.lock|.*/patches/.*|.*hdr|.*svg|.*png|.*jpg +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.6.0 + hooks: + - id: check-added-large-files + args: ["--maxkb=1024"] + - id: check-ast + - id: check-builtin-literals + - id: check-case-conflict + - id: check-docstring-first + - id: check-executables-have-shebangs + - id: check-merge-conflict + - id: check-shebang-scripts-are-executable + exclude: \.rs$ + - id: check-symlinks + - id: check-toml + - id: check-vcs-permalinks + - id: check-xml + - id: check-yaml + - id: debug-statements + - id: destroyed-symlinks + - id: detect-private-key + - id: end-of-file-fixer + - id: fix-byte-order-marker + - id: mixed-line-ending + - id: name-tests-test + - id: trailing-whitespace + + - repo: https://github.com/lovesegfault/beautysh + rev: v6.2.1 + hooks: + - id: beautysh + + - repo: https://github.com/psf/black + rev: 24.10.0 + hooks: + - id: black + + - repo: https://github.com/codespell-project/codespell + rev: v2.3.0 + hooks: + - id: codespell + args: ["--ignore-words-list", "crate"] + + - repo: https://github.com/hadolint/hadolint + rev: v2.13.1-beta + hooks: + - id: hadolint-docker + + - repo: https://github.com/pycqa/isort + rev: 5.13.2 + hooks: + - id: isort + args: ["--profile", "black"] + + - repo: https://github.com/executablebooks/mdformat + rev: 0.7.18 + hooks: + - id: mdformat + + - repo: https://github.com/AndrejOrsula/pre-commit-cargo + rev: 0.3.0 + hooks: + - id: cargo-fmt + - id: cargo-update + - id: cargo-clippy + args: ["--workspace", "--all-targets", "--", "--deny=warnings"] + - id: cargo-check + args: ["--workspace", "--all-targets"] + - id: cargo-test + args: ["--workspace", "--all-targets"] + - id: cargo-test-doc + args: ["--workspace"] + - id: cargo-doc + args: ["--workspace", "--no-deps", "--document-private-items"] + - id: cargo-deny-check diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..5561ba2 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,582 @@ +{ + "files.associations": { + "*.kit": "toml" + }, + "python.defaultInterpreterPath": "~/isaac-sim/kit/python/bin/python3", + "python.analysis.extraPaths": [ + "scripts", + "${workspaceFolder}/../isaaclab/source/extensions/omni.isaac.lab", + "${workspaceFolder}/../isaaclab/source/extensions/omni.isaac.lab_assets", + "${workspaceFolder}/../isaaclab/source/extensions/omni.isaac.lab_tasks", + "~/isaac-sim/exts/omni.exporter.urdf", + "~/isaac-sim/exts/omni.exporter.urdf/pip_prebundle", + "~/isaac-sim/exts/omni.isaac.app.selector", + "~/isaac-sim/exts/omni.isaac.app.setup", + "~/isaac-sim/exts/omni.isaac.articulation_inspector", + "~/isaac-sim/exts/omni.isaac.asset_browser", + "~/isaac-sim/exts/omni.isaac.assets_check", + "~/isaac-sim/exts/omni.isaac.benchmark.services", + "~/isaac-sim/exts/omni.isaac.benchmarks", + "~/isaac-sim/exts/omni.isaac.block_world", + "~/isaac-sim/exts/omni.isaac.camera_inspector", + "~/isaac-sim/exts/omni.isaac.cloner", + "~/isaac-sim/exts/omni.isaac.common_includes", + "~/isaac-sim/exts/omni.isaac.conveyor", + "~/isaac-sim/exts/omni.isaac.conveyor.ui", + "~/isaac-sim/exts/omni.isaac.core", + "~/isaac-sim/exts/omni.isaac.core_archive", + "~/isaac-sim/exts/omni.isaac.core_archive/pip_prebundle", + "~/isaac-sim/exts/omni.isaac.core_nodes", + "~/isaac-sim/exts/omni.isaac.cortex", + "~/isaac-sim/exts/omni.isaac.cortex.sample_behaviors", + "~/isaac-sim/exts/omni.isaac.cortex_sync", + "~/isaac-sim/exts/omni.isaac.debug_draw", + "~/isaac-sim/exts/omni.isaac.doctest", + "~/isaac-sim/exts/omni.isaac.dynamic_control", + "~/isaac-sim/exts/omni.isaac.examples", + "~/isaac-sim/exts/omni.isaac.examples_nodes", + "~/isaac-sim/exts/omni.isaac.extension_templates", + "~/isaac-sim/exts/omni.isaac.franka", + "~/isaac-sim/exts/omni.isaac.gain_tuner", + "~/isaac-sim/exts/omni.isaac.grasp_editor", + "~/isaac-sim/exts/omni.isaac.import_wizard", + "~/isaac-sim/exts/omni.isaac.internal_tools", + "~/isaac-sim/exts/omni.isaac.jupyter_notebook", + "~/isaac-sim/exts/omni.isaac.kit", + "~/isaac-sim/exts/omni.isaac.lula", + "~/isaac-sim/exts/omni.isaac.lula/pip_prebundle", + "~/isaac-sim/exts/omni.isaac.lula_test_widget", + "~/isaac-sim/exts/omni.isaac.manipulators", + "~/isaac-sim/exts/omni.isaac.manipulators.ui", + "~/isaac-sim/exts/omni.isaac.menu", + "~/isaac-sim/exts/omni.isaac.merge_mesh", + "~/isaac-sim/exts/omni.isaac.ml_archive", + "~/isaac-sim/exts/omni.isaac.ml_archive/pip_prebundle", + "~/isaac-sim/exts/omni.isaac.motion_generation", + "~/isaac-sim/exts/omni.isaac.nucleus", + "~/isaac-sim/exts/omni.isaac.occupancy_map", + "~/isaac-sim/exts/omni.isaac.occupancy_map.ui", + "~/isaac-sim/exts/omni.isaac.ocs2", + "~/isaac-sim/exts/omni.isaac.physics_inspector", + "~/isaac-sim/exts/omni.isaac.physics_utilities", + "~/isaac-sim/exts/omni.isaac.proximity_sensor", + "~/isaac-sim/exts/omni.isaac.quadruped", + "~/isaac-sim/exts/omni.isaac.range_sensor", + "~/isaac-sim/exts/omni.isaac.range_sensor.examples", + "~/isaac-sim/exts/omni.isaac.range_sensor.ui", + "~/isaac-sim/exts/omni.isaac.repl", + "~/isaac-sim/exts/omni.isaac.robot_assembler", + "~/isaac-sim/exts/omni.isaac.robot_description_editor", + "~/isaac-sim/exts/omni.isaac.ros2_bridge", + "~/isaac-sim/exts/omni.isaac.ros2_bridge.robot_description", + "~/isaac-sim/exts/omni.isaac.ros_bridge", + "~/isaac-sim/exts/omni.isaac.scene_blox", + "~/isaac-sim/exts/omni.isaac.sensor", + "~/isaac-sim/exts/omni.isaac.surface_gripper", + "~/isaac-sim/exts/omni.isaac.surface_gripper.ui", + "~/isaac-sim/exts/omni.isaac.synthetic_recorder", + "~/isaac-sim/exts/omni.isaac.tests", + "~/isaac-sim/exts/omni.isaac.tf_viewer", + "~/isaac-sim/exts/omni.isaac.throttling", + "~/isaac-sim/exts/omni.isaac.ui", + "~/isaac-sim/exts/omni.isaac.ui_template", + "~/isaac-sim/exts/omni.isaac.universal_robots", + "~/isaac-sim/exts/omni.isaac.utils", + "~/isaac-sim/exts/omni.isaac.version", + "~/isaac-sim/exts/omni.isaac.vscode", + "~/isaac-sim/exts/omni.isaac.wheeled_robots", + "~/isaac-sim/exts/omni.isaac.wheeled_robots.ui", + "~/isaac-sim/exts/omni.isaac.window.about", + "~/isaac-sim/exts/omni.kit.loop-isaac", + "~/isaac-sim/exts/omni.kit.property.isaac", + "~/isaac-sim/exts/omni.pip.cloud", + "~/isaac-sim/exts/omni.pip.cloud/pip_prebundle", + "~/isaac-sim/exts/omni.pip.compute", + "~/isaac-sim/exts/omni.pip.compute/pip_prebundle", + "~/isaac-sim/exts/omni.replicator.isaac", + "~/isaac-sim/exts/omni.usd.schema.isaac", + "~/isaac-sim/extsPhysics/omni.convexdecomposition", + "~/isaac-sim/extsPhysics/omni.kit.property.physx", + "~/isaac-sim/extsPhysics/omni.kvdb", + "~/isaac-sim/extsPhysics/omni.localcache", + "~/isaac-sim/extsPhysics/omni.physics.tensors", + "~/isaac-sim/extsPhysics/omni.physics.tensors.tests", + "~/isaac-sim/extsPhysics/omni.physx", + "~/isaac-sim/extsPhysics/omni.physx.bundle", + "~/isaac-sim/extsPhysics/omni.physx.camera", + "~/isaac-sim/extsPhysics/omni.physx.cct", + "~/isaac-sim/extsPhysics/omni.physx.commands", + "~/isaac-sim/extsPhysics/omni.physx.cooking", + "~/isaac-sim/extsPhysics/omni.physx.demos", + "~/isaac-sim/extsPhysics/omni.physx.fabric", + "~/isaac-sim/extsPhysics/omni.physx.forcefields", + "~/isaac-sim/extsPhysics/omni.physx.foundation", + "~/isaac-sim/extsPhysics/omni.physx.graph", + "~/isaac-sim/extsPhysics/omni.physx.internal", + "~/isaac-sim/extsPhysics/omni.physx.pvd", + "~/isaac-sim/extsPhysics/omni.physx.stageupdate", + "~/isaac-sim/extsPhysics/omni.physx.supportui", + "~/isaac-sim/extsPhysics/omni.physx.telemetry", + "~/isaac-sim/extsPhysics/omni.physx.tensors", + "~/isaac-sim/extsPhysics/omni.physx.tests", + "~/isaac-sim/extsPhysics/omni.physx.tests.mini", + "~/isaac-sim/extsPhysics/omni.physx.tests.visual", + "~/isaac-sim/extsPhysics/omni.physx.ui", + "~/isaac-sim/extsPhysics/omni.physx.vehicle", + "~/isaac-sim/extsPhysics/omni.physx.vehicle.tests", + "~/isaac-sim/extsPhysics/omni.physx.zerogravity", + "~/isaac-sim/extsPhysics/omni.usd.schema.forcefield", + "~/isaac-sim/extsPhysics/omni.usd.schema.physx", + "~/isaac-sim/extsPhysics/omni.usdphysics", + "~/isaac-sim/extsPhysics/omni.usdphysics.tests", + "~/isaac-sim/extsPhysics/omni.usdphysics.ui", + "~/isaac-sim/extscache/carb.audio-0.1.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/carb.imaging.python-0.1.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/carb.windowing.plugins-1.0.0+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.activity.core-1.0.1+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.anim.asset-106.1.0+106.1.0.lx64.r", + "~/isaac-sim/extscache/omni.anim.behavior.schema-106.1.0+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.anim.curve.bundle-1.2.3+106.1.0.ub3f", + "~/isaac-sim/extscache/omni.anim.curve.core-1.1.14+106.1.0.lx64.r.cp310.ub3f", + "~/isaac-sim/extscache/omni.anim.curve.ui-1.3.17+106.1.0.ub3f", + "~/isaac-sim/extscache/omni.anim.curve_editor-105.17.10+106.0.1.ub3f", + "~/isaac-sim/extscache/omni.anim.graph.bundle-106.1.0+106.1.0", + "~/isaac-sim/extscache/omni.anim.graph.core-106.1.2+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.anim.graph.schema-106.1.0+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.anim.graph.ui-106.1.1+106.1.0", + "~/isaac-sim/extscache/omni.anim.navigation.bundle-106.1.0+106.1.0", + "~/isaac-sim/extscache/omni.anim.navigation.core-106.1.3+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.anim.navigation.navmesh.recast-106.1.4+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.anim.navigation.schema-106.1.0+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.anim.navigation.ui-106.1.1+106.1.0", + "~/isaac-sim/extscache/omni.anim.people-0.5.0", + "~/isaac-sim/extscache/omni.anim.retarget.bundle-106.1.0+106.1.0", + "~/isaac-sim/extscache/omni.anim.retarget.core-106.1.2+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.anim.retarget.ui-106.1.1+106.1.0", + "~/isaac-sim/extscache/omni.anim.shared.core-106.0.1+106.0.1.lx64.r.cp310", + "~/isaac-sim/extscache/omni.anim.skelJoint-106.1.2+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.anim.timeline-105.0.23+106.1.0.lx64.r.cp310.ub3f", + "~/isaac-sim/extscache/omni.anim.window.timeline-105.13.5+106.0.1.ub3f", + "~/isaac-sim/extscache/omni.appwindow-1.1.8+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.asset_validator.core-0.11.8", + "~/isaac-sim/extscache/omni.asset_validator.ui-0.11.8", + "~/isaac-sim/extscache/omni.blobkey-1.1.2+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.command.usd-1.0.3+10a4b5c0", + "~/isaac-sim/extscache/omni.cuopt.examples-1.0.0+106.0.0", + "~/isaac-sim/extscache/omni.cuopt.service-1.0.0+106.0.0", + "~/isaac-sim/extscache/omni.cuopt.visualization-1.0.0+106.0.0", + "~/isaac-sim/extscache/omni.curve.creator-105.0.4+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.curve.manipulator-105.2.8+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.datastore-0.0.0+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.debugdraw-0.1.3+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.extended.materials-105.0.9", + "~/isaac-sim/extscache/omni.fabric.commands-1.1.5+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.flowusd-106.1.1+106.1.0.lx64.r.cp310.ub3f", + "~/isaac-sim/extscache/omni.genproc.core-106.1.0+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.gpu_foundation-0.0.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.gpu_foundation.shadercache.vulkan-1.0.0+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.gpucompute.plugins-0.0.0+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.graph-1.140.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.graph.action-1.102.1+106.0.3", + "~/isaac-sim/extscache/omni.graph.action_core-1.1.6+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.graph.action_nodes-1.24.0+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.graph.bundle.action-2.4.1+106.0.3", + "~/isaac-sim/extscache/omni.graph.core-2.179.2+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.graph.exec-0.9.4+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.graph.image.core-0.4.5+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.graph.image.nodes-1.1.0+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.graph.io-1.9.1+106.0.3", + "~/isaac-sim/extscache/omni.graph.nodes-1.146.1+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.graph.scriptnode-1.20.1+106.1.0", + "~/isaac-sim/extscache/omni.graph.telemetry-2.15.0+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.graph.tools-1.79.0+10a4b5c0", + "~/isaac-sim/extscache/omni.graph.tutorials-1.29.0+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.graph.ui-1.70.2+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.graph.ui_nodes-1.26.0+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.graph.visualization.nodes-2.1.1", + "~/isaac-sim/extscache/omni.graph.window.action-1.28.0+106.1.0", + "~/isaac-sim/extscache/omni.graph.window.core-1.113.1+106.1.0", + "~/isaac-sim/extscache/omni.graph.window.generic-1.26.0+106.1.0", + "~/isaac-sim/extscache/omni.hsscclient-1.1.1+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.hydra.engine.stats-1.0.2+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.hydra.iray.shadercache.vulkan-1.0.0+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.hydra.rtx-1.0.0+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.hydra.rtx.shadercache.vulkan-1.0.0+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.hydra.scene_api-0.1.2+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.hydra.scene_delegate-0.3.3+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.hydra.usdrt_delegate-7.5.1+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.importer.mjcf-1.1.1+106.0.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.importer.onshape-0.7.3+106.0.0", + "~/isaac-sim/extscache/omni.importer.urdf-1.14.1+106.0.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.inspect-1.0.1+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.iray.libs-0.0.0+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.kit.actions.core-1.0.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.actions.window-1.1.1+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.agent.watcher-0.2.1", + "~/isaac-sim/extscache/omni.kit.asset_converter-2.1.21+lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.audiodeviceenum-1.0.1+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.browser.asset-1.3.10", + "~/isaac-sim/extscache/omni.kit.browser.asset_provider.local-1.0.9", + "~/isaac-sim/extscache/omni.kit.browser.core-2.3.11", + "~/isaac-sim/extscache/omni.kit.browser.deepsearch-1.1.8", + "~/isaac-sim/extscache/omni.kit.browser.folder.core-1.9.13", + "~/isaac-sim/extscache/omni.kit.browser.material-1.6.0", + "~/isaac-sim/extscache/omni.kit.browser.sample-1.4.8", + "~/isaac-sim/extscache/omni.kit.browser.texture-1.2.1", + "~/isaac-sim/extscache/omni.kit.capture.viewport-1.5.1", + "~/isaac-sim/extscache/omni.kit.clipboard-1.0.4+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.collaboration.channel_manager-1.0.12+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.collaboration.presence_layer-1.0.9+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.collaboration.telemetry-1.0.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.commands-1.4.9+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.context_menu-1.8.1+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.converter.cad-201.1.0+106.0.1", + "~/isaac-sim/extscache/omni.kit.converter.common-500.0.10+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.converter.dgn-500.0.6+106.1.0", + "~/isaac-sim/extscache/omni.kit.converter.dgn_core-500.0.19+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.converter.geojson-0.0.10+lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.converter.hoops-500.0.7+106.1.0", + "~/isaac-sim/extscache/omni.kit.converter.hoops_core-500.0.17+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.converter.jt-500.0.7+106.1.0", + "~/isaac-sim/extscache/omni.kit.converter.jt_core-500.0.16+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.converter.lib3mf-1.1.3", + "~/isaac-sim/extscache/omni.kit.converter.ogc-1.1.22+lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.converter.stl-0.1.1+105.1", + "~/isaac-sim/extscache/omni.kit.converter.vtk-2.3.1+105.2.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.core.collection-0.1.7", + "~/isaac-sim/extscache/omni.kit.data2ui.core-1.0.27+106.0", + "~/isaac-sim/extscache/omni.kit.data2ui.usd-1.0.27+106.0", + "~/isaac-sim/extscache/omni.kit.environment.core-1.3.14", + "~/isaac-sim/extscache/omni.kit.exec.core-0.13.4+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.gfn-106.0.5+106.0.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.graph.delegate.default-1.2.2", + "~/isaac-sim/extscache/omni.kit.graph.delegate.modern-1.10.6", + "~/isaac-sim/extscache/omni.kit.graph.delegate.neo-1.1.3", + "~/isaac-sim/extscache/omni.kit.graph.editor.core-1.5.3", + "~/isaac-sim/extscache/omni.kit.graph.editor.example-1.0.24", + "~/isaac-sim/extscache/omni.kit.graph.usd.commands-1.3.1", + "~/isaac-sim/extscache/omni.kit.graph.widget.variables-2.1.0", + "~/isaac-sim/extscache/omni.kit.helper.file_utils-0.1.8+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.hotkeys.core-1.3.5+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.hotkeys.window-1.4.5+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.hydra_texture-1.3.9+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.livestream.core-3.2.0+105.2.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.livestream.core-4.3.6+106.0.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.livestream.messaging-1.1.1", + "~/isaac-sim/extscache/omni.kit.livestream.native-4.1.0+106.0.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.livestream.webrtc-4.1.1+106.0.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.mainwindow-1.0.3+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.manipulator.camera-105.0.5+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.manipulator.prim-107.0.0+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.manipulator.prim.core-107.0.4+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.manipulator.prim.fabric-107.0.3+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.manipulator.prim.usd-107.0.2+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.manipulator.selection-106.0.1+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.manipulator.selector-1.1.1+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.manipulator.tool.mesh_snap-1.4.5+106.0.0", + "~/isaac-sim/extscache/omni.kit.manipulator.tool.snap-1.5.11+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.manipulator.transform-104.7.4+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.manipulator.viewport-107.0.0+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.material.library-1.5.6+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.menu.aov-1.1.4+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.menu.common-1.1.7+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.menu.core-1.0.4+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.menu.create-1.0.16+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.menu.edit-1.1.24+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.menu.file-1.1.14+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.menu.stage-1.2.5", + "~/isaac-sim/extscache/omni.kit.menu.utils-1.5.27+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.mesh.raycast-105.4.0+106.0.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.ngsearch-0.3.3", + "~/isaac-sim/extscache/omni.kit.notification_manager-1.0.9+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.numpy.common-0.1.2+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.pip_archive-0.0.0+10a4b5c0.lx64.cp310", + "~/isaac-sim/extscache/omni.kit.pip_archive-0.0.0+10a4b5c0.lx64.cp310/pip_prebundle", + "~/isaac-sim/extscache/omni.kit.pipapi-0.0.0+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.playlist.core-1.3.4", + "~/isaac-sim/extscache/omni.kit.pointclouds-1.4.3+cp310", + "~/isaac-sim/extscache/omni.kit.preferences.animation-1.1.8+106.1.0.ub3f", + "~/isaac-sim/extscache/omni.kit.prim.icon-1.0.13", + "~/isaac-sim/extscache/omni.kit.primitive.mesh-1.0.17+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.profiler.tracy-1.1.5+106.0.0.lx64", + "~/isaac-sim/extscache/omni.kit.profiler.window-2.2.3", + "~/isaac-sim/extscache/omni.kit.property.adapter.core-1.0.1+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.property.adapter.fabric-1.0.1+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.property.adapter.usd-1.0.1+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.property.audio-1.0.14+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.property.bundle-1.3.1+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.property.camera-1.0.8+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.property.collection-0.1.17", + "~/isaac-sim/extscache/omni.kit.property.environment-1.2.1", + "~/isaac-sim/extscache/omni.kit.property.geometry-1.3.1+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.property.layer-1.1.8+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.property.light-1.0.10+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.property.material-1.10.8+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.property.render-1.1.2+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.property.sbsar-107.0.0", + "~/isaac-sim/extscache/omni.kit.property.transform-1.5.9+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.property.usd-4.2.8+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.property.visualization-104.0.5", + "~/isaac-sim/extscache/omni.kit.quicklayout-1.0.7+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.raycast.query-1.0.5+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.renderer.capture-0.0.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.renderer.core-1.0.2+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.renderer.cuda_interop-1.0.1+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.renderer.imgui-1.0.1+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.renderer.init-0.0.0+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.kit.scripting-106.1.1+106.1.0", + "~/isaac-sim/extscache/omni.kit.search.files-1.0.4", + "~/isaac-sim/extscache/omni.kit.search.service-0.1.12", + "~/isaac-sim/extscache/omni.kit.search_core-1.0.5+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.selection-0.1.4+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.sequencer.core-103.4.2+106.0.0", + "~/isaac-sim/extscache/omni.kit.sequencer.usd-103.4.5+106.0.0", + "~/isaac-sim/extscache/omni.kit.stage.copypaste-1.2.2+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.stage.mdl_converter-1.0.6+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.stage_column.payload-2.0.0", + "~/isaac-sim/extscache/omni.kit.stage_column.variant-1.0.13", + "~/isaac-sim/extscache/omni.kit.stage_template.core-1.1.22+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.stage_templates-1.2.5+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.stagerecorder.bundle-105.0.2+106.1.0", + "~/isaac-sim/extscache/omni.kit.stagerecorder.core-105.0.5+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.stagerecorder.ui-105.0.6+106.1.0", + "~/isaac-sim/extscache/omni.kit.streamsdk.plugins-3.2.1+105.2.lx64.r", + "~/isaac-sim/extscache/omni.kit.streamsdk.plugins-4.5.2+106.0.0.lx64.r", + "~/isaac-sim/extscache/omni.kit.streamsdk.plugins-4.5.3+106.0.0.lx64.r", + "~/isaac-sim/extscache/omni.kit.telemetry-0.5.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.test-1.1.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.test_helpers_gfx-0.0.0+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.thumbnails.mdl-1.0.24", + "~/isaac-sim/extscache/omni.kit.thumbnails.usd-1.0.9", + "~/isaac-sim/extscache/omni.kit.timeline.minibar-1.2.9", + "~/isaac-sim/extscache/omni.kit.tool.asset_exporter-1.3.3", + "~/isaac-sim/extscache/omni.kit.tool.asset_importer-2.5.5", + "~/isaac-sim/extscache/omni.kit.tool.collect-2.2.14", + "~/isaac-sim/extscache/omni.kit.tool.measure-105.2.6+106.0.1", + "~/isaac-sim/extscache/omni.kit.tool.remove_unused.controller-0.1.3", + "~/isaac-sim/extscache/omni.kit.tool.remove_unused.core-0.1.2", + "~/isaac-sim/extscache/omni.kit.ui.actions-1.0.2+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.ui_test-1.3.2+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.uiapp-0.0.0+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.usd.collect-2.2.21+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.usd.layers-2.1.36+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.usd.mdl-1.0.1+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.usd_undo-0.1.8+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.usda_edit-1.1.14+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.variant.editor-106.1.0+106.1.0", + "~/isaac-sim/extscache/omni.kit.variant.presenter-106.0.0", + "~/isaac-sim/extscache/omni.kit.viewport.actions-106.0.2+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.viewport.bundle-104.0.1+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.viewport.legacy_gizmos-1.1.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.viewport.manipulator.transform-107.0.2+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.viewport.menubar.camera-107.0.2+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.viewport.menubar.core-106.1.0+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.viewport.menubar.display-107.0.2+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.viewport.menubar.lighting-106.0.2+ub3f", + "~/isaac-sim/extscache/omni.kit.viewport.menubar.render-106.1.1+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.viewport.menubar.settings-107.0.3+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.viewport.menubar.waypoint-104.2.16", + "~/isaac-sim/extscache/omni.kit.viewport.registry-104.0.6+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.viewport.rtx-104.0.1+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.viewport.utility-1.0.17+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.viewport.window-107.0.6+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.viewport_widgets_manager-1.0.8+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.waypoint.bundle-1.0.4", + "~/isaac-sim/extscache/omni.kit.waypoint.core-1.4.54", + "~/isaac-sim/extscache/omni.kit.waypoint.playlist-1.0.8", + "~/isaac-sim/extscache/omni.kit.widget.browser_bar-2.0.10+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.cache_indicator-2.0.10+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.calendar-1.0.8", + "~/isaac-sim/extscache/omni.kit.widget.collection-0.1.18", + "~/isaac-sim/extscache/omni.kit.widget.context_menu-1.2.2+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.extended_searchfield-1.0.28", + "~/isaac-sim/extscache/omni.kit.widget.filebrowser-2.10.51+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.filter-1.1.4+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.graph-1.12.15+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.highlight_label-1.0.2+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.imageview-1.0.3+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.inspector-1.0.3+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.widget.layers-1.8.0+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.live-2.1.8+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.live_session_management-1.2.20+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.live_session_management.ui-1.0.1+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.material_preview-1.0.16", + "~/isaac-sim/extscache/omni.kit.widget.nucleus_connector-1.1.9+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.nucleus_info-1.0.2+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.opengl-1.0.8+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.widget.options_button-1.0.3+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.options_menu-1.1.6+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.path_field-2.0.10+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.prompt-1.0.7+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.search_delegate-1.0.5+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.searchable_combobox-1.0.6+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.searchfield-1.1.8+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.settings-1.2.2+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.sliderbar-1.0.10", + "~/isaac-sim/extscache/omni.kit.widget.stage-2.11.2+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.stage_icons-1.0.5+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.text_editor-1.0.2+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.widget.timeline-105.0.1+105.0", + "~/isaac-sim/extscache/omni.kit.widget.toolbar-1.7.2+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.versioning-1.4.7+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.viewport-106.1.0+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.widget.zoombar-1.0.5", + "~/isaac-sim/extscache/omni.kit.widgets.custom-1.0.8", + "~/isaac-sim/extscache/omni.kit.window.collection-0.1.22", + "~/isaac-sim/extscache/omni.kit.window.commands-0.2.6+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.window.console-0.2.13+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.window.content_browser-2.9.18+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.window.content_browser_registry-0.0.6+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.window.cursor-1.1.2+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.window.drop_support-1.0.3+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.window.extensions-1.4.11+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.window.file-1.3.54+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.window.file_exporter-1.0.30+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.window.file_importer-1.1.12+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.window.filepicker-2.10.40+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.window.material-1.6.2", + "~/isaac-sim/extscache/omni.kit.window.material_graph-1.8.18", + "~/isaac-sim/extscache/omni.kit.window.movie_capture-2.4.2", + "~/isaac-sim/extscache/omni.kit.window.popup_dialog-2.0.24+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.window.preferences-1.6.0+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.window.property-1.11.3+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.window.quicksearch-2.4.4", + "~/isaac-sim/extscache/omni.kit.window.script_editor-1.7.6+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.window.section-107.0.1", + "~/isaac-sim/extscache/omni.kit.window.stage-2.5.10+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.window.stats-0.1.6+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.window.status_bar-0.1.7+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.window.title-1.1.5+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.window.toolbar-1.6.1+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.window.usd_paths-1.0.7+10a4b5c0", + "~/isaac-sim/extscache/omni.kit.window.usddebug-1.0.2", + "~/isaac-sim/extscache/omni.kit.xr.advertise-106.1.24+106.1.0", + "~/isaac-sim/extscache/omni.kit.xr.core-106.1.24+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.xr.example.usd_scene_ui-106.0.116+106.0.1", + "~/isaac-sim/extscache/omni.kit.xr.profile.ar-106.1.24+106.1.0", + "~/isaac-sim/extscache/omni.kit.xr.profile.common-106.1.24+106.1.0", + "~/isaac-sim/extscache/omni.kit.xr.profile.tabletar-106.1.24+106.1.0", + "~/isaac-sim/extscache/omni.kit.xr.profile.vr-106.1.24+106.1.0", + "~/isaac-sim/extscache/omni.kit.xr.scene_view.core-106.1.24+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.xr.scene_view.utils-106.1.24+106.1.0", + "~/isaac-sim/extscache/omni.kit.xr.system.cloudxr-106.1.24+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.xr.system.cloudxr41-106.1.24+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.xr.system.openxr-106.1.24+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.xr.system.playback-106.1.24+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.xr.system.steamvr-106.1.24+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.xr.telemetry-106.1.24+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.kit.xr.ui.config.common-106.1.24+106.1.0", + "~/isaac-sim/extscache/omni.kit.xr.ui.config.generic-106.1.24+106.1.0", + "~/isaac-sim/extscache/omni.kit.xr.ui.config.htcvive-106.1.24+106.1.0", + "~/isaac-sim/extscache/omni.kit.xr.ui.config.magicleap-106.1.24+106.1.0", + "~/isaac-sim/extscache/omni.kit.xr.ui.config.metaquest-106.1.24+106.1.0", + "~/isaac-sim/extscache/omni.kit.xr.ui.stage.common-106.1.24+106.1.0", + "~/isaac-sim/extscache/omni.kit.xr.ui.window.profile-106.1.24+106.1.0", + "~/isaac-sim/extscache/omni.kit.xr.ui.window.viewport-106.1.24+106.1.0", + "~/isaac-sim/extscache/omni.materialx.libs-1.0.4+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.mdl-55.0.1+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.mdl.neuraylib-0.2.8+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.mdl.usd_converter-1.0.21+10a4b5c0", + "~/isaac-sim/extscache/omni.net-0.0.1-isaac-1+lx64.r", + "~/isaac-sim/extscache/omni.no_code_ui.bundle-1.0.27+106.0", + "~/isaac-sim/extscache/omni.product_configurator.panel-1.0.15", + "~/isaac-sim/extscache/omni.product_configurator.utils-1.2.2", + "~/isaac-sim/extscache/omni.ramp-105.1.15+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.replicator.agent.camera_calibration-0.4.0", + "~/isaac-sim/extscache/omni.replicator.agent.core-0.4.1", + "~/isaac-sim/extscache/omni.replicator.agent.ui-0.4.0", + "~/isaac-sim/extscache/omni.replicator.core-1.11.20+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.replicator.object-0.3.8+lx64", + "~/isaac-sim/extscache/omni.replicator.replicator_yaml-2.0.6+lx64", + "~/isaac-sim/extscache/omni.resourcemonitor-105.0.1+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.rtx.settings.core-0.6.3+10a4b5c0", + "~/isaac-sim/extscache/omni.rtx.shadercache.vulkan-1.0.0+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.rtx.window.settings-0.6.17+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.scene.optimizer.bundle-106.1.1+106.1.0", + "~/isaac-sim/extscache/omni.scene.optimizer.core-106.1.1+106.1.0.lx64.r.cp310.ub3f", + "~/isaac-sim/extscache/omni.scene.optimizer.ui-106.1.1+106.1.0", + "~/isaac-sim/extscache/omni.scene.visualization.bundle-105.1.0+106.1.0", + "~/isaac-sim/extscache/omni.scene.visualization.core-105.4.13+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.scene.visualization.ui-105.1.2+106.1.0", + "~/isaac-sim/extscache/omni.sensors.nv.camera-0.10.0-isaac-1+lx64.r", + "~/isaac-sim/extscache/omni.sensors.nv.common-1.2.2-isaac-1+lx64.r.cp310", + "~/isaac-sim/extscache/omni.sensors.nv.ids-1.1.0-isaac-1+lx64.r.cp310", + "~/isaac-sim/extscache/omni.sensors.nv.lidar-1.2.2-isaac-1+lx64.r.cp310", + "~/isaac-sim/extscache/omni.sensors.nv.materials-1.2.1-isaac-1+lx64.r.cp310", + "~/isaac-sim/extscache/omni.sensors.nv.radar-1.2.1-isaac-1+lx64.r.cp310", + "~/isaac-sim/extscache/omni.sensors.nv.ultrasonic-1.2.1-isaac-1+lx64.r.cp310", + "~/isaac-sim/extscache/omni.sensors.nv.visualizer-1.0.1-isaac-1+lx64.r.cp310", + "~/isaac-sim/extscache/omni.sensors.nv.wpm-1.2.1-isaac-1+lx64.r", + "~/isaac-sim/extscache/omni.sensors.tiled-0.0.6+106.1.0.lx64.r", + "~/isaac-sim/extscache/omni.services.browser.asset-1.3.3+106.0.0", + "~/isaac-sim/extscache/omni.services.carb.event_stream-1.0.0", + "~/isaac-sim/extscache/omni.services.client-0.5.3", + "~/isaac-sim/extscache/omni.services.core-1.9.0", + "~/isaac-sim/extscache/omni.services.facilities.base-1.0.4", + "~/isaac-sim/extscache/omni.services.facilities.workqueue-1.1.2", + "~/isaac-sim/extscache/omni.services.pip_archive-0.13.6+lx64", + "~/isaac-sim/extscache/omni.services.starfleet.auth-0.1.5", + "~/isaac-sim/extscache/omni.services.streamclient.webrtc-1.3.8", + "~/isaac-sim/extscache/omni.services.streamclient.websocket-2.0.0", + "~/isaac-sim/extscache/omni.services.streaming.manager-0.3.10", + "~/isaac-sim/extscache/omni.services.thumbnails.images-1.3.2", + "~/isaac-sim/extscache/omni.services.transport.client.base-1.2.4", + "~/isaac-sim/extscache/omni.services.transport.client.http_async-1.3.6", + "~/isaac-sim/extscache/omni.services.transport.server.base-1.1.1", + "~/isaac-sim/extscache/omni.services.transport.server.http-1.3.1", + "~/isaac-sim/extscache/omni.services.transport.server.zeroconf-1.0.9", + "~/isaac-sim/extscache/omni.services.usd-1.1.0", + "~/isaac-sim/extscache/omni.simready.explorer-1.1.1", + "~/isaac-sim/extscache/omni.slangnode-106.1.0+106.1.0.lx64.r.cp310.ub3f", + "~/isaac-sim/extscache/omni.stats-1.0.1+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.syntheticdata-0.6.9+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.timeline-1.0.10+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.timeline.live_session-1.0.8+10a4b5c0", + "~/isaac-sim/extscache/omni.tools.array-105.0.4", + "~/isaac-sim/extscache/omni.ui-2.25.22+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.ui.scene-1.10.3+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.ui_query-1.1.4+10a4b5c0", + "~/isaac-sim/extscache/omni.uiaudio-1.0.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.ujitso.client-0.0.0+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.ujitso.default-1.0.0+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.ujitso.processor.texture-1.0.0+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.usd-1.12.2+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.usd.config-1.0.4+10a4b5c0", + "~/isaac-sim/extscache/omni.usd.core-1.4.1+10a4b5c0.lx64.r", + "~/isaac-sim/extscache/omni.usd.fileformat.e57-1.2.4+106.0.0.lx64.r.cp310.ub3f", + "~/isaac-sim/extscache/omni.usd.fileformat.pts-106.0.1+106.0.0.lx64.r.cp310.ub3f", + "~/isaac-sim/extscache/omni.usd.fileformat.sbsar-107.0.2+lx64.r.cp310.ub3f", + "~/isaac-sim/extscache/omni.usd.libs-1.0.1+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.usd.metrics.assembler-106.1.0+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.usd.metrics.assembler.physics-106.1.0+106.1.0", + "~/isaac-sim/extscache/omni.usd.metrics.assembler.ui-106.1.0+106.1.0", + "~/isaac-sim/extscache/omni.usd.schema.anim-0.0.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.usd.schema.audio-0.0.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.usd.schema.destruction-0.7.0+106.0.1.lx64.r", + "~/isaac-sim/extscache/omni.usd.schema.flow-106.0.8+106.0.0.ub3f", + "~/isaac-sim/extscache/omni.usd.schema.geospatial-0.0.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.usd.schema.metrics.assembler-106.1.0+106.1.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.usd.schema.omnigraph-1.0.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.usd.schema.omniscripting-1.0.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.usd.schema.scene.visualization-2.0.2+106.1.0", + "~/isaac-sim/extscache/omni.usd.schema.semantics-0.0.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.usd.schema.sequence-2.3.1+106.0.0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.usd_resolver-1.0.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.vdb_timesample_editor-0.1.10", + "~/isaac-sim/extscache/omni.volume-0.5.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/extscache/omni.warehouse_creator-0.4.2", + "~/isaac-sim/extscache/omni.warp-1.2.1", + "~/isaac-sim/extscache/omni.warp.core-1.2.1", + "~/isaac-sim/extscache/semantics.schema.editor-0.3.8", + "~/isaac-sim/extscache/semantics.schema.property-1.0.4", + "~/isaac-sim/extscache/usdrt.scenegraph-7.5.0+10a4b5c0.lx64.r.cp310", + "~/isaac-sim/kit/extscore/omni.assets.plugins", + "~/isaac-sim/kit/extscore/omni.client", + "~/isaac-sim/kit/extscore/omni.kit.async_engine", + "~/isaac-sim/kit/extscore/omni.kit.registry.nucleus", + "~/isaac-sim/kit/kernel/py", + "~/isaac-sim/kit/plugins/bindings-python", + "~/isaac-sim/kit/python/lib/python3.10/site-packages", + "~/isaac-sim/python_packages", + ] +} diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..6bc78d5 --- /dev/null +++ b/CMakeLists.txt @@ -0,0 +1,31 @@ +cmake_minimum_required(VERSION 3.5) +project(space_robotics_bench) + +## Find dependencies +find_package(ament_cmake REQUIRED) +find_package(ament_cmake_python REQUIRED) + +## Install Python package +ament_python_install_package(space_robotics_bench) + +## Install Python scripts +set(SCRIPTS_DIR scripts) +install(PROGRAMS + ${SCRIPTS_DIR}/blender/procgen_assets.py + ${SCRIPTS_DIR}/utils/clean_procgen_cache.py + ${SCRIPTS_DIR}/utils/tensorboard.bash + ${SCRIPTS_DIR}/utils/update_assets.bash + ${SCRIPTS_DIR}/gui.bash + ${SCRIPTS_DIR}/list_envs.py + ${SCRIPTS_DIR}/random_agent.py + ${SCRIPTS_DIR}/ros2.py + ${SCRIPTS_DIR}/teleop.py + ${SCRIPTS_DIR}/zero_agent.py + DESTINATION lib/${PROJECT_NAME} +) + +## Install directories +install(DIRECTORY config launch DESTINATION share/${PROJECT_NAME}) + +## Setup the project +ament_package() diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..7fb31cb --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,4873 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 3 + +[[package]] +name = "ab_glyph" +version = "0.2.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec3672c180e71eeaaac3a541fbbc5f5ad4def8b747c595ad30d674e43049f7b0" +dependencies = [ + "ab_glyph_rasterizer", + "owned_ttf_parser", +] + +[[package]] +name = "ab_glyph_rasterizer" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c71b1793ee61086797f5c80b6efa2b8ffa6d5dd703f118545808a7f2e27f7046" + +[[package]] +name = "accesskit" +version = "0.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "99b76d84ee70e30a4a7e39ab9018e2b17a6a09e31084176cc7c0b2dec036ba45" +dependencies = [ + "enumn", + "serde", +] + +[[package]] +name = "accesskit_atspi_common" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5393c75d4666f580f4cac0a968bc97c36076bb536a129f28210dac54ee127ed" +dependencies = [ + "accesskit", + "accesskit_consumer", + "atspi-common", + "serde", + "thiserror", + "zvariant", +] + +[[package]] +name = "accesskit_consumer" +version = "0.24.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a12dc159d52233c43d9fe5415969433cbdd52c3d6e0df51bda7d447427b9986" +dependencies = [ + "accesskit", + "immutable-chunkmap", +] + +[[package]] +name = "accesskit_macos" +version = "0.17.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49509c722e8da39e6c569f7d13ec51620208988913e738abbc67e3c65f06e6d5" +dependencies = [ + "accesskit", + "accesskit_consumer", + "objc2", + "objc2-app-kit", + "objc2-foundation", + "once_cell", +] + +[[package]] +name = "accesskit_unix" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be7f5cf6165be10a54b2655fa2e0e12b2509f38ed6fc43e11c31fdb7ee6230bb" +dependencies = [ + "accesskit", + "accesskit_atspi_common", + "async-channel", + "async-executor", + "async-task", + "atspi", + "futures-lite", + "futures-util", + "serde", + "zbus", +] + +[[package]] +name = "accesskit_windows" +version = "0.23.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "974e96c347384d9133427167fb8a58c340cb0496988dacceebdc1ed27071023b" +dependencies = [ + "accesskit", + "accesskit_consumer", + "paste", + "static_assertions", + "windows 0.58.0", + "windows-core 0.58.0", +] + +[[package]] +name = "accesskit_winit" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9987e852fe7e4e5a493f8c8af0b729b47da2750f0dea10a4c8984fe1117ebaec" +dependencies = [ + "accesskit", + "accesskit_macos", + "accesskit_unix", + "accesskit_windows", + "raw-window-handle", + "winit", +] + +[[package]] +name = "adler2" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "512761e0bb2578dd7380c6baaa0f4ce03e84f95e960231d1dec8bf4d7d6e2627" + +[[package]] +name = "ahash" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e89da841a80418a9b391ebaea17f5c112ffaaa96f621d2c285b5174da76b9011" +dependencies = [ + "cfg-if", + "getrandom", + "once_cell", + "serde", + "version_check", + "zerocopy", +] + +[[package]] +name = "aho-corasick" +version = "1.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916" +dependencies = [ + "memchr", +] + +[[package]] +name = "aligned" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "377e4c0ba83e4431b10df45c1d4666f178ea9c552cac93e60c3a88bf32785923" +dependencies = [ + "as-slice", +] + +[[package]] +name = "aligned-array" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e05c92d086290f52938013f6242ac62bf7d401fab8ad36798a609faa65c3fd2c" +dependencies = [ + "generic-array", +] + +[[package]] +name = "android-activity" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef6978589202a00cd7e118380c448a08b6ed394c3a8df3a430d0898e3a42d046" +dependencies = [ + "android-properties", + "bitflags 2.6.0", + "cc", + "cesu8", + "jni", + "jni-sys", + "libc", + "log", + "ndk", + "ndk-context", + "ndk-sys", + "num_enum", + "thiserror", +] + +[[package]] +name = "android-properties" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7eb209b1518d6bb87b283c20095f5228ecda460da70b44f0802523dea6da04" + +[[package]] +name = "android-tzdata" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e999941b234f3131b00bc13c22d06e8c5ff726d1b6318ac7eb276997bbb4fef0" + +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + +[[package]] +name = "arboard" +version = "3.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df099ccb16cd014ff054ac1bf392c67feeef57164b05c42f037cd40f5d4357f4" +dependencies = [ + "clipboard-win", + "log", + "objc2", + "objc2-app-kit", + "objc2-foundation", + "parking_lot", + "x11rb", +] + +[[package]] +name = "arrayref" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" + +[[package]] +name = "arrayvec" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" + +[[package]] +name = "as-raw-xcb-connection" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "175571dd1d178ced59193a6fc02dde1b972eb0bc56c892cde9beeceac5bf0f6b" + +[[package]] +name = "as-slice" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "516b6b4f0e40d50dcda9365d53964ec74560ad4284da2e7fc97122cd83174516" +dependencies = [ + "stable_deref_trait", +] + +[[package]] +name = "async-broadcast" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20cd0e2e25ea8e5f7e9df04578dc6cf5c83577fd09b1a46aaf5c85e1c33f2a7e" +dependencies = [ + "event-listener", + "event-listener-strategy", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-channel" +version = "2.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89b47800b0be77592da0afd425cc03468052844aff33b84e33cc696f64e77b6a" +dependencies = [ + "concurrent-queue", + "event-listener-strategy", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-executor" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30ca9a001c1e8ba5149f91a74362376cc6bc5b919d92d988668657bd570bdcec" +dependencies = [ + "async-task", + "concurrent-queue", + "fastrand", + "futures-lite", + "slab", +] + +[[package]] +name = "async-fs" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebcd09b382f40fcd159c2d695175b2ae620ffa5f3bd6f664131efff4e8b9e04a" +dependencies = [ + "async-lock", + "blocking", + "futures-lite", +] + +[[package]] +name = "async-io" +version = "2.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "444b0228950ee6501b3568d3c93bf1176a1fdbc3b758dcd9475046d30f4dc7e8" +dependencies = [ + "async-lock", + "cfg-if", + "concurrent-queue", + "futures-io", + "futures-lite", + "parking", + "polling", + "rustix", + "slab", + "tracing", + "windows-sys 0.59.0", +] + +[[package]] +name = "async-lock" +version = "3.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff6e472cdea888a4bd64f342f09b3f50e1886d32afe8df3d663c01140b811b18" +dependencies = [ + "event-listener", + "event-listener-strategy", + "pin-project-lite", +] + +[[package]] +name = "async-process" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63255f1dc2381611000436537bbedfe83183faa303a5a0edaf191edef06526bb" +dependencies = [ + "async-channel", + "async-io", + "async-lock", + "async-signal", + "async-task", + "blocking", + "cfg-if", + "event-listener", + "futures-lite", + "rustix", + "tracing", +] + +[[package]] +name = "async-recursion" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b43422f69d8ff38f95f1b2bb76517c91589a924d1559a0e935d7c8ce0274c11" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "async-signal" +version = "0.2.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "637e00349800c0bdf8bfc21ebbc0b6524abea702b0da4168ac00d070d0c0b9f3" +dependencies = [ + "async-io", + "async-lock", + "atomic-waker", + "cfg-if", + "futures-core", + "futures-io", + "rustix", + "signal-hook-registry", + "slab", + "windows-sys 0.59.0", +] + +[[package]] +name = "async-task" +version = "4.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b75356056920673b02621b35afd0f7dda9306d03c79a30f5c56c44cf256e3de" + +[[package]] +name = "async-trait" +version = "0.1.83" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "721cae7de5c34fbb2acd27e21e6d2cf7b886dce0c27388d46c4e6c47ea4318dd" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "atomic" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d818003e740b63afc82337e3160717f4f63078720a810b7b903e70a5d1d2994" +dependencies = [ + "bytemuck", +] + +[[package]] +name = "atomic-waker" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + +[[package]] +name = "atspi" +version = "0.22.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be534b16650e35237bb1ed189ba2aab86ce65e88cc84c66f4935ba38575cecbf" +dependencies = [ + "atspi-common", + "atspi-connection", + "atspi-proxies", +] + +[[package]] +name = "atspi-common" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1909ed2dc01d0a17505d89311d192518507e8a056a48148e3598fef5e7bb6ba7" +dependencies = [ + "enumflags2", + "serde", + "static_assertions", + "zbus", + "zbus-lockstep", + "zbus-lockstep-macros", + "zbus_names", + "zvariant", +] + +[[package]] +name = "atspi-connection" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "430c5960624a4baaa511c9c0fcc2218e3b58f5dbcc47e6190cafee344b873333" +dependencies = [ + "atspi-common", + "atspi-proxies", + "futures-lite", + "zbus", +] + +[[package]] +name = "atspi-proxies" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5e6c5de3e524cf967569722446bcd458d5032348554d9a17d7d72b041ab7496" +dependencies = [ + "atspi-common", + "serde", + "zbus", + "zvariant", +] + +[[package]] +name = "autocfg" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26" + +[[package]] +name = "az" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b7e4c2464d97fe331d41de9d5db0def0a96f4d823b8b32a2efd503578988973" + +[[package]] +name = "base64" +version = "0.21.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567" + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + +[[package]] +name = "bincode" +version = "1.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad" +dependencies = [ + "serde", +] + +[[package]] +name = "bindgen" +version = "0.63.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36d860121800b2a9a94f9b5604b332d5cffb234ce17609ea479d723dbc9d3885" +dependencies = [ + "bitflags 1.3.2", + "cexpr", + "clang-sys", + "lazy_static", + "lazycell", + "log", + "peeking_take_while", + "proc-macro2", + "quote", + "regex", + "rustc-hash", + "shlex", + "syn 1.0.109", + "which", +] + +[[package]] +name = "bit-set" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0700ddab506f33b20a03b13996eccd309a48e5ff77d0d95926aa0210fb4e95f1" +dependencies = [ + "bit-vec", +] + +[[package]] +name = "bit-vec" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb" + +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + +[[package]] +name = "bitflags" +version = "2.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de" +dependencies = [ + "serde", +] + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "block2" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c132eebf10f5cad5289222520a4a058514204aed6d791f1cf4fe8088b82d15f" +dependencies = [ + "objc2", +] + +[[package]] +name = "blocking" +version = "1.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "703f41c54fc768e63e091340b424302bb1c29ef4aa0c7f10fe849dfb114d29ea" +dependencies = [ + "async-channel", + "async-task", + "futures-io", + "futures-lite", + "piper", +] + +[[package]] +name = "bumpalo" +version = "3.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79296716171880943b8470b5f8d03aa55eb2e645a4874bdbb28adb49162e012c" + +[[package]] +name = "bytemuck" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8334215b81e418a0a7bdb8ef0849474f40bb10c8b71f1c4ed315cff49f32494d" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bcfcc3cd946cb52f0bbfdbbcfa2f4e24f75ebb6c0e1002f7c25904fada18b9ec" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + +[[package]] +name = "byteorder-lite" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f1fe948ff07f4bd06c30984e69f5b4899c516a3ef74f34df92a2df2ab535495" + +[[package]] +name = "bytes" +version = "1.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "428d9aa8fbc0670b7b8d6030a7fadd0f86151cae55e4dbbece15f3780a3dfaf3" + +[[package]] +name = "calloop" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b99da2f8558ca23c71f4fd15dc57c906239752dd27ff3c00a1d56b685b7cbfec" +dependencies = [ + "bitflags 2.6.0", + "log", + "polling", + "rustix", + "slab", + "thiserror", +] + +[[package]] +name = "calloop-wayland-source" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95a66a987056935f7efce4ab5668920b5d0dac4a7c99991a67395f13702ddd20" +dependencies = [ + "calloop", + "rustix", + "wayland-backend", + "wayland-client", +] + +[[package]] +name = "cc" +version = "1.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2e7962b54006dcfcc61cb72735f4d89bb97061dd6a7ed882ec6b8ee53714c6f" +dependencies = [ + "jobserver", + "libc", + "shlex", +] + +[[package]] +name = "cesu8" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d43a04d8753f35258c91f8ec639f792891f748a1edbd759cf1dcea3382ad83c" + +[[package]] +name = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom", +] + +[[package]] +name = "cfg-if" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" + +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + +[[package]] +name = "cgl" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ced0551234e87afee12411d535648dd89d2e7f34c78b753395567aff3d447ff" +dependencies = [ + "libc", +] + +[[package]] +name = "chrono" +version = "0.4.38" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a21f936df1771bf62b77f047b726c4625ff2e8aa607c01ec06e5a05bd8463401" +dependencies = [ + "android-tzdata", + "iana-time-zone", + "js-sys", + "num-traits", + "wasm-bindgen", + "windows-targets 0.52.6", +] + +[[package]] +name = "clang-sys" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[package]] +name = "clipboard-win" +version = "5.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15efe7a882b08f34e38556b14f2fb3daa98769d06c7f0c1b076dfd0d983bc892" +dependencies = [ + "error-code", +] + +[[package]] +name = "color_quant" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d7b894f5411737b7867f4827955924d7c254fc9f4d91a6aad6b097804b1018b" + +[[package]] +name = "combine" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" +dependencies = [ + "bytes", + "memchr", +] + +[[package]] +name = "concurrent-queue" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "const_format" +version = "0.2.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50c655d81ff1114fb0dcdea9225ea9f0cc712a6f8d189378e82bdf62a473a64b" +dependencies = [ + "const_format_proc_macros", + "konst", +] + +[[package]] +name = "const_format_proc_macros" +version = "0.2.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eff1a44b93f47b1bac19a27932f5c591e43d1ba357ee4f61526c8a25603f0eb1" +dependencies = [ + "proc-macro2", + "quote", + "unicode-xid", +] + +[[package]] +name = "core-foundation" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e195e091a93c46f7102ec7818a2aa394e1e1771c3ab4825963fa03e45afb8f" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b55271e5c8c478ad3f38ad24ef34923091e0548492a266d19b3c0b4d82574c63" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "core-graphics" +version = "0.23.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c07782be35f9e1140080c6b96f0d44b739e2278479f64e02fdab4e32dfd8b081" +dependencies = [ + "bitflags 1.3.2", + "core-foundation 0.9.4", + "core-graphics-types", + "foreign-types", + "libc", +] + +[[package]] +name = "core-graphics-types" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45390e6114f68f718cc7a830514a96f903cccd70d02a8f6d9f643ac4ba45afaf" +dependencies = [ + "bitflags 1.3.2", + "core-foundation 0.9.4", + "libc", +] + +[[package]] +name = "cpufeatures" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "608697df725056feaccfa42cffdaeeec3fccc4ffc38358ecd19b243e716a78e0" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a97769d94ddab943e4510d138150169a2758b5ef3eb191a9ee688de3e23ef7b3" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80" + +[[package]] +name = "crunchy" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7" + +[[package]] +name = "crypto-common" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "cursor-icon" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96a6ac251f4a2aca6b3f91340350eab87ae57c3f127ffeb585e92bd336717991" + +[[package]] +name = "data-url" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c297a1c74b71ae29df00c3e22dd9534821d60eb9af5a0192823fa2acea70c2a" + +[[package]] +name = "deranged" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b42b6fa04a440b495c8b04d0e71b707c585f83cb9cb28cf8cd0d976c315e31b4" +dependencies = [ + "powerfmt", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "dispatch" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd0c93bb4b0c6d9b77f4435b0ae98c24d17f1c45b2ff844c6151a07256ca923b" + +[[package]] +name = "display_json" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e8619be5104f2a5fe15929a69a960db3dc3f71a879f15e1c4255083d9b57a1c" +dependencies = [ + "proc-macro2", + "quote", + "serde", + "serde_json", + "syn 1.0.109", +] + +[[package]] +name = "divrem" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69dde51e8fef5e12c1d65e0929b03d66e4c0c18282bc30ed2ca050ad6f44dd82" + +[[package]] +name = "dlib" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330c60081dcc4c72131f8eb70510f1ac07223e5d4163db481a04a0befcffa412" +dependencies = [ + "libloading", +] + +[[package]] +name = "doc-comment" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fea41bba32d969b513997752735605054bc0dfa92b4c56bf1189f2e174be7a10" + +[[package]] +name = "document-features" +version = "0.2.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb6969eaabd2421f8a2775cfd2471a2b634372b4a25d41e3bd647b79912850a0" +dependencies = [ + "litrs", +] + +[[package]] +name = "downcast-rs" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75b325c5dbd37f80359721ad39aca5a29fb04c89279657cffdda8736d0c0b9d2" + +[[package]] +name = "dpi" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f25c0e292a7ca6d6498557ff1df68f32c99850012b6ea401cf8daf771f22ff53" + +[[package]] +name = "ecolor" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "775cfde491852059e386c4e1deb4aef381c617dc364184c6f6afee99b87c402b" +dependencies = [ + "bytemuck", + "emath", + "serde", +] + +[[package]] +name = "eframe" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ac2645a9bf4826eb4e91488b1f17b8eaddeef09396706b2f14066461338e24f" +dependencies = [ + "ahash", + "bytemuck", + "document-features", + "egui", + "egui-winit", + "egui_glow", + "glow", + "glutin", + "glutin-winit", + "home", + "image", + "js-sys", + "log", + "objc2", + "objc2-app-kit", + "objc2-foundation", + "parking_lot", + "percent-encoding", + "raw-window-handle", + "ron", + "serde", + "static_assertions", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", + "web-time", + "winapi", + "windows-sys 0.52.0", + "winit", +] + +[[package]] +name = "egui" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53eafabcce0cb2325a59a98736efe0bf060585b437763f8c476957fb274bb974" +dependencies = [ + "accesskit", + "ahash", + "emath", + "epaint", + "log", + "nohash-hasher", + "ron", + "serde", +] + +[[package]] +name = "egui-winit" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0a9c430f4f816340e8e8c1b20eec274186b1be6bc4c7dfc467ed50d57abc36c6" +dependencies = [ + "accesskit_winit", + "ahash", + "arboard", + "egui", + "log", + "raw-window-handle", + "serde", + "smithay-clipboard", + "web-time", + "webbrowser", + "winit", +] + +[[package]] +name = "egui_commonmark" +version = "0.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f169228b94d1c8eb9330b7ea1b5f65b1193b6bea95159c87f574ed4aff8c172" +dependencies = [ + "egui", + "egui_commonmark_backend", + "egui_commonmark_macros", + "egui_extras", + "pulldown-cmark", +] + +[[package]] +name = "egui_commonmark_backend" +version = "0.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dcd08f95abeb137e59c9bfdd0880d362bff74a83afe13805fde7a2d014ef773d" +dependencies = [ + "egui", + "egui_extras", + "pulldown-cmark", + "syntect", +] + +[[package]] +name = "egui_commonmark_macros" +version = "0.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb1af8e42cb5ab90dc11600ac50fc81f4f87e46f1f6fce6b602f70b45d625fdd" +dependencies = [ + "egui", + "egui_commonmark_backend", + "proc-macro-crate", + "proc-macro2", + "pulldown-cmark", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "egui_extras" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf3c1f5cd8dfe2ade470a218696c66cf556fcfd701e7830fa2e9f4428292a2a1" +dependencies = [ + "ahash", + "egui", + "ehttp", + "enum-map", + "image", + "log", + "mime_guess2", + "resvg", + "syntect", +] + +[[package]] +name = "egui_glow" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e39bccc683cd43adab530d8f21a13eb91e80de10bcc38c3f1c16601b6f62b26" +dependencies = [ + "ahash", + "bytemuck", + "egui", + "glow", + "log", + "memoffset", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "ehttp" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59a81c221a1e4dad06cb9c9deb19aea1193a5eea084e8cd42d869068132bf876" +dependencies = [ + "document-features", + "js-sys", + "ureq", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + +[[package]] +name = "either" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0" + +[[package]] +name = "elapsed" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f4e5af126dafd0741c2ad62d47f68b28602550102e5f0dd45c8a97fc8b49c29" + +[[package]] +name = "emath" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1fe0049ce51d0fb414d029e668dd72eb30bc2b739bf34296ed97bd33df544f3" +dependencies = [ + "bytemuck", + "serde", +] + +[[package]] +name = "endi" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3d8a32ae18130a3c84dd492d4215c3d913c3b07c6b63c2eb3eb7ff1101ab7bf" + +[[package]] +name = "enum-map" +version = "2.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6866f3bfdf8207509a033af1a75a7b08abda06bbaaeae6669323fd5a097df2e9" +dependencies = [ + "enum-map-derive", + "serde", +] + +[[package]] +name = "enum-map-derive" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f282cfdfe92516eb26c2af8589c274c7c17681f5ecc03c18255fe741c6aa64eb" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "enumflags2" +version = "0.7.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d232db7f5956f3f14313dc2f87985c58bd2c695ce124c8cdd984e08e15ac133d" +dependencies = [ + "enumflags2_derive", + "serde", +] + +[[package]] +name = "enumflags2_derive" +version = "0.7.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de0d48a183585823424a4ce1aa132d174a6a81bd540895822eb4c8373a8e49e8" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "enumn" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f9ed6b3789237c8a0c1c505af1c7eb2c560df6186f01b098c3a1064ea532f38" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "epaint" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a32af8da821bd4f43f2c137e295459ee2e1661d87ca8779dfa0eaf45d870e20f" +dependencies = [ + "ab_glyph", + "ahash", + "bytemuck", + "ecolor", + "emath", + "log", + "nohash-hasher", + "parking_lot", + "serde", +] + +[[package]] +name = "equivalent" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" + +[[package]] +name = "errno" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "534c5cf6194dfab3db3242765c03bbe257cf92f22b38f6bc0c58d59108a820ba" +dependencies = [ + "libc", + "windows-sys 0.52.0", +] + +[[package]] +name = "error-code" +version = "3.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5d9305ccc6942a704f4335694ecd3de2ea531b114ac2d51f5f843750787a92f" + +[[package]] +name = "event-listener" +version = "5.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6032be9bd27023a771701cc49f9f053c751055f71efb2e0ae5c15809093675ba" +dependencies = [ + "concurrent-queue", + "parking", + "pin-project-lite", +] + +[[package]] +name = "event-listener-strategy" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f214dc438f977e6d4e3500aaa277f5ad94ca83fbbd9b1a15713ce2344ccc5a1" +dependencies = [ + "event-listener", + "pin-project-lite", +] + +[[package]] +name = "fancy-regex" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b95f7c0680e4142284cf8b22c14a476e87d61b004a3a0861872b32ef7ead40a2" +dependencies = [ + "bit-set", + "regex", +] + +[[package]] +name = "fast_poisson" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2472baa9796d2ee497bd61690e3093a26935390d8ce0dd0ddc2db9b47a65898f" +dependencies = [ + "kiddo", + "rand", + "rand_distr", + "rand_xoshiro", +] + +[[package]] +name = "fastrand" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8c02a5121d4ea3eb16a80748c74f5549a5665e4c21333c6098f283870fbdea6" + +[[package]] +name = "fdeflate" +version = "0.3.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8090f921a24b04994d9929e204f50b498a33ea6ba559ffaa05e04f7ee7fb5ab" +dependencies = [ + "simd-adler32", +] + +[[package]] +name = "figment" +version = "0.10.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8cb01cd46b0cf372153850f4c6c272d9cbea2da513e07538405148f95bd789f3" +dependencies = [ + "atomic", + "parking_lot", + "pear", + "serde", + "serde_json", + "serde_yaml", + "tempfile", + "toml", + "uncased", + "version_check", +] + +[[package]] +name = "fixed" +version = "1.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85c6e0b89bf864acd20590dbdbad56f69aeb898abfc9443008fd7bd48b2cc85a" +dependencies = [ + "az", + "bytemuck", + "half", + "num-traits", + "typenum", +] + +[[package]] +name = "flate2" +version = "1.0.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1b589b4dc103969ad3cf85c950899926ec64300a1a46d76c03a6072957036f0" +dependencies = [ + "crc32fast", + "miniz_oxide", +] + +[[package]] +name = "float-cmp" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98de4bbd547a563b716d8dfa9aad1cb19bfab00f4fa09a6a4ed21dbcf44ce9c4" + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "force-send-sync" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32688dcc448aa684426ecc398f9af78c55b0769515ccd12baa565daf6f32feb6" + +[[package]] +name = "foreign-types" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d737d9aa519fb7b749cbc3b962edcf310a8dd1f4b67c91c4f83975dbdd17d965" +dependencies = [ + "foreign-types-macros", + "foreign-types-shared", +] + +[[package]] +name = "foreign-types-macros" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a5c6c585bc94aaf2c7b51dd4c2ba22680844aba4c687be581871a6f518c5742" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "foreign-types-shared" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa9a19cbb55df58761df49b23516a86d432839add4af60fc256da840f66ed35b" + +[[package]] +name = "form_urlencoded" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13624c2627564efccf4934284bdd98cbaa14e79b0b5a141218e507b3a823456" +dependencies = [ + "percent-encoding", +] + +[[package]] +name = "futures" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2dff15bf788c671c1934e366d07e30c1814a8ef514e1af724a602e8a2fbe1b10" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e" + +[[package]] +name = "futures-executor" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e28d1d997f585e54aebc3f97d39e72338912123a67330d723fdbb564d646c9f" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" + +[[package]] +name = "futures-lite" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52527eb5074e35e9339c6b4e8d12600c7128b68fb25dcb9fa9dec18f7c25f3a5" +dependencies = [ + "fastrand", + "futures-core", + "futures-io", + "parking", + "pin-project-lite", +] + +[[package]] +name = "futures-macro" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "futures-sink" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e575fab7d1e0dcb8d0c7bcf9a63ee213816ab51902e6d244a95819acacf1d4f7" + +[[package]] +name = "futures-task" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" + +[[package]] +name = "futures-util" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "pin-utils", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "gethostname" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0176e0459c2e4a1fe232f984bca6890e681076abb9934f6cea7c326f3fc47818" +dependencies = [ + "libc", + "windows-targets 0.48.5", +] + +[[package]] +name = "getrandom" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + +[[package]] +name = "gif" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fb2d69b19215e18bb912fa30f7ce15846e301408695e44e0ef719f1da9e19f2" +dependencies = [ + "color_quant", + "weezl", +] + +[[package]] +name = "gl_generator" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a95dfc23a2b4a9a2f5ab41d194f8bfda3cabec42af4e39f08c339eb2a0c124d" +dependencies = [ + "khronos_api", + "log", + "xml-rs", +] + +[[package]] +name = "glob" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2fabcfbdc87f4758337ca535fb41a6d701b65693ce38287d856d1674551ec9b" + +[[package]] +name = "glow" +version = "0.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d51fa363f025f5c111e03f13eda21162faeacb6911fe8caa0c0349f9cf0c4483" +dependencies = [ + "js-sys", + "slotmap", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "glutin" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec69412a0bf07ea7607e638b415447857a808846c2b685a43c8aa18bc6d5e499" +dependencies = [ + "bitflags 2.6.0", + "cfg_aliases", + "cgl", + "core-foundation 0.9.4", + "dispatch", + "glutin_egl_sys", + "glutin_glx_sys", + "glutin_wgl_sys", + "libloading", + "objc2", + "objc2-app-kit", + "objc2-foundation", + "once_cell", + "raw-window-handle", + "wayland-sys", + "windows-sys 0.52.0", + "x11-dl", +] + +[[package]] +name = "glutin-winit" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85edca7075f8fc728f28cb8fbb111a96c3b89e930574369e3e9c27eb75d3788f" +dependencies = [ + "cfg_aliases", + "glutin", + "raw-window-handle", + "winit", +] + +[[package]] +name = "glutin_egl_sys" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cae99fff4d2850dbe6fb8c1fa8e4fead5525bab715beaacfccf3fb994e01c827" +dependencies = [ + "gl_generator", + "windows-sys 0.52.0", +] + +[[package]] +name = "glutin_glx_sys" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c2b2d3918e76e18e08796b55eb64e8fe6ec67d5a6b2e2a7e2edce224ad24c63" +dependencies = [ + "gl_generator", + "x11-dl", +] + +[[package]] +name = "glutin_wgl_sys" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0a4e1951bbd9434a81aa496fe59ccc2235af3820d27b85f9314e279609211e2c" +dependencies = [ + "gl_generator", +] + +[[package]] +name = "half" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dd08c532ae367adf81c312a4580bc67f1d0fe8bc9c460520283f4c0ff277888" +dependencies = [ + "cfg-if", + "crunchy", +] + +[[package]] +name = "hashbrown" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e087f84d4f86bf4b218b927129862374b72199ae7d8657835f1e89000eea4fb" + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "hermit-abi" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fbf6a919d6cf397374f7dfeeea91d974c7c0a7221d0d0f4f20d859d329e53fcc" + +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + +[[package]] +name = "home" +version = "0.5.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3d1354bf6b7235cb4a0576c2619fd4ed18183f689b12b006a0ee7329eeff9a5" +dependencies = [ + "windows-sys 0.52.0", +] + +[[package]] +name = "iana-time-zone" +version = "0.1.61" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "235e081f3925a06703c2d0117ea8b91f042756fd6e7a6e5d901e8ca1a996b220" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "wasm-bindgen", + "windows-core 0.52.0", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + +[[package]] +name = "idna" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "634d9b1461af396cad843f47fdba5597a4f9e6ddd4bfb6ff5d85028c25cb12f6" +dependencies = [ + "unicode-bidi", + "unicode-normalization", +] + +[[package]] +name = "image" +version = "0.25.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc144d44a31d753b02ce64093d532f55ff8dc4ebf2ffb8a63c0dda691385acae" +dependencies = [ + "bytemuck", + "byteorder-lite", + "color_quant", + "gif", + "num-traits", + "png", + "zune-core", + "zune-jpeg", +] + +[[package]] +name = "imagesize" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "029d73f573d8e8d63e6d5020011d3255b28c3ba85d6cf870a07184ed23de9284" + +[[package]] +name = "immutable-chunkmap" +version = "2.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4419f022e55cc63d5bbd6b44b71e1d226b9c9480a47824c706e9d54e5c40c5eb" +dependencies = [ + "arrayvec", +] + +[[package]] +name = "indexmap" +version = "2.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "707907fe3c25f5424cce2cb7e1cbcafee6bdbe735ca90ef77c29e84591e5b9da" +dependencies = [ + "equivalent", + "hashbrown", +] + +[[package]] +name = "indoc" +version = "2.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b248f5224d1d606005e02c97f5aa4e88eeb230488bcc03bc9ca4d7991399f2b5" + +[[package]] +name = "inlinable_string" +version = "0.1.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8fae54786f62fb2918dcfae3d568594e50eb9b5c25bf04371af6fe7516452fb" + +[[package]] +name = "itertools" +version = "0.10.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" +dependencies = [ + "either", +] + +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" + +[[package]] +name = "jni" +version = "0.21.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a87aa2bb7d2af34197c04845522473242e1aa17c12f4935d5856491a7fb8c97" +dependencies = [ + "cesu8", + "cfg-if", + "combine", + "jni-sys", + "log", + "thiserror", + "walkdir", + "windows-sys 0.45.0", +] + +[[package]] +name = "jni-sys" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8eaf4bc02d17cbdd7ff4c7438cafcdf7fb9a4613313ad11b4f8fefe7d3fa0130" + +[[package]] +name = "jobserver" +version = "0.1.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48d1dbcbbeb6a7fec7e059840aa538bd62aaccf972c7346c4d9d2059312853d0" +dependencies = [ + "libc", +] + +[[package]] +name = "js-sys" +version = "0.3.72" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a88f1bda2bd75b0452a14784937d796722fdebfe50df998aeb3f0b7603019a9" +dependencies = [ + "wasm-bindgen", +] + +[[package]] +name = "khronos_api" +version = "3.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2db585e1d738fc771bf08a151420d3ed193d9d895a36df7f6f8a9456b911ddc" + +[[package]] +name = "kiddo" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e1c5ea778d68eacd5c33f29537ba0b7b6c2595e74ee013a69cedc20ab4d3177" +dependencies = [ + "aligned", + "aligned-array", + "az", + "divrem", + "doc-comment", + "elapsed", + "fixed", + "log", + "min-max-heap", + "num-traits", + "rand", + "rayon", +] + +[[package]] +name = "konst" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330f0e13e6483b8c34885f7e6c9f19b1a7bd449c673fbb948a51c99d66ef74f4" +dependencies = [ + "konst_macro_rules", +] + +[[package]] +name = "konst_macro_rules" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4933f3f57a8e9d9da04db23fb153356ecaf00cbd14aee46279c33dc80925c37" + +[[package]] +name = "kurbo" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd85a5776cd9500c2e2059c8c76c3b01528566b7fcbaf8098b55a33fc298849b" +dependencies = [ + "arrayvec", +] + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.161" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9489c2807c139ffd9c1794f4af0ebe86a828db53ecdc7fea2111d0fed085d1" + +[[package]] +name = "libloading" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4979f22fdb869068da03c9f7528f8297c6fd2606bc3a4affe42e6a823fdb8da4" +dependencies = [ + "cfg-if", + "windows-targets 0.52.6", +] + +[[package]] +name = "libm" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058" + +[[package]] +name = "libredox" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0ff37bd590ca25063e35af745c343cb7a0271906fb7b37e4813e8f79f00268d" +dependencies = [ + "bitflags 2.6.0", + "libc", + "redox_syscall 0.5.7", +] + +[[package]] +name = "linked-hash-map" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0717cef1bc8b636c6e1c1bbdefc09e6322da8a9321966e8928ef80d20f7f770f" + +[[package]] +name = "linux-raw-sys" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78b3ae25bc7c8c38cec158d1f2757ee79e9b3740fbc7ccf0e59e4b08d793fa89" + +[[package]] +name = "litrs" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ce301924b7887e9d637144fdade93f9dfff9b60981d4ac161db09720d39aa5" + +[[package]] +name = "lock_api" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07af8b9cdd281b7915f413fa73f29ebd5d55d0d3f0155584dade1ff18cea1b17" +dependencies = [ + "autocfg", + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7a70ba024b9dc04c27ea2f0c0548feb474ec5c54bba33a7f72f873a39d07b24" + +[[package]] +name = "memchr" +version = "2.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" + +[[package]] +name = "memmap2" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd3f7eed9d3848f8b98834af67102b720745c4ec028fcd0aa0239277e7de374f" +dependencies = [ + "libc", +] + +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + +[[package]] +name = "mime" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" + +[[package]] +name = "mime_guess2" +version = "2.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25a3333bb1609500601edc766a39b4c1772874a4ce26022f4d866854dc020c41" +dependencies = [ + "mime", + "unicase", +] + +[[package]] +name = "min-max-heap" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2687e6cf9c00f48e9284cf9fd15f2ef341d03cc7743abf9df4c5f07fdee50b18" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "miniz_oxide" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2d80299ef12ff69b16a84bb182e3b9df68b5a91574d3d4fa6e41b65deec4df1" +dependencies = [ + "adler2", + "simd-adler32", +] + +[[package]] +name = "ndk" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3f42e7bbe13d351b6bead8286a43aac9534b82bd3cc43e47037f012ebfd62d4" +dependencies = [ + "bitflags 2.6.0", + "jni-sys", + "log", + "ndk-sys", + "num_enum", + "raw-window-handle", + "thiserror", +] + +[[package]] +name = "ndk-context" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27b02d87554356db9e9a873add8782d4ea6e3e58ea071a9adb9a2e8ddb884a8b" + +[[package]] +name = "ndk-sys" +version = "0.6.0+11769913" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee6cda3051665f1fb8d9e08fc35c96d5a244fb1be711a03b71118828afc9a873" +dependencies = [ + "jni-sys", +] + +[[package]] +name = "nix" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46" +dependencies = [ + "bitflags 2.6.0", + "cfg-if", + "cfg_aliases", + "libc", + "memoffset", +] + +[[package]] +name = "nohash-hasher" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2bf50223579dc7cdcfb3bfcacf7069ff68243f8c363f62ffa99cf000a6b9c451" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "ntapi" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8a3895c6391c39d7fe7ebc444a87eb2991b2a0bc718fdabd071eec617fc68e4" +dependencies = [ + "winapi", +] + +[[package]] +name = "nu-ansi-term" +version = "0.46.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84" +dependencies = [ + "overload", + "winapi", +] + +[[package]] +name = "num-conv" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51d515d32fb182ee37cda2ccdcb92950d6a3c2893aa280e540671c2cd0f3b1d9" + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", + "libm", +] + +[[package]] +name = "num_enum" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e613fc340b2220f734a8595782c551f1250e969d87d3be1ae0579e8d4065179" +dependencies = [ + "num_enum_derive", +] + +[[package]] +name = "num_enum_derive" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af1844ef2428cc3e1cb900be36181049ef3d3193c63e43026cfe202983b27a56" +dependencies = [ + "proc-macro-crate", + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "objc-sys" +version = "0.3.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdb91bdd390c7ce1a8607f35f3ca7151b65afc0ff5ff3b34fa350f7d7c7e4310" + +[[package]] +name = "objc2" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46a785d4eeff09c14c487497c162e92766fbb3e4059a71840cecc03d9a50b804" +dependencies = [ + "objc-sys", + "objc2-encode", +] + +[[package]] +name = "objc2-app-kit" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4e89ad9e3d7d297152b17d39ed92cd50ca8063a89a9fa569046d41568891eff" +dependencies = [ + "bitflags 2.6.0", + "block2", + "libc", + "objc2", + "objc2-core-data", + "objc2-core-image", + "objc2-foundation", + "objc2-quartz-core", +] + +[[package]] +name = "objc2-cloud-kit" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74dd3b56391c7a0596a295029734d3c1c5e7e510a4cb30245f8221ccea96b009" +dependencies = [ + "bitflags 2.6.0", + "block2", + "objc2", + "objc2-core-location", + "objc2-foundation", +] + +[[package]] +name = "objc2-contacts" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5ff520e9c33812fd374d8deecef01d4a840e7b41862d849513de77e44aa4889" +dependencies = [ + "block2", + "objc2", + "objc2-foundation", +] + +[[package]] +name = "objc2-core-data" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "617fbf49e071c178c0b24c080767db52958f716d9eabdf0890523aeae54773ef" +dependencies = [ + "bitflags 2.6.0", + "block2", + "objc2", + "objc2-foundation", +] + +[[package]] +name = "objc2-core-image" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55260963a527c99f1819c4f8e3b47fe04f9650694ef348ffd2227e8196d34c80" +dependencies = [ + "block2", + "objc2", + "objc2-foundation", + "objc2-metal", +] + +[[package]] +name = "objc2-core-location" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "000cfee34e683244f284252ee206a27953279d370e309649dc3ee317b37e5781" +dependencies = [ + "block2", + "objc2", + "objc2-contacts", + "objc2-foundation", +] + +[[package]] +name = "objc2-encode" +version = "4.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7891e71393cd1f227313c9379a26a584ff3d7e6e7159e988851f0934c993f0f8" + +[[package]] +name = "objc2-foundation" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ee638a5da3799329310ad4cfa62fbf045d5f56e3ef5ba4149e7452dcf89d5a8" +dependencies = [ + "bitflags 2.6.0", + "block2", + "dispatch", + "libc", + "objc2", +] + +[[package]] +name = "objc2-link-presentation" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1a1ae721c5e35be65f01a03b6d2ac13a54cb4fa70d8a5da293d7b0020261398" +dependencies = [ + "block2", + "objc2", + "objc2-app-kit", + "objc2-foundation", +] + +[[package]] +name = "objc2-metal" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd0cba1276f6023976a406a14ffa85e1fdd19df6b0f737b063b95f6c8c7aadd6" +dependencies = [ + "bitflags 2.6.0", + "block2", + "objc2", + "objc2-foundation", +] + +[[package]] +name = "objc2-quartz-core" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e42bee7bff906b14b167da2bac5efe6b6a07e6f7c0a21a7308d40c960242dc7a" +dependencies = [ + "bitflags 2.6.0", + "block2", + "objc2", + "objc2-foundation", + "objc2-metal", +] + +[[package]] +name = "objc2-symbols" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0a684efe3dec1b305badae1a28f6555f6ddd3bb2c2267896782858d5a78404dc" +dependencies = [ + "objc2", + "objc2-foundation", +] + +[[package]] +name = "objc2-ui-kit" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8bb46798b20cd6b91cbd113524c490f1686f4c4e8f49502431415f3512e2b6f" +dependencies = [ + "bitflags 2.6.0", + "block2", + "objc2", + "objc2-cloud-kit", + "objc2-core-data", + "objc2-core-image", + "objc2-core-location", + "objc2-foundation", + "objc2-link-presentation", + "objc2-quartz-core", + "objc2-symbols", + "objc2-uniform-type-identifiers", + "objc2-user-notifications", +] + +[[package]] +name = "objc2-uniform-type-identifiers" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44fa5f9748dbfe1ca6c0b79ad20725a11eca7c2218bceb4b005cb1be26273bfe" +dependencies = [ + "block2", + "objc2", + "objc2-foundation", +] + +[[package]] +name = "objc2-user-notifications" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76cfcbf642358e8689af64cee815d139339f3ed8ad05103ed5eaf73db8d84cb3" +dependencies = [ + "bitflags 2.6.0", + "block2", + "objc2", + "objc2-core-location", + "objc2-foundation", +] + +[[package]] +name = "once_cell" +version = "1.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1261fe7e33c73b354eab43b1273a57c8f967d0391e80353e51f764ac02cf6775" + +[[package]] +name = "orbclient" +version = "0.3.48" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba0b26cec2e24f08ed8bb31519a9333140a6599b867dac464bb150bdb796fd43" +dependencies = [ + "libredox", +] + +[[package]] +name = "ordered-stream" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9aa2b01e1d916879f73a53d01d1d6cee68adbb31d6d9177a8cfce093cced1d50" +dependencies = [ + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "os_str_bytes" +version = "6.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2355d85b9a3786f481747ced0e0ff2ba35213a1f9bd406ed906554d7af805a1" +dependencies = [ + "memchr", +] + +[[package]] +name = "overload" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" + +[[package]] +name = "owned_ttf_parser" +version = "0.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22ec719bbf3b2a81c109a4e20b1f129b5566b7dce654bc3872f6a05abf82b2c4" +dependencies = [ + "ttf-parser", +] + +[[package]] +name = "parking" +version = "2.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" + +[[package]] +name = "parking_lot" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1bf18183cf54e8d6059647fc3063646a1801cf30896933ec2311622cc4b9a27" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e401f977ab385c9e4e3ab30627d6f26d00e2c73eef317493c4ec6d468726cf8" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall 0.5.7", + "smallvec", + "windows-targets 0.52.6", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "pear" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bdeeaa00ce488657faba8ebf44ab9361f9365a97bd39ffb8a60663f57ff4b467" +dependencies = [ + "inlinable_string", + "pear_codegen", + "yansi", +] + +[[package]] +name = "pear_codegen" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4bab5b985dc082b345f812b7df84e1bef27e7207b39e448439ba8bd69c93f147" +dependencies = [ + "proc-macro2", + "proc-macro2-diagnostics", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "peeking_take_while" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19b17cddbe7ec3f8bc800887bab5e717348c95ea2ca0b1bf0837fb964dc67099" + +[[package]] +name = "percent-encoding" +version = "2.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" + +[[package]] +name = "phf" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ade2d8b8f33c7333b51bcf0428d37e217e9f32192ae4772156f65063b8ce03dc" +dependencies = [ + "phf_macros", + "phf_shared", +] + +[[package]] +name = "phf_generator" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48e4cc64c2ad9ebe670cb8fd69dd50ae301650392e81c05f9bfcb2d5bdbc24b0" +dependencies = [ + "phf_shared", + "rand", +] + +[[package]] +name = "phf_macros" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3444646e286606587e49f3bcf1679b8cef1dc2c5ecc29ddacaffc305180d464b" +dependencies = [ + "phf_generator", + "phf_shared", + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "phf_shared" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90fcb95eef784c2ac79119d1dd819e162b5da872ce6f3c3abe1e8ca1c082f72b" +dependencies = [ + "siphasher", +] + +[[package]] +name = "pico-args" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5be167a7af36ee22fe3115051bc51f6e6c7054c9348e28deb4f49bd6f705a315" + +[[package]] +name = "pin-project" +version = "1.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf123a161dde1e524adf36f90bc5d8d3462824a9c43553ad07a8183161189ec" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4502d8515ca9f32f1fb543d987f63d95a14934883db45bdb48060b6b69257f8" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bda66fc9667c18cb2758a2ac84d1167245054bcf85d5d1aaa6923f45801bdd02" + +[[package]] +name = "pin-utils" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" + +[[package]] +name = "piper" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96c8c490f422ef9a4efd2cb5b42b76c8613d7e7dfc1caf667b8a3350a5acc066" +dependencies = [ + "atomic-waker", + "fastrand", + "futures-io", +] + +[[package]] +name = "pkg-config" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "953ec861398dccce10c670dfeaf3ec4911ca479e9c02154b3a215178c5f566f2" + +[[package]] +name = "plist" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42cf17e9a1800f5f396bc67d193dc9411b59012a5876445ef450d449881e1016" +dependencies = [ + "base64 0.22.1", + "indexmap", + "quick-xml 0.32.0", + "serde", + "time", +] + +[[package]] +name = "png" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52f9d46a34a05a6a57566bc2bfae066ef07585a6e3fa30fbbdff5936380623f0" +dependencies = [ + "bitflags 1.3.2", + "crc32fast", + "fdeflate", + "flate2", + "miniz_oxide", +] + +[[package]] +name = "polling" +version = "3.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc2790cd301dec6cd3b7a025e4815cf825724a51c98dccfe6a3e55f05ffb6511" +dependencies = [ + "cfg-if", + "concurrent-queue", + "hermit-abi", + "pin-project-lite", + "rustix", + "tracing", + "windows-sys 0.59.0", +] + +[[package]] +name = "portable-atomic" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc9c68a3f6da06753e9335d63e27f6b9754dd1920d941135b7ea8224f141adb2" + +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + +[[package]] +name = "ppv-lite86" +version = "0.2.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "prettyplease" +version = "0.2.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "910d41a655dac3b764f1ade94821093d3610248694320cd072303a8eedcf221d" +dependencies = [ + "proc-macro2", + "syn 2.0.82", +] + +[[package]] +name = "proc-macro-crate" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecf48c7ca261d60b74ab1a7b20da18bede46776b2e55535cb958eb595c5fa7b" +dependencies = [ + "toml_edit", +] + +[[package]] +name = "proc-macro2" +version = "1.0.88" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c3a7fc5db1e57d5a779a352c8cdb57b29aa4c40cc69c3a68a7fedc815fbf2f9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "proc-macro2-diagnostics" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af066a9c399a26e020ada66a034357a868728e72cd426f3adcd35f80d88d88c8" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", + "version_check", + "yansi", +] + +[[package]] +name = "pulldown-cmark" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f86ba2052aebccc42cbbb3ed234b8b13ce76f75c3551a303cb2bcffcff12bb14" +dependencies = [ + "bitflags 2.6.0", + "memchr", + "unicase", +] + +[[package]] +name = "pyo3" +version = "0.22.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d922163ba1f79c04bc49073ba7b32fd5a8d3b76a87c955921234b8e77333c51" +dependencies = [ + "cfg-if", + "indoc", + "libc", + "memoffset", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", + "unindent", +] + +[[package]] +name = "pyo3-build-config" +version = "0.22.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc38c5feeb496c8321091edf3d63e9a6829eab4b863b4a6a65f26f3e9cc6b179" +dependencies = [ + "once_cell", + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.22.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94845622d88ae274d2729fcefc850e63d7a3ddff5e3ce11bd88486db9f1d357d" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.22.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e655aad15e09b94ffdb3ce3d217acf652e26bbc37697ef012f5e5e348c716e5e" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.22.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae1e3f09eecd94618f60a455a23def79f79eba4dc561a97324bf9ac8c6df30ce" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "quick-xml" +version = "0.30.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eff6510e86862b57b210fd8cbe8ed3f0d7d600b9c2863cd4549a2e033c66e956" +dependencies = [ + "memchr", + "serde", +] + +[[package]] +name = "quick-xml" +version = "0.32.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d3a6e5838b60e0e8fa7a43f22ade549a37d61f8bdbe636d0d7816191de969c2" +dependencies = [ + "memchr", +] + +[[package]] +name = "quick-xml" +version = "0.36.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7649a7b4df05aed9ea7ec6f628c67c9953a43869b8bc50929569b2999d443fe" +dependencies = [ + "memchr", +] + +[[package]] +name = "quote" +version = "1.0.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5b9d34b8991d19d98081b46eacdd8eb58c6f2b201139f7c5f643cc155a633af" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r2r" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "71a6ecd5ceee58172c691c22ecfc013ef4621d38eb6804eb5bea0271f6f53daa" +dependencies = [ + "force-send-sync", + "futures", + "indexmap", + "lazy_static", + "log", + "phf", + "prettyplease", + "proc-macro2", + "quote", + "r2r_actions", + "r2r_common", + "r2r_macros", + "r2r_msg_gen", + "r2r_rcl", + "rayon", + "serde", + "serde_json", + "syn 2.0.82", + "thiserror", + "uuid", +] + +[[package]] +name = "r2r_actions" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57386bf86b3db4e77bd91d6faa33bce0dba4389cefd4ba808a5b3dc5f870dc96" +dependencies = [ + "bindgen", + "r2r_common", + "r2r_msg_gen", + "r2r_rcl", +] + +[[package]] +name = "r2r_common" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f663cf4827e641a11b670844ce1cfa67d7ecf625ada3b91246515ee26c896b18" +dependencies = [ + "bindgen", + "os_str_bytes", + "regex", + "sha2", +] + +[[package]] +name = "r2r_macros" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7092302580ef2a6febeda5db5fd6af4eb28400eb5e0da6c729dcc8b41990d137" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "r2r_msg_gen" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c538acb9f8150028d1a8163eda1c48008551374f6a317cf52c30a3f06d4bef5b" +dependencies = [ + "bindgen", + "force-send-sync", + "itertools 0.10.5", + "phf", + "proc-macro2", + "quote", + "r2r_common", + "r2r_rcl", + "rayon", + "syn 2.0.82", +] + +[[package]] +name = "r2r_rcl" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6a333ef3afc5cda5ddbf2c0d0b0c0ff01e5f8dae542b109112dad0142b47531" +dependencies = [ + "bindgen", + "paste", + "r2r_common", + "widestring", +] + +[[package]] +name = "rand" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +dependencies = [ + "libc", + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom", +] + +[[package]] +name = "rand_distr" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32cb0b9bc82b0a0876c2dd994a7e7a2683d3e7390ca40e6886785ef0c7e3ee31" +dependencies = [ + "num-traits", + "rand", +] + +[[package]] +name = "rand_xoshiro" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f97cdb2a36ed4183de61b2f824cc45c9f1037f28afe0a322e9fff4c108b5aaa" +dependencies = [ + "rand_core", +] + +[[package]] +name = "raw-window-handle" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539" + +[[package]] +name = "rayon" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b418a60154510ca1a002a752ca9714984e21e4241e804d32555251faf8b78ffa" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1465873a3dfdaa8ae7cb14b4383657caab0b3e8a0aa9ae8e04b044854c8dfce2" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "rctree" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b42e27ef78c35d3998403c1d26f3efd9e135d3e5121b0a4845cc5cc27547f4f" + +[[package]] +name = "redox_syscall" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4722d768eff46b75989dd134e5c353f0d6296e5aaa3132e776cbdb56be7731aa" +dependencies = [ + "bitflags 1.3.2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b6dfecf2c74bce2466cabf93f6664d6998a69eb21e39f4207930065b27b771f" +dependencies = [ + "bitflags 2.6.0", +] + +[[package]] +name = "regex" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38200e5ee88914975b69f657f0801b6f6dccafd44fd9326302a4aaeecfacb1d8" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "368758f23274712b504848e9d5a6f010445cc8b87a7cdb4d7cbee666c1288da3" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c" + +[[package]] +name = "resvg" +version = "0.37.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cadccb3d99a9efb8e5e00c16fbb732cbe400db2ec7fc004697ee7d97d86cf1f4" +dependencies = [ + "log", + "pico-args", + "rgb", + "svgtypes", + "tiny-skia", + "usvg", +] + +[[package]] +name = "rgb" +version = "0.8.50" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57397d16646700483b67d2dd6511d79318f9d057fdbd21a4066aeac8b41d310a" +dependencies = [ + "bytemuck", +] + +[[package]] +name = "ring" +version = "0.17.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c17fa4cb658e3583423e915b9f3acc01cceaee1860e33d59ebae66adc3a2dc0d" +dependencies = [ + "cc", + "cfg-if", + "getrandom", + "libc", + "spin", + "untrusted", + "windows-sys 0.52.0", +] + +[[package]] +name = "ron" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b91f7eff05f748767f183df4320a63d6936e9c6107d97c9e6bdd9784f4289c94" +dependencies = [ + "base64 0.21.7", + "bitflags 2.6.0", + "serde", + "serde_derive", +] + +[[package]] +name = "roxmltree" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3cd14fd5e3b777a7422cca79358c57a8f6e3a703d9ac187448d0daf220c2407f" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustix" +version = "0.38.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8acb788b847c24f28525660c4d7758620a7210875711f79e7f663cc152726811" +dependencies = [ + "bitflags 2.6.0", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.52.0", +] + +[[package]] +name = "rustls" +version = "0.23.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fbb44d7acc4e873d613422379f69f237a1b141928c02f6bc6ccfddddc2d7993" +dependencies = [ + "log", + "once_cell", + "ring", + "rustls-pki-types", + "rustls-webpki", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-pki-types" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16f1201b3c9a7ee8039bcadc17b7e605e2945b27eee7631788c1bd2b0643674b" + +[[package]] +name = "rustls-webpki" +version = "0.102.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64ca1bc8749bd4cf37b5ce386cc146580777b4e8572c7b97baf22c83f444bee9" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + +[[package]] +name = "ryu" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f" + +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "scoped-tls" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1cf6437eb19a8f4a6cc0f7dca544973b0b78843adbfeb3683d1a94a0024a294" + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.210" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8e3592472072e6e22e0a54d5904d9febf8508f65fb8552499a1abc7d1078c3a" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.210" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "243902eda00fad750862fc144cea25caca5e20d615af0a81bee94ca738f1df1f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "serde_json" +version = "1.0.132" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d726bfaff4b320266d395898905d0eba0345aae23b54aee3a737e260fd46db03" +dependencies = [ + "itoa", + "memchr", + "ryu", + "serde", +] + +[[package]] +name = "serde_repr" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c64451ba24fc7a6a2d60fc75dd9c83c90903b19028d4eff35e88fc1e86564e9" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "serde_spanned" +version = "0.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87607cb1398ed59d48732e575a4c28a7a8ebf2454b964fe3f224f2afc07909e1" +dependencies = [ + "serde", +] + +[[package]] +name = "serde_yaml" +version = "0.9.34+deprecated" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8b1a1a2ebf674015cc02edccce75287f1a0130d394307b36743c2f5d504b47" +dependencies = [ + "indexmap", + "itoa", + "ryu", + "serde", + "unsafe-libyaml", +] + +[[package]] +name = "sha1" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "sha2" +version = "0.10.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "793db75ad2bcafc3ffa7c68b215fee268f537982cd901d132f89c6343f3a3dc8" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "signal-hook-registry" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9e9e0b4211b72e7b8b6e85c807d36c212bdb33ea8587f7569562a84df5465b1" +dependencies = [ + "libc", +] + +[[package]] +name = "simd-adler32" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d66dc143e6b11c1eddc06d5c423cfc97062865baf299914ab64caa38182078fe" + +[[package]] +name = "simplecss" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a11be7c62927d9427e9f40f3444d5499d868648e2edbc4e2116de69e7ec0e89d" +dependencies = [ + "log", +] + +[[package]] +name = "siphasher" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38b58827f4464d87d377d175e90bf58eb00fd8716ff0a62f80356b5e61555d0d" + +[[package]] +name = "slab" +version = "0.4.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f92a496fb766b417c996b9c5e57daf2f7ad3b0bebe1ccfca4856390e3d3bb67" +dependencies = [ + "autocfg", +] + +[[package]] +name = "slotmap" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dbff4acf519f630b3a3ddcfaea6c06b42174d9a44bc70c620e9ed1649d58b82a" +dependencies = [ + "version_check", +] + +[[package]] +name = "smallvec" +version = "1.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" + +[[package]] +name = "smithay-client-toolkit" +version = "0.19.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3457dea1f0eb631b4034d61d4d8c32074caa6cd1ab2d59f2327bd8461e2c0016" +dependencies = [ + "bitflags 2.6.0", + "calloop", + "calloop-wayland-source", + "cursor-icon", + "libc", + "log", + "memmap2", + "rustix", + "thiserror", + "wayland-backend", + "wayland-client", + "wayland-csd-frame", + "wayland-cursor", + "wayland-protocols", + "wayland-protocols-wlr", + "wayland-scanner", + "xkeysym", +] + +[[package]] +name = "smithay-clipboard" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc8216eec463674a0e90f29e0ae41a4db573ec5b56b1c6c1c71615d249b6d846" +dependencies = [ + "libc", + "smithay-client-toolkit", + "wayland-backend", +] + +[[package]] +name = "smol_str" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd538fb6910ac1099850255cf94a94df6551fbdd602454387d0adb2d1ca6dead" +dependencies = [ + "serde", +] + +[[package]] +name = "space_robotics_bench" +version = "0.0.1" +dependencies = [ + "display_json", + "fast_poisson", + "figment", + "paste", + "pyo3", + "rand", + "rand_xoshiro", + "rayon", + "serde", + "serde_json", + "serde_yaml", + "thiserror", + "toml", +] + +[[package]] +name = "space_robotics_bench_gui" +version = "0.0.1" +dependencies = [ + "chrono", + "const_format", + "display_json", + "eframe", + "egui", + "egui_commonmark", + "egui_extras", + "home", + "image", + "itertools 0.13.0", + "nix", + "paste", + "r2r", + "serde", + "serde_json", + "space_robotics_bench", + "subprocess", + "sysinfo", + "tracing", + "tracing-subscriber", + "typed-builder", +] + +[[package]] +name = "space_robotics_bench_py" +version = "0.0.1" +dependencies = [ + "const_format", + "pyo3", + "space_robotics_bench", +] + +[[package]] +name = "spin" +version = "0.9.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" + +[[package]] +name = "stable_deref_trait" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" + +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + +[[package]] +name = "strict-num" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6637bab7722d379c8b41ba849228d680cc12d0a45ba1fa2b48f2a30577a06731" +dependencies = [ + "float-cmp", +] + +[[package]] +name = "subprocess" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c2e86926081dda636c546d8c5e641661049d7562a68f5488be4a1f7f66f6086" +dependencies = [ + "libc", + "winapi", +] + +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + +[[package]] +name = "svgtypes" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e44e288cd960318917cbd540340968b90becc8bc81f171345d706e7a89d9d70" +dependencies = [ + "kurbo", + "siphasher", +] + +[[package]] +name = "syn" +version = "1.0.109" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "2.0.82" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83540f837a8afc019423a8edb95b52a8effe46957ee402287f4292fae35be021" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syntect" +version = "5.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "874dcfa363995604333cf947ae9f751ca3af4522c60886774c4963943b4746b1" +dependencies = [ + "bincode", + "bitflags 1.3.2", + "fancy-regex", + "flate2", + "fnv", + "once_cell", + "plist", + "regex-syntax", + "serde", + "serde_derive", + "serde_json", + "thiserror", + "walkdir", + "yaml-rust", +] + +[[package]] +name = "sysinfo" +version = "0.32.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3b5ae3f4f7d64646c46c4cae4e3f01d1c5d255c7406fdd7c7f999a94e488791" +dependencies = [ + "core-foundation-sys", + "libc", + "memchr", + "ntapi", + "rayon", + "windows 0.57.0", +] + +[[package]] +name = "target-lexicon" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" + +[[package]] +name = "tempfile" +version = "3.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0f2c9fc62d0beef6951ccffd757e241266a2c833136efbe35af6cd2567dca5b" +dependencies = [ + "cfg-if", + "fastrand", + "once_cell", + "rustix", + "windows-sys 0.59.0", +] + +[[package]] +name = "thiserror" +version = "1.0.64" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d50af8abc119fb8bb6dbabcfa89656f46f84aa0ac7688088608076ad2b459a84" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.64" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08904e7672f5eb876eaaf87e0ce17857500934f4981c4a0ab2b4aa98baac7fc3" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "thread_local" +version = "1.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b9ef9bad013ada3808854ceac7b46812a6465ba368859a37e2100283d2d719c" +dependencies = [ + "cfg-if", + "once_cell", +] + +[[package]] +name = "time" +version = "0.3.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5dfd88e563464686c916c7e46e623e520ddc6d79fa6641390f2e3fa86e83e885" +dependencies = [ + "deranged", + "itoa", + "num-conv", + "powerfmt", + "serde", + "time-core", + "time-macros", +] + +[[package]] +name = "time-core" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef927ca75afb808a4d64dd374f00a2adf8d0fcff8e7b184af886c3c87ec4a3f3" + +[[package]] +name = "time-macros" +version = "0.2.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f252a68540fde3a3877aeea552b832b40ab9a69e318efd078774a01ddee1ccf" +dependencies = [ + "num-conv", + "time-core", +] + +[[package]] +name = "tiny-skia" +version = "0.11.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83d13394d44dae3207b52a326c0c85a8bf87f1541f23b0d143811088497b09ab" +dependencies = [ + "arrayref", + "arrayvec", + "bytemuck", + "cfg-if", + "log", + "png", + "tiny-skia-path", +] + +[[package]] +name = "tiny-skia-path" +version = "0.11.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c9e7fc0c2e86a30b117d0462aa261b72b7a99b7ebd7deb3a14ceda95c5bdc93" +dependencies = [ + "arrayref", + "bytemuck", + "strict-num", +] + +[[package]] +name = "tinyvec" +version = "1.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "445e881f4f6d382d5f27c034e25eb92edd7c784ceab92a0937db7f2e9471b938" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + +[[package]] +name = "toml" +version = "0.8.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1ed1f98e3fdc28d6d910e6737ae6ab1a93bf1985935a1193e68f93eeb68d24e" +dependencies = [ + "serde", + "serde_spanned", + "toml_datetime", + "toml_edit", +] + +[[package]] +name = "toml_datetime" +version = "0.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0dd7358ecb8fc2f8d014bf86f6f638ce72ba252a2c3a2572f2a795f1d23efb41" +dependencies = [ + "serde", +] + +[[package]] +name = "toml_edit" +version = "0.22.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ae48d6208a266e853d946088ed816055e556cc6028c5e8e2b84d9fa5dd7c7f5" +dependencies = [ + "indexmap", + "serde", + "serde_spanned", + "toml_datetime", + "winnow", +] + +[[package]] +name = "tracing" +version = "0.1.40" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3523ab5a71916ccf420eebdf5521fcef02141234bbc0b8a49f2fdc4544364ef" +dependencies = [ + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34704c8d6ebcbc939824180af020566b01a7c01f80641264eba0999f6c2b6be7" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "tracing-core" +version = "0.1.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c06d3da6113f116aaee68e4d601191614c9053067f9ab7f6edbcb161237daa54" +dependencies = [ + "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad0f048c97dbd9faa9b7df56362b8ebcaa52adb06b498c050d2f4e32f90a7a8b" +dependencies = [ + "nu-ansi-term", + "sharded-slab", + "smallvec", + "thread_local", + "tracing-core", + "tracing-log", +] + +[[package]] +name = "ttf-parser" +version = "0.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5902c5d130972a0000f60860bfbf46f7ca3db5391eddfedd1b8728bd9dc96c0e" + +[[package]] +name = "typed-builder" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e14ed59dc8b7b26cacb2a92bad2e8b1f098806063898ab42a3bd121d7d45e75" +dependencies = [ + "typed-builder-macro", +] + +[[package]] +name = "typed-builder-macro" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "560b82d656506509d43abe30e0ba64c56b1953ab3d4fe7ba5902747a7a3cedd5" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "typenum" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" + +[[package]] +name = "uds_windows" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89daebc3e6fd160ac4aa9fc8b3bf71e1f74fbf92367ae71fb83a037e8bf164b9" +dependencies = [ + "memoffset", + "tempfile", + "winapi", +] + +[[package]] +name = "uncased" +version = "0.9.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1b88fcfe09e89d3866a5c11019378088af2d24c3fbd4f0543f96b479ec90697" +dependencies = [ + "version_check", +] + +[[package]] +name = "unicase" +version = "2.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e51b68083f157f853b6379db119d1c1be0e6e4dec98101079dec41f6f5cf6df" + +[[package]] +name = "unicode-bidi" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ab17db44d7388991a428b2ee655ce0c212e862eff1768a455c58f9aad6e7893" + +[[package]] +name = "unicode-ident" +version = "1.0.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91b56cd4cadaeb79bbf1a5645f6b4f8dc5bde8834ad5894a8db35fda9efa1fe" + +[[package]] +name = "unicode-normalization" +version = "0.1.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5033c97c4262335cded6d6fc3e5c18ab755e1a3dc96376350f3d8e9f009ad956" +dependencies = [ + "tinyvec", +] + +[[package]] +name = "unicode-segmentation" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" + +[[package]] +name = "unicode-xid" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" + +[[package]] +name = "unindent" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7de7d73e1754487cb58364ee906a499937a0dfabd86bcb980fa99ec8c8fa2ce" + +[[package]] +name = "unsafe-libyaml" +version = "0.2.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861" + +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + +[[package]] +name = "ureq" +version = "2.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b74fc6b57825be3373f7054754755f03ac3a8f5d70015ccad699ba2029956f4a" +dependencies = [ + "base64 0.22.1", + "flate2", + "log", + "once_cell", + "rustls", + "rustls-pki-types", + "url", + "webpki-roots", +] + +[[package]] +name = "url" +version = "2.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22784dbdf76fdde8af1aeda5622b546b422b6fc585325248a2bf9f5e41e94d6c" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", +] + +[[package]] +name = "usvg" +version = "0.37.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38b0a51b72ab80ca511d126b77feeeb4fb1e972764653e61feac30adc161a756" +dependencies = [ + "base64 0.21.7", + "log", + "pico-args", + "usvg-parser", + "usvg-tree", + "xmlwriter", +] + +[[package]] +name = "usvg-parser" +version = "0.37.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9bd4e3c291f45d152929a31f0f6c819245e2921bfd01e7bd91201a9af39a2bdc" +dependencies = [ + "data-url", + "flate2", + "imagesize", + "kurbo", + "log", + "roxmltree", + "simplecss", + "siphasher", + "svgtypes", + "usvg-tree", +] + +[[package]] +name = "usvg-tree" +version = "0.37.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ee3d202ebdb97a6215604b8f5b4d6ef9024efd623cf2e373a6416ba976ec7d3" +dependencies = [ + "rctree", + "strict-num", + "svgtypes", + "tiny-skia-path", +] + +[[package]] +name = "uuid" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8c5f0a0af699448548ad1a2fbf920fb4bee257eae39953ba95cb84891a0446a" +dependencies = [ + "getrandom", + "serde", +] + +[[package]] +name = "valuable" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + +[[package]] +name = "wasi" +version = "0.11.0+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" + +[[package]] +name = "wasm-bindgen" +version = "0.2.95" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "128d1e363af62632b8eb57219c8fd7877144af57558fb2ef0368d0087bddeb2e" +dependencies = [ + "cfg-if", + "once_cell", + "wasm-bindgen-macro", +] + +[[package]] +name = "wasm-bindgen-backend" +version = "0.2.95" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb6dd4d3ca0ddffd1dd1c9c04f94b868c37ff5fac97c30b97cff2d74fce3a358" +dependencies = [ + "bumpalo", + "log", + "once_cell", + "proc-macro2", + "quote", + "syn 2.0.82", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-futures" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc7ec4f8827a71586374db3e87abdb5a2bb3a15afed140221307c3ec06b1f63b" +dependencies = [ + "cfg-if", + "js-sys", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.95" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e79384be7f8f5a9dd5d7167216f022090cf1f9ec128e6e6a482a2cb5c5422c56" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.95" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26c6ab57572f7a24a4985830b120de1594465e5d500f24afe89e16b4e833ef68" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", + "wasm-bindgen-backend", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.95" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65fc09f10666a9f147042251e0dda9c18f166ff7de300607007e96bdebc1068d" + +[[package]] +name = "wayland-backend" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "056535ced7a150d45159d3a8dc30f91a2e2d588ca0b23f70e56033622b8016f6" +dependencies = [ + "cc", + "downcast-rs", + "rustix", + "scoped-tls", + "smallvec", + "wayland-sys", +] + +[[package]] +name = "wayland-client" +version = "0.31.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3f45d1222915ef1fd2057220c1d9d9624b7654443ea35c3877f7a52bd0a5a2d" +dependencies = [ + "bitflags 2.6.0", + "rustix", + "wayland-backend", + "wayland-scanner", +] + +[[package]] +name = "wayland-csd-frame" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625c5029dbd43d25e6aa9615e88b829a5cad13b2819c4ae129fdbb7c31ab4c7e" +dependencies = [ + "bitflags 2.6.0", + "cursor-icon", + "wayland-backend", +] + +[[package]] +name = "wayland-cursor" +version = "0.31.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a94697e66e76c85923b0d28a0c251e8f0666f58fc47d316c0f4da6da75d37cb" +dependencies = [ + "rustix", + "wayland-client", + "xcursor", +] + +[[package]] +name = "wayland-protocols" +version = "0.32.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b5755d77ae9040bb872a25026555ce4cb0ae75fd923e90d25fba07d81057de0" +dependencies = [ + "bitflags 2.6.0", + "wayland-backend", + "wayland-client", + "wayland-scanner", +] + +[[package]] +name = "wayland-protocols-plasma" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a0a41a6875e585172495f7a96dfa42ca7e0213868f4f15c313f7c33221a7eff" +dependencies = [ + "bitflags 2.6.0", + "wayland-backend", + "wayland-client", + "wayland-protocols", + "wayland-scanner", +] + +[[package]] +name = "wayland-protocols-wlr" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dad87b5fd1b1d3ca2f792df8f686a2a11e3fe1077b71096f7a175ab699f89109" +dependencies = [ + "bitflags 2.6.0", + "wayland-backend", + "wayland-client", + "wayland-protocols", + "wayland-scanner", +] + +[[package]] +name = "wayland-scanner" +version = "0.31.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "597f2001b2e5fc1121e3d5b9791d3e78f05ba6bfa4641053846248e3a13661c3" +dependencies = [ + "proc-macro2", + "quick-xml 0.36.2", + "quote", +] + +[[package]] +name = "wayland-sys" +version = "0.31.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "efa8ac0d8e8ed3e3b5c9fc92c7881406a268e11555abe36493efabe649a29e09" +dependencies = [ + "dlib", + "log", + "once_cell", + "pkg-config", +] + +[[package]] +name = "web-sys" +version = "0.3.72" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6488b90108c040df0fe62fa815cbdee25124641df01814dd7282749234c6112" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "webbrowser" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e5f07fb9bc8de2ddfe6b24a71a75430673fd679e568c48b52716cef1cfae923" +dependencies = [ + "block2", + "core-foundation 0.10.0", + "home", + "jni", + "log", + "ndk-context", + "objc2", + "objc2-foundation", + "url", + "web-sys", +] + +[[package]] +name = "webpki-roots" +version = "0.26.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841c67bff177718f1d4dfefde8d8f0e78f9b6589319ba88312f567fc5841a958" +dependencies = [ + "rustls-pki-types", +] + +[[package]] +name = "weezl" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53a85b86a771b1c87058196170769dd264f66c0782acf1ae6cc51bfd64b39082" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "widestring" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7219d36b6eac893fa81e84ebe06485e7dcbb616177469b142df14f1f4deb1311" + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-util" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" +dependencies = [ + "windows-sys 0.59.0", +] + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + +[[package]] +name = "windows" +version = "0.57.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "12342cb4d8e3b046f3d80effd474a7a02447231330ef77d71daa6fbc40681143" +dependencies = [ + "windows-core 0.57.0", + "windows-targets 0.52.6", +] + +[[package]] +name = "windows" +version = "0.58.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd04d41d93c4992d421894c18c8b43496aa748dd4c081bac0dc93eb0489272b6" +dependencies = [ + "windows-core 0.58.0", + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-core" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ab640c8d7e35bf8ba19b884ba838ceb4fba93a4e8c65a9059d08afcfc683d9" +dependencies = [ + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-core" +version = "0.57.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2ed2439a290666cd67ecce2b0ffaad89c2a56b976b736e6ece670297897832d" +dependencies = [ + "windows-implement 0.57.0", + "windows-interface 0.57.0", + "windows-result 0.1.2", + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-core" +version = "0.58.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ba6d44ec8c2591c134257ce647b7ea6b20335bf6379a27dac5f1641fcf59f99" +dependencies = [ + "windows-implement 0.58.0", + "windows-interface 0.58.0", + "windows-result 0.2.0", + "windows-strings", + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-implement" +version = "0.57.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9107ddc059d5b6fbfbffdfa7a7fe3e22a226def0b2608f72e9d552763d3e1ad7" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "windows-implement" +version = "0.58.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2bbd5b46c938e506ecbce286b6628a02171d56153ba733b6c741fc627ec9579b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "windows-interface" +version = "0.57.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29bee4b38ea3cde66011baa44dba677c432a78593e202392d1e9070cf2a7fca7" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "windows-interface" +version = "0.58.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053c4c462dc91d3b1504c6fe5a726dd15e216ba718e84a0e46a88fbe5ded3515" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "windows-result" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e383302e8ec8515204254685643de10811af0ed97ea37210dc26fb0032647f8" +dependencies = [ + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-result" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d1043d8214f791817bab27572aaa8af63732e11bf84aa21a45a78d6c317ae0e" +dependencies = [ + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-strings" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cd9b125c486025df0eabcb585e62173c6c9eddcec5d117d3b6e8c30e2ee4d10" +dependencies = [ + "windows-result 0.2.0", + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-sys" +version = "0.45.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75283be5efb2831d37ea142365f009c02ec203cd29a3ebecbc093d52315b66d0" +dependencies = [ + "windows-targets 0.42.2", +] + +[[package]] +name = "windows-sys" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +dependencies = [ + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-targets" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e5180c00cd44c9b1c88adb3693291f1cd93605ded80c250a75d472756b4d071" +dependencies = [ + "windows_aarch64_gnullvm 0.42.2", + "windows_aarch64_msvc 0.42.2", + "windows_i686_gnu 0.42.2", + "windows_i686_msvc 0.42.2", + "windows_x86_64_gnu 0.42.2", + "windows_x86_64_gnullvm 0.42.2", + "windows_x86_64_msvc 0.42.2", +] + +[[package]] +name = "windows-targets" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" +dependencies = [ + "windows_aarch64_gnullvm 0.48.5", + "windows_aarch64_msvc 0.48.5", + "windows_i686_gnu 0.48.5", + "windows_i686_msvc 0.48.5", + "windows_x86_64_gnu 0.48.5", + "windows_x86_64_gnullvm 0.48.5", + "windows_x86_64_msvc 0.48.5", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm 0.52.6", + "windows_aarch64_msvc 0.52.6", + "windows_i686_gnu 0.52.6", + "windows_i686_gnullvm", + "windows_i686_msvc 0.52.6", + "windows_x86_64_gnu 0.52.6", + "windows_x86_64_gnullvm 0.52.6", + "windows_x86_64_msvc 0.52.6", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8" + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f" + +[[package]] +name = "windows_i686_gnu" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060" + +[[package]] +name = "windows_i686_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "winit" +version = "0.30.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0be9e76a1f1077e04a411f0b989cbd3c93339e1771cb41e71ac4aee95bfd2c67" +dependencies = [ + "ahash", + "android-activity", + "atomic-waker", + "bitflags 2.6.0", + "block2", + "bytemuck", + "calloop", + "cfg_aliases", + "concurrent-queue", + "core-foundation 0.9.4", + "core-graphics", + "cursor-icon", + "dpi", + "js-sys", + "libc", + "memmap2", + "ndk", + "objc2", + "objc2-app-kit", + "objc2-foundation", + "objc2-ui-kit", + "orbclient", + "percent-encoding", + "pin-project", + "raw-window-handle", + "redox_syscall 0.4.1", + "rustix", + "smithay-client-toolkit", + "smol_str", + "tracing", + "unicode-segmentation", + "wasm-bindgen", + "wasm-bindgen-futures", + "wayland-backend", + "wayland-client", + "wayland-protocols", + "wayland-protocols-plasma", + "web-sys", + "web-time", + "windows-sys 0.52.0", + "x11-dl", + "x11rb", + "xkbcommon-dl", +] + +[[package]] +name = "winnow" +version = "0.6.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36c1fec1a2bb5866f07c25f68c26e565c4c200aebb96d7e55710c19d3e8ac49b" +dependencies = [ + "memchr", +] + +[[package]] +name = "x11-dl" +version = "2.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38735924fedd5314a6e548792904ed8c6de6636285cb9fec04d5b1db85c1516f" +dependencies = [ + "libc", + "once_cell", + "pkg-config", +] + +[[package]] +name = "x11rb" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d91ffca73ee7f68ce055750bf9f6eca0780b8c85eff9bc046a3b0da41755e12" +dependencies = [ + "as-raw-xcb-connection", + "gethostname", + "libc", + "libloading", + "once_cell", + "rustix", + "x11rb-protocol", +] + +[[package]] +name = "x11rb-protocol" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec107c4503ea0b4a98ef47356329af139c0a4f7750e621cf2973cd3385ebcb3d" + +[[package]] +name = "xcursor" +version = "0.3.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ef33da6b1660b4ddbfb3aef0ade110c8b8a781a3b6382fa5f2b5b040fd55f61" + +[[package]] +name = "xdg-home" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec1cdab258fb55c0da61328dc52c8764709b249011b2cad0454c72f0bf10a1f6" +dependencies = [ + "libc", + "windows-sys 0.59.0", +] + +[[package]] +name = "xkbcommon-dl" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d039de8032a9a8856a6be89cea3e5d12fdd82306ab7c94d74e6deab2460651c5" +dependencies = [ + "bitflags 2.6.0", + "dlib", + "log", + "once_cell", + "xkeysym", +] + +[[package]] +name = "xkeysym" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9cc00251562a284751c9973bace760d86c0276c471b4be569fe6b068ee97a56" + +[[package]] +name = "xml-rs" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af4e2e2f7cba5a093896c1e150fbfe177d1883e7448200efb81d40b9d339ef26" + +[[package]] +name = "xmlwriter" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec7a2a501ed189703dba8b08142f057e887dfc4b2cc4db2d343ac6376ba3e0b9" + +[[package]] +name = "yaml-rust" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56c1936c4cc7a1c9ab21a1ebb602eb942ba868cbd44a99cb7cdc5892335e1c85" +dependencies = [ + "linked-hash-map", +] + +[[package]] +name = "yansi" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfe53a6657fd280eaa890a3bc59152892ffa3e30101319d168b781ed6529b049" + +[[package]] +name = "zbus" +version = "4.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb97012beadd29e654708a0fdb4c84bc046f537aecfde2c3ee0a9e4b4d48c725" +dependencies = [ + "async-broadcast", + "async-executor", + "async-fs", + "async-io", + "async-lock", + "async-process", + "async-recursion", + "async-task", + "async-trait", + "blocking", + "enumflags2", + "event-listener", + "futures-core", + "futures-sink", + "futures-util", + "hex", + "nix", + "ordered-stream", + "rand", + "serde", + "serde_repr", + "sha1", + "static_assertions", + "tracing", + "uds_windows", + "windows-sys 0.52.0", + "xdg-home", + "zbus_macros", + "zbus_names", + "zvariant", +] + +[[package]] +name = "zbus-lockstep" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca2c5dceb099bddaade154055c926bb8ae507a18756ba1d8963fd7b51d8ed1d" +dependencies = [ + "zbus_xml", + "zvariant", +] + +[[package]] +name = "zbus-lockstep-macros" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709ab20fc57cb22af85be7b360239563209258430bccf38d8b979c5a2ae3ecce" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", + "zbus-lockstep", + "zbus_xml", + "zvariant", +] + +[[package]] +name = "zbus_macros" +version = "4.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "267db9407081e90bbfa46d841d3cbc60f59c0351838c4bc65199ecd79ab1983e" +dependencies = [ + "proc-macro-crate", + "proc-macro2", + "quote", + "syn 2.0.82", + "zvariant_utils", +] + +[[package]] +name = "zbus_names" +version = "3.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b9b1fef7d021261cc16cba64c351d291b715febe0fa10dc3a443ac5a5022e6c" +dependencies = [ + "serde", + "static_assertions", + "zvariant", +] + +[[package]] +name = "zbus_xml" +version = "4.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab3f374552b954f6abb4bd6ce979e6c9b38fb9d0cd7cc68a7d796e70c9f3a233" +dependencies = [ + "quick-xml 0.30.0", + "serde", + "static_assertions", + "zbus_names", + "zvariant", +] + +[[package]] +name = "zerocopy" +version = "0.7.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0" +dependencies = [ + "byteorder", + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.7.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] + +[[package]] +name = "zeroize" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ced3678a2879b30306d323f4542626697a464a97c0a07c9aebf7ebca65cd4dde" + +[[package]] +name = "zune-core" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f423a2c17029964870cfaabb1f13dfab7d092a62a29a89264f4d36990ca414a" + +[[package]] +name = "zune-jpeg" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16099418600b4d8f028622f73ff6e3deaabdff330fb9a2a131dea781ee8b0768" +dependencies = [ + "zune-core", +] + +[[package]] +name = "zvariant" +version = "4.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2084290ab9a1c471c38fc524945837734fbf124487e105daec2bb57fd48c81fe" +dependencies = [ + "endi", + "enumflags2", + "serde", + "static_assertions", + "zvariant_derive", +] + +[[package]] +name = "zvariant_derive" +version = "4.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73e2ba546bda683a90652bac4a279bc146adad1386f25379cf73200d2002c449" +dependencies = [ + "proc-macro-crate", + "proc-macro2", + "quote", + "syn 2.0.82", + "zvariant_utils", +] + +[[package]] +name = "zvariant_utils" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c51bcff7cc3dbb5055396bcf774748c3dab426b4b8659046963523cee4808340" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.82", +] diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..910c153 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,57 @@ +[workspace.package] +authors = ["Andrej Orsula "] +description = "Comprehensive benchmark for space robotics" +categories = ["science::robotics"] +keywords = ["benchmark", "robotics", "simulation", "space"] +readme = "README.md" +license = "MIT OR Apache-2.0" +homepage = "https://github.com/AndrejOrsula/space_robotics_bench" +repository = "https://github.com/AndrejOrsula/space_robotics_bench" +documentation = "https://andrejorsula.github.io/space_robotics_bench" +edition = "2021" +rust-version = "1.80" +version = "0.0.1" +publish = false + +[workspace] +resolver = "2" +members = [ + "crates/space_robotics_bench", + "crates/space_robotics_bench_py", + "crates/space_robotics_bench_gui", +] +default-members = [ + "crates/space_robotics_bench", + "crates/space_robotics_bench_py", +] + +[workspace.dependencies] +space_robotics_bench = { path = "crates/space_robotics_bench", version = "0.0.1" } + +chrono = { version = "0.4" } +const_format = { version = "0.2", features = ["more_str_macros"] } +display_json = { version = "0.2" } +fast_poisson = { version = "1.0", features = ["single_precision"] } +figment = { version = "0.10", features = ["env", "test"] } +home = { version = "0.5.9" } +image = { version = "0.25", default-features = false, features = [ + "jpeg", + "png", +] } +itertools = { version = "0.13" } +nix = { version = "0.29", features = ["signal"] } +paste = { version = "1.0" } +pyo3 = { version = "0.22", features = ["abi3-py310", "auto-initialize"] } +rand = { version = "0.8" } +rand_xoshiro = { version = "0.6" } +rayon = { version = "1.5" } +serde = { version = "1", features = ["derive"] } +serde_json = { version = "1" } +serde_yaml = { version = "0.9" } +subprocess = { version = "0.2" } +sysinfo = { version = "0.32" } +thiserror = { version = "1" } +toml = { version = "0.8" } +tracing = { version = "0.1" } +tracing-subscriber = { version = "0.3" } +typed-builder = { version = "0.20" } diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..c1f29de --- /dev/null +++ b/Dockerfile @@ -0,0 +1,223 @@ +## Base +ARG BASE_IMAGE_NAME="ubuntu" +ARG BASE_IMAGE_TAG="jammy" + +## Isaac Sim +## Label as isaac-sim for copying into the final image +ARG ISAAC_SIM_IMAGE_NAME="nvcr.io/nvidia/isaac-sim" +ARG ISAAC_SIM_IMAGE_TAG="4.2.0" +FROM ${ISAAC_SIM_IMAGE_NAME}:${ISAAC_SIM_IMAGE_TAG} AS isaac-sim + +## Continue with the base image +FROM ${BASE_IMAGE_NAME}:${BASE_IMAGE_TAG} + +## Use bash as the default shell +SHELL ["/bin/bash", "-o", "pipefail", "-c"] + +## Create a barebones entrypoint that is conditionally updated throughout the Dockerfile +RUN echo "#!/usr/bin/env bash" >> /entrypoint.bash && \ + chmod +x /entrypoint.bash + +## Copy Isaac Sim into the base image +ARG ISAAC_SIM_PATH="/root/isaac-sim" +ENV ISAAC_SIM_PYTHON="${ISAAC_SIM_PATH}/python.sh" +COPY --from=isaac-sim /isaac-sim "${ISAAC_SIM_PATH}" +COPY --from=isaac-sim /root/.nvidia-omniverse/config /root/.nvidia-omniverse/config +COPY --from=isaac-sim /etc/vulkan/icd.d/nvidia_icd.json /etc/vulkan/icd.d/nvidia_icd.json +RUN ISAAC_SIM_VERSION="$(cut -d'-' -f1 < "${ISAAC_SIM_PATH}/VERSION")" && \ + echo -e "\n# Isaac Sim ${ISAAC_SIM_VERSION}" >> /entrypoint.bash && \ + echo "export ISAAC_SIM_PATH=\"${ISAAC_SIM_PATH}\"" >> /entrypoint.bash && \ + echo "export OMNI_KIT_ALLOW_ROOT=\"1\"" >> /entrypoint.bash +## Fix cosmetic issues in `isaac-sim/setup_python_env.sh` that append nonsense paths to `PYTHONPATH` and `LD_LIBRARY_PATH` +# hadolint ignore=SC2016 +RUN sed -i 's|$SCRIPT_DIR/../../../$LD_LIBRARY_PATH:||' "${ISAAC_SIM_PATH}/setup_python_env.sh" && \ + sed -i 's|$SCRIPT_DIR/../../../$PYTHONPATH:||' "${ISAAC_SIM_PATH}/setup_python_env.sh" + +## Optimization: Build Python to improve the runtime performance of training +ARG PYTHON_VERSION="3.10.14" +ENV PYTHONEXE="/usr/local/bin/python${PYTHON_VERSION%%.*}" +# hadolint ignore=DL3003,DL3008 +RUN PYTHON_DL_PATH="/tmp/Python-${PYTHON_VERSION}.tar.xz" && \ + PYTHON_SRC_DIR="/tmp/python${PYTHON_VERSION}" && \ + apt-get update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -yq --no-install-recommends \ + build-essential \ + ca-certificates \ + curl \ + libbz2-dev \ + libdb4o-cil-dev \ + libgdm-dev \ + libhidapi-dev \ + liblzma-dev \ + libncurses5-dev \ + libpcap-dev \ + libreadline-dev \ + libsqlite3-dev \ + libssl-dev \ + libtk8.6 \ + lzma \ + xz-utils && \ + rm -rf /var/lib/apt/lists/* && \ + curl --proto "=https" --tlsv1.2 -sSfL "https://www.python.org/ftp/python/${PYTHON_VERSION}/Python-${PYTHON_VERSION}.tar.xz" -o "${PYTHON_DL_PATH}" && \ + mkdir -p "${PYTHON_SRC_DIR}" && \ + tar xf "${PYTHON_DL_PATH}" -C "${PYTHON_SRC_DIR}" --strip-components=1 && \ + rm "${PYTHON_DL_PATH}" && \ + cd "${PYTHON_SRC_DIR}" && \ + "${PYTHON_SRC_DIR}/configure" --enable-shared --enable-optimizations --with-lto --prefix="/usr/local" && \ + make -j "$(nproc)" && \ + make install && \ + cd - && \ + rm -rf "${PYTHON_SRC_DIR}" +## Fix `PYTHONEXE` by disabling the append of "isaac-sim/kit/kernel/plugins" to `LD_LIBRARY_PATH` inside `isaac-sim/setup_python_env.sh` +# hadolint ignore=SC2016 +RUN sed -i 's|$SCRIPT_DIR/kit/kernel/plugins:||' "${ISAAC_SIM_PATH}/setup_python_env.sh" + +## Install system dependencies +# hadolint ignore=DL3008 +RUN apt-get update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -yq --no-install-recommends \ + # Common + build-essential \ + ca-certificates \ + cmake \ + curl \ + git \ + mold \ + xz-utils \ + # Isaac Sim + libgl1 \ + libglu1 \ + libxt-dev \ + # Blender + libegl1 \ + # egui + libxkbcommon-x11-0 \ + # r2r + clang \ + # Video recording/processing + ffmpeg && \ + rm -rf /var/lib/apt/lists/* + +## Upgrade pip +RUN "${ISAAC_SIM_PYTHON}" -m pip install --no-input --no-cache-dir --upgrade pip + +## Install Rust +ARG RUST_VERSION="1.80" +RUN echo -e "\n# Rust ${RUST_VERSION}" >> /entrypoint.bash && \ + echo "export PATH=\"${HOME}/.cargo/bin\${PATH:+:\${PATH}}\"" >> /entrypoint.bash && \ + echo "export CARGO_TARGET_DIR=\"${HOME}/.cargo/target\"" >> /entrypoint.bash && \ + echo "export CARGO_TARGET_X86_64_UNKNOWN_LINUX_GNU_RUSTFLAGS=\"-Clink-arg=-fuse-ld=mold -Ctarget-cpu=native\"" >> /entrypoint.bash && \ + echo -e "\n# PyO3" >> /entrypoint.bash && \ + echo "export PYO3_PYTHON=\"${ISAAC_SIM_PYTHON}\"" >> /entrypoint.bash && \ + curl --proto "=https" --tlsv1.2 -sSfL "https://sh.rustup.rs" | sh -s -- --no-modify-path -y --default-toolchain "${RUST_VERSION}" --profile default + +## Install ROS +ARG ROS_DISTRO="humble" +# hadolint ignore=SC1091,DL3008 +RUN curl --proto "=https" --tlsv1.2 -sSfL "https://raw.githubusercontent.com/ros/rosdistro/master/ros.key" -o /usr/share/keyrings/ros-archive-keyring.gpg && \ + echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/ros-archive-keyring.gpg] http://packages.ros.org/ros2/ubuntu $(source /etc/os-release && echo "${UBUNTU_CODENAME}") main" > /etc/apt/sources.list.d/ros2.list && \ + apt-get update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -yq --no-install-recommends \ + ros-dev-tools \ + "ros-${ROS_DISTRO}-ros-base" \ + "ros-${ROS_DISTRO}-rmw-cyclonedds-cpp" && \ + rm -rf /var/lib/apt/lists/* && \ + "${ISAAC_SIM_PYTHON}" -m pip install --no-input --no-cache-dir catkin_pkg && \ + rosdep init --rosdistro "${ROS_DISTRO}" && \ + echo -e "\n# ROS ${ROS_DISTRO^}" >> /entrypoint.bash && \ + echo "source \"/opt/ros/${ROS_DISTRO}/setup.bash\" --" >> /entrypoint.bash + +## Install Blender +ARG BLENDER_PATH="/root/blender" +ARG BLENDER_VERSION="4.2.2" +# hadolint ignore=SC2016 +RUN echo -e "\n# Blender ${BLENDER_VERSION}" >> /entrypoint.bash && \ + echo "export PATH=\"${BLENDER_PATH}\${PATH:+:\${PATH}}\"" >> /entrypoint.bash && \ + curl --proto "=https" --tlsv1.2 -sSfL "https://download.blender.org/release/Blender${BLENDER_VERSION%.*}/blender-${BLENDER_VERSION}-linux-x64.tar.xz" -o "/tmp/blender_${BLENDER_VERSION}.tar.xz" && \ + mkdir -p "${BLENDER_PATH}" && \ + tar xf "/tmp/blender_${BLENDER_VERSION}.tar.xz" -C "${BLENDER_PATH}" --strip-components=1 && \ + rm "/tmp/blender_${BLENDER_VERSION}.tar.xz" + +## Install Isaac Lab +ARG ISAACLAB_PATH="/root/isaaclab" +ARG ISAACLAB_REMOTE="https://github.com/isaac-sim/IsaacLab.git" +ARG ISAACLAB_BRANCH="main" +ARG ISAACLAB_COMMIT_SHA="0ef582badf6f257bb3c320c63d5d6d899604a138" # Oct 5, 2024 +# hadolint ignore=SC2044 +RUN echo -e "\n# Isaac Lab ${ISAACLAB_COMMIT_SHA}" >> /entrypoint.bash && \ + echo "export ISAACLAB_PATH=\"${ISAACLAB_PATH}\"" >> /entrypoint.bash && \ + git clone "${ISAACLAB_REMOTE}" "${ISAACLAB_PATH}" --branch "${ISAACLAB_BRANCH}" && \ + git -C "${ISAACLAB_PATH}" reset --hard "${ISAACLAB_COMMIT_SHA}" && \ + for extension in $(find -L "${ISAACLAB_PATH}/source/extensions" -mindepth 1 -maxdepth 1 -type d); do \ + if [ -f "${extension}/setup.py" ]; then \ + "${ISAAC_SIM_PYTHON}" -m pip install --no-input --no-cache-dir --editable "${extension}" ; \ + fi ; \ + done + +## Finalize the entrypoint +# hadolint ignore=SC2016 +RUN echo -e "\n# Execute command" >> /entrypoint.bash && \ + echo -en 'exec "${@}"\n' >> /entrypoint.bash && \ + sed -i '$a source /entrypoint.bash --' ~/.bashrc +ENTRYPOINT ["/entrypoint.bash"] + +## Define the workspace of the project +ARG WS="/root/ws" +WORKDIR "${WS}" + +## [Optional] Install Python dependencies in advance to cache the layers (speeds up rebuilds) +COPY ./pyproject.toml "${WS}/" +# hadolint ignore=SC2046 +RUN "${ISAAC_SIM_PYTHON}" -m pip install --no-input --no-cache-dir toml==0.10.2 && \ + "${ISAAC_SIM_PYTHON}" -m pip install --no-input --no-cache-dir $("${ISAAC_SIM_PYTHON}" -c "import toml, itertools; data = toml.load(open('${WS}/pyproject.toml')); project_name = data['project']['name']; deps = [dep for dep in data['project'].get('dependencies', []) if project_name not in dep]; opt_deps = list(itertools.chain(*[opt for opt in data['project'].get('optional-dependencies', {}).values() if not any(project_name in dep for dep in opt)])); print(' '.join(deps + opt_deps))") + +## Install ROS dependencies in advance to cache the layers (speeds up rebuilds) +COPY ./package.xml "${WS}/" +RUN apt-get update && \ + rosdep update --rosdistro "${ROS_DISTRO}" && \ + DEBIAN_FRONTEND=noninteractive rosdep install --default-yes --ignore-src --rosdistro "${ROS_DISTRO}" --from-paths "${WS}" && \ + rm -rf /var/lib/apt/lists/* /root/.ros/rosdep/sources.cache + +## Copy the source code into the image +COPY . "${WS}" + +## Build Rust targets +# hadolint ignore=SC1091 +RUN source /entrypoint.bash -- && \ + cargo build --release --workspace --all-targets + +## Install project as editable Python module +# hadolint ignore=SC1091 +RUN source /entrypoint.bash -- && \ + "${ISAAC_SIM_PYTHON}" -m pip install --no-input --no-cache-dir --editable "${WS}[all]" + +## Install project as ROS package +ARG ROS_WS="/opt/ros/${ROS_DISTRO}/ws" +# hadolint ignore=SC1091 +RUN source /entrypoint.bash -- && \ + colcon build --merge-install --symlink-install --cmake-args -DPython3_EXECUTABLE=${ISAAC_SIM_PYTHON} --paths "${WS}" --build-base "${ROS_WS}/build" --install-base "${ROS_WS}/install" && \ + rm -rf ./log && \ + sed -i "s|source \"/opt/ros/${ROS_DISTRO}/setup.bash\" --|source \"${ROS_WS}/install/setup.bash\" --|" /entrypoint.bash + +## Set the default command +CMD ["bash"] + +############ +### Misc ### +############ + +## Skip writing Python bytecode to the disk to avoid polluting mounted host volume with `__pycache__` directories +ENV PYTHONDONTWRITEBYTECODE=1 + +############### +### Develop ### +############### + +# ## Install DreamerV3 locally to enable mounting the source code into the container +# ARG DREAMERV3_PATH="/root/dreamerv3" +# ARG DREAMERV3_REMOTE="https://github.com/AndrejOrsula/dreamerv3.git" +# ARG DREAMERV3_BRANCH="dev" +# RUN git clone "${DREAMERV3_REMOTE}" "${DREAMERV3_PATH}" --branch "${DREAMERV3_BRANCH}" && \ +# "${ISAAC_SIM_PYTHON}" -m pip install --no-input --no-cache-dir -r "${DREAMERV3_PATH}/embodied/requirements.txt" && \ +# "${ISAAC_SIM_PYTHON}" -m pip install --no-input --no-cache-dir -r "${DREAMERV3_PATH}/dreamerv3/requirements.txt" -f "https://storage.googleapis.com/jax-releases/jax_cuda_releases.html" && \ +# "${ISAAC_SIM_PYTHON}" -m pip install --no-input --no-cache-dir --editable "${DREAMERV3_PATH}" diff --git a/LICENSE-APACHE b/LICENSE-APACHE new file mode 100644 index 0000000..d9a10c0 --- /dev/null +++ b/LICENSE-APACHE @@ -0,0 +1,176 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS diff --git a/LICENSE-MIT b/LICENSE-MIT new file mode 100644 index 0000000..1133824 --- /dev/null +++ b/LICENSE-MIT @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 Andrej Orsula + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..2e44ab7 --- /dev/null +++ b/README.md @@ -0,0 +1,26 @@ +# Space Robotics Bench + +![](docs/src/_images/srb_multi_env.jpg) + +

+ Rust + Python + Docker + Docs +

+ +The **Space Robotics Bench** aims to be a comprehensive collection of environments and tasks for robotics research in the challenging domain of space. The benchmark covers a wide range of applications and scenarios while providing a unified framework for experimenting with new tasks. Although the primary focus is on the application of robot learning techniques, the benchmark is designed to be flexible and extensible to accommodate a variety of research directions. + +## Documentation + +The full documentation is available in its raw form inside the [docs](docs) directory. The compiled version is hosted [online](https://AndrejOrsula.github.io/space_robotics_bench) in a more accessible format. + + HTML + +## License + +This project is dual-licensed under either the [MIT](LICENSE-MIT) or [Apache 2.0](LICENSE-APACHE) licenses. + +All assets created by contributors of this repository and those generated from the included procedural pipelines are licensed under the [CC0 1.0 Universal](https://github.com/AndrejOrsula/srb_assets/blob/main/LICENSE-CC0) license. Some assets are based on modified third-party resources, which might require you to give appropriate credit to the original author. Please review [`srb_assets` repository](https://github.com/AndrejOrsula/srb_assets) for more information. + + diff --git a/apps/srb.headless.kit b/apps/srb.headless.kit new file mode 100644 index 0000000..3b6e71b --- /dev/null +++ b/apps/srb.headless.kit @@ -0,0 +1,171 @@ +[package] +title = "Space Robotics Bench | Headless" +description = "Headless experience for the Space Robotics Bench" +version = "0.0.1" +keywords = ["benchmark", "robotics", "simulation", "space"] + + +[dependencies] +## Essentials +"omni.isaac.cloner" = {} +"omni.isaac.core" = {} +"omni.kit.loop-isaac" = {} +"omni.kit.viewport.rtx" = {} +"omni.kit.window.title" = {} +"omni.physx.bundle" = {} + + +[settings] +## Name +app.name = "Space Robotics Bench" +app.window.title = "Space Robotics Bench" +crashreporter.data.experience = "Space Robotics Bench" +exts."omni.kit.window.modifier.titlebar".titleFormatString = "Space Robotics Bench" + +## Extensions +app.exts.folders = [ + "${app}", + "${exe-path}/exts", + "${exe-path}/extscore", + "${exe-path}/../exts", + "${exe-path}/../extscache", + "${exe-path}/../extsPhysics", + "${exe-path}/../isaacsim/exts", + "${exe-path}/../isaacsim/extscache", + "${exe-path}/../isaacsim/extsPhysics", +] + +## UI +app.window.hideUi = false +app.window.iconPath = "${app}/../docs/theme/favicon.png" +app.file.ignoreUnsavedOnExit = true + +## Viewport +app.viewport.defaultCamPos.x = 5 +app.viewport.defaultCamPos.y = 5 +app.viewport.defaultCamPos.z = 5 + +## Rendering +app.asyncRendering = false +app.asyncRenderingLowLatency = false +app.hydraEngine.waitIdle = true +app.renderer.resolution.height = 720 +app.renderer.resolution.width = 1280 +app.renderer.skipWhileMinimized = false +app.renderer.sleepMsOnFocus = 0 +app.renderer.sleepMsOutOfFocus = 0 +app.renderer.waitIdle = true +app.updateOrder.checkForHydraRenderComplete = 1000 +app.vsync = false +ngx.enabled = true +omni.replicator.asyncRendering = false +renderer.active = "rtx" +renderer.enabled = "rtx" +renderer.multiGpu.autoEnable = false +renderer.multiGpu.enabled = false +renderer.multiGpu.maxGpuCount = 1 +rtx-transient.dlssg.enabled = false +rtx-transient.resourcemanager.enableTextureStreaming = true +rtx.hydra.enableSemanticSchema = true +rtx.indirectDiffuse.enabled = false +rtx.mdltranslator.mdlDistilling = false +rtx.newDenoiser.enabled = false +rtx.pathtracing.optixDenoiser.enabled = false +rtx.post.dlss.execMode = 0 # 0 (Performance) | 1 (Balanced) | 2 (Quality) | 3 (Auto) +rtx.post.histogram.enabled = true # Auto exposure +rtx.raytracing.fractionalCutoutOpacity = false +rtx.raytracing.subsurface.denoiser.enabled = false +rtx.reflections.enabled = false +rtx.sceneDb.ambientLightIntensity = 1.0 +rtx.transient.dlssg.enabled = false +rtx.translucency.enabled = true +rtx.translucency.worldEps = 0.005 + +# Note: Sampled lighting works better for static scenes +rtx.directLighting.sampledLighting.enabled = false +rtx.directLighting.sampledLighting.samplesPerPixel = 1 + +## Physics +physics.autoPopupSimulationOutputWindow = false +physics.fabricUpdateForceSensors = false +physics.fabricUpdateJointStates = false +physics.fabricUpdateTransformations = false +physics.fabricUpdateVelocities = false +physics.outputVelocitiesLocalSpace = false +physics.updateForceSensorsToUsd = false +physics.updateParticlesToUsd = false +physics.updateToUsd = false +physics.updateVelocitiesToUsd = false +physics.updateVelocitiesToUsd = false +physics.useFastCache = false +physics.visualizationDisplayJoints = false +physics.visualizationSimulationOutput = false + +## Thread settings +app.runLoops.main.rateLimitEnabled = false +app.runLoops.main.rateLimitFrequency = 120 +app.runLoops.main.rateLimitUsePrecisionSleep = true +app.runLoops.main.syncToPresent = false +app.runLoops.present.rateLimitFrequency = 120 +app.runLoops.present.rateLimitUsePrecisionSleep = true +app.runLoops.rendering_0.rateLimitFrequency = 120 +app.runLoops.rendering_0.rateLimitUsePrecisionSleep = true +app.runLoops.rendering_0.syncToPresent = false +app.runLoops.rendering_1.rateLimitFrequency = 120 +app.runLoops.rendering_1.rateLimitUsePrecisionSleep = true +app.runLoops.rendering_1.syncToPresent = false +app.runLoopsGlobal.syncToPresent = false + +## Audio +app.audio.enabled = false + +## Suppress errors and warnings +app.settings.fabricDefaultStageFrameHistoryCount = 3 +app.vulkan = true +rtx.pathtracing.maxSamplesPerLaunch = 1000000 + +## Performance +exts."omni.replicator.core".Orchestrator.enabled = false +exts."omni.kit.renderer.core".present.enabled = false +exts."omni.kit.renderer.core".present.presentAfterRendering = false + + +app.settings.persistent = true +[settings.persistent] +## Stage +app.primCreation.DefaultXformOpOrder = "xformOp:translate, xformOp:orient, xformOp:scale" +app.primCreation.DefaultXformOpType = "Scale, Orient, Translate" +app.primCreation.typedDefaults.camera.clippingRange = [0.01, 10000000.0] +app.stage.instanceableOnCreatingReference = false +app.stage.materialStrength = "weakerThanDescendants" +app.stage.movePrimInPlace = false +app.stage.upAxis = "Z" +simulation.defaultMetersPerUnit = 1.0 +simulation.minFrameRate = 15 + +## UI +app.file.recentFiles = [] +app.window.uiStyle = "NvidiaDark" + +## Viewport +app.transform.gizmoUseSRT = true +app.viewport.camMoveVelocity = 0.05 +app.viewport.defaults.tickRate = 120 +app.viewport.displayOptions = 31887 +app.viewport.gizmo.scale = 0.01 +app.viewport.grid.scale = 1.0 +app.viewport.pickingMode = "kind:model.ALL" +app.viewport.previewOnPeek = false +app.viewport.snapToSurface = false + +## Rendering +omni.replicator.captureOnPlay = true +renderer.startupMessageDisplayed = false + +## Omnigraph +omnigraph.disablePrimNodes = true +omnigraph.updateToUsd = false +omnigraph.useSchemaPrims = true + +## Miscellaneous +resourcemonitor.timeBetweenQueries = 100 diff --git a/apps/srb.headless.rendering.kit b/apps/srb.headless.rendering.kit new file mode 100644 index 0000000..948ad1c --- /dev/null +++ b/apps/srb.headless.rendering.kit @@ -0,0 +1,34 @@ +[package] +title = "Space Robotics Bench | Headless + Rendering" +description = "Headless experience with rendering support for the Space Robotics Bench" +version = "0.0.1" +keywords = ["benchmark", "robotics", "simulation", "space"] + + +[dependencies] +## Rendering +"omni.replicator.core" = {} +"omni.replicator.isaac" = {} + +## Inheritance +"srb.headless" = {} + +[settings] +## Extensions +app.exts.folders = [ + "${app}", + "${exe-path}/exts", + "${exe-path}/extscore", + "${exe-path}/../exts", + "${exe-path}/../extscache", + "${exe-path}/../extsPhysics", + "${exe-path}/../isaacsim/exts", + "${exe-path}/../isaacsim/extscache", + "${exe-path}/../isaacsim/extsPhysics", +] + +## Rendering +isaaclab.cameras_enabled = true +rtx.raytracing.cached.enabled = false +rtx.raytracing.lightcache.spatialCache.enabled = false +rtx.descriptorSets = 30000 # Note: Increase for scenes with more cameras diff --git a/apps/srb.kit b/apps/srb.kit new file mode 100644 index 0000000..664dc14 --- /dev/null +++ b/apps/srb.kit @@ -0,0 +1,44 @@ +[package] +title = "Space Robotics Bench" +description = "Regular experience for the Space Robotics Bench" +version = "0.0.1" +keywords = ["benchmark", "robotics", "simulation", "space"] + + +[dependencies] +## Basic UI +"omni.kit.uiapp" = {} +"omni.kit.window.stage" = {} + +## Menu bar +"omni.kit.viewport.menubar.camera" = {} +"omni.kit.viewport.menubar.display" = {} +"omni.kit.viewport.menubar.lighting" = {} +"omni.kit.viewport.menubar.render" = {} +"omni.kit.viewport.menubar.settings" = {} + +## Property widgets +"omni.kit.property.bundle" = {} + +## Basic manipulation +"omni.kit.manipulator.camera" = {} +"omni.kit.manipulator.prim" = {} +"omni.kit.manipulator.selection" = {} + +## Inheritance +"srb.headless" = {} + + +[settings] +## Extensions +app.exts.folders = [ + "${app}", + "${exe-path}/exts", + "${exe-path}/extscore", + "${exe-path}/../exts", + "${exe-path}/../extscache", + "${exe-path}/../extsPhysics", + "${exe-path}/../isaacsim/exts", + "${exe-path}/../isaacsim/extscache", + "${exe-path}/../isaacsim/extsPhysics", +] diff --git a/apps/srb.rendering.kit b/apps/srb.rendering.kit new file mode 100644 index 0000000..5e21673 --- /dev/null +++ b/apps/srb.rendering.kit @@ -0,0 +1,35 @@ +[package] +title = "Space Robotics Bench | Rendering" +description = "Experience with rendering support for the Space Robotics Bench" +version = "0.0.1" +keywords = ["benchmark", "robotics", "simulation", "space"] + + +[dependencies] +## Rendering +"omni.replicator.core" = {} +"omni.replicator.isaac" = {} + +## Inheritance +"srb" = {} + + +[settings] +## Extensions +app.exts.folders = [ + "${app}", + "${exe-path}/exts", + "${exe-path}/extscore", + "${exe-path}/../exts", + "${exe-path}/../extscache", + "${exe-path}/../extsPhysics", + "${exe-path}/../isaacsim/exts", + "${exe-path}/../isaacsim/extscache", + "${exe-path}/../isaacsim/extsPhysics", +] + +## Rendering +isaaclab.cameras_enabled = true +rtx.raytracing.cached.enabled = false +rtx.raytracing.lightcache.spatialCache.enabled = false +rtx.descriptorSets = 30000 # Note: Increase for scenes with more cameras diff --git a/assets/srb_assets b/assets/srb_assets new file mode 160000 index 0000000..dc48f1c --- /dev/null +++ b/assets/srb_assets @@ -0,0 +1 @@ +Subproject commit dc48f1ce0b28d8cdc7a6cde563309c28aee479b2 diff --git a/config/env.yaml b/config/env.yaml new file mode 100644 index 0000000..a4d1982 --- /dev/null +++ b/config/env.yaml @@ -0,0 +1,12 @@ +seed: 42 # SRB_SEED [int] +scenario: mars # SRB_SCENARIO [mars, moon, orbit] +detail: 0.5 # SRB_DETAIL [float] +assets: + robot: + variant: dataset # SRB_ASSETS_ROBOT_VARIANT [dataset] + object: + variant: procedural # SRB_ASSETS_OBJECT_VARIANT [primitive, dataset, procedural] + terrain: + variant: procedural # SRB_ASSETS_TERRAIN_VARIANT [none, primitive, dataset, procedural] + vehicle: + variant: dataset # SRB_ASSETS_VEHICLE_VARIANT [none, dataset] diff --git a/config/rviz/default.rviz b/config/rviz/default.rviz new file mode 100644 index 0000000..9a5ce4b --- /dev/null +++ b/config/rviz/default.rviz @@ -0,0 +1,310 @@ +Panels: + - Class: rviz_common/Displays + Help Height: 0 + Name: Displays + Property Tree Widget: + Expanded: ~ + Splitter Ratio: 0.75 + Tree Height: 294 + - Class: rviz_common/Selection + Name: Selection + - Class: rviz_common/Tool Properties + Expanded: + - /2D Goal Pose1 + - /Publish Point1 + Name: Tool Properties + Splitter Ratio: 0.6 + - Class: rviz_common/Views + Expanded: + - /Current View1 + Name: Views + Splitter Ratio: 0.5 + - Class: rviz_common/Time + Experimental: false + Name: Time + SyncMode: 0 + SyncSource: "" +Visualization Manager: + Class: "" + Displays: + - Alpha: 0.25 + Cell Size: 1 + Class: rviz_default_plugins/Grid + Color: 127; 127; 127 + Enabled: true + Line Style: + Line Width: 0.029999999329447746 + Value: Lines + Name: Grid + Normal Cell Count: 0 + Offset: + X: 0 + Y: 0 + Z: 0 + Plane: XY + Plane Cell Count: 100 + Reference Frame: + Value: true + - Class: rviz_default_plugins/TF + Enabled: true + Frame Timeout: 15 + Frames: + All Enabled: true + robot/base: + Value: true + world: + Value: true + Marker Scale: 1 + Name: TF + Show Arrows: true + Show Axes: true + Show Names: false + Tree: + world: + robot/base: {} + Update Interval: 0 + Value: true + - Class: rviz_default_plugins/Image + Enabled: true + Max Value: 1 + Median window: 5 + Min Value: 0 + Name: Scene Camera (RGB) + Normalize Range: true + Topic: + Depth: 5 + Durability Policy: Volatile + History Policy: Keep Last + Reliability Policy: Reliable + Value: /camera_scene/image_raw + Value: true + - Class: rviz_default_plugins/Image + Enabled: true + Max Value: 1 + Median window: 5 + Min Value: 0 + Name: Scene Camera (Depth) + Normalize Range: true + Topic: + Depth: 5 + Durability Policy: Volatile + History Policy: Keep Last + Reliability Policy: Reliable + Value: /camera_scene/depth/image_raw + Value: true + - Class: rviz_default_plugins/Image + Enabled: true + Max Value: 1 + Median window: 5 + Min Value: 0 + Name: Robot Base Camera (RGB) + Normalize Range: true + Topic: + Depth: 5 + Durability Policy: Volatile + History Policy: Keep Last + Reliability Policy: Reliable + Value: /robot/camera_base/image_raw + Value: true + - Class: rviz_default_plugins/Image + Enabled: true + Max Value: 1 + Median window: 5 + Min Value: 0 + Name: Robot Base Camera (Depth) + Normalize Range: true + Topic: + Depth: 5 + Durability Policy: Volatile + History Policy: Keep Last + Reliability Policy: Reliable + Value: /robot/camera_base/depth/image_raw + Value: true + - Class: rviz_default_plugins/Image + Enabled: true + Max Value: 1 + Median window: 5 + Min Value: 0 + Name: Wrist Camera (RGB) + Normalize Range: true + Topic: + Depth: 5 + Durability Policy: Volatile + History Policy: Keep Last + Reliability Policy: Reliable + Value: /robot/camera_wrist/image_raw + Value: true + - Class: rviz_default_plugins/Image + Enabled: true + Max Value: 1 + Median window: 5 + Min Value: 0 + Name: Wrist Camera (Depth) + Normalize Range: true + Topic: + Depth: 5 + Durability Policy: Volatile + History Policy: Keep Last + Reliability Policy: Reliable + Value: /robot/camera_wrist/depth/image_raw + Value: true + - Class: rviz_default_plugins/Image + Enabled: true + Max Value: 1 + Median window: 5 + Min Value: 0 + Name: Front Camera (RGB) + Normalize Range: true + Topic: + Depth: 5 + Durability Policy: Volatile + History Policy: Keep Last + Reliability Policy: Reliable + Value: /robot/camera_front/image_raw + Value: true + - Class: rviz_default_plugins/Image + Enabled: true + Max Value: 1 + Median window: 5 + Min Value: 0 + Name: Front Camera (Depth) + Normalize Range: true + Topic: + Depth: 5 + Durability Policy: Volatile + History Policy: Keep Last + Reliability Policy: Reliable + Value: /robot/camera_front/depth/image_raw + Value: true + - Class: rviz_default_plugins/Image + Enabled: true + Max Value: 1 + Median window: 5 + Min Value: 0 + Name: Bottom Camera (RGB) + Normalize Range: true + Topic: + Depth: 5 + Durability Policy: Volatile + History Policy: Keep Last + Reliability Policy: Reliable + Value: /robot/camera_bottom/image_raw + Value: true + - Class: rviz_default_plugins/Image + Enabled: true + Max Value: 1 + Median window: 5 + Min Value: 0 + Name: Bottom Camera (Depth) + Normalize Range: true + Topic: + Depth: 5 + Durability Policy: Volatile + History Policy: Keep Last + Reliability Policy: Reliable + Value: /robot/camera_bottom/depth/image_raw + Value: true + Enabled: true + Global Options: + Background Color: 31; 31; 31 + Fixed Frame: world + Frame Rate: 30 + Name: root + Tools: + - Class: rviz_default_plugins/Interact + Hide Inactive Objects: true + - Class: rviz_default_plugins/MoveCamera + - Class: rviz_default_plugins/Select + - Class: rviz_default_plugins/FocusCamera + - Class: rviz_default_plugins/Measure + Line color: 128; 128; 0 + - Class: rviz_default_plugins/SetInitialPose + Covariance x: 0.25 + Covariance y: 0.25 + Covariance yaw: 0.06853891909122467 + Topic: + Depth: 5 + Durability Policy: Volatile + History Policy: Keep Last + Reliability Policy: Reliable + Value: /initialpose + - Class: rviz_default_plugins/SetGoal + Topic: + Depth: 5 + Durability Policy: Volatile + History Policy: Keep Last + Reliability Policy: Reliable + Value: /goal_pose + - Class: rviz_default_plugins/PublishPoint + Single click: true + Topic: + Depth: 5 + Durability Policy: Volatile + History Policy: Keep Last + Reliability Policy: Reliable + Value: /clicked_point + Transformation: + Current: + Class: rviz_default_plugins/TF + Value: true + Views: + Current: + Class: rviz_default_plugins/Orbit + Distance: 15 + Enable Stereo Rendering: + Stereo Eye Separation: 0.05999999865889549 + Stereo Focal Distance: 1 + Swap Stereo Eyes: false + Value: false + Focal Point: + X: 0 + Y: 0 + Z: 0 + Focal Shape Fixed Size: true + Focal Shape Size: 0.05000000074505806 + Invert Z Axis: false + Name: Current View + Near Clip Distance: 0.009999999776482582 + Pitch: 0.785398006439209 + Target Frame: + Value: Orbit (rviz) + Yaw: 0.785398006439209 + Saved: ~ +Window Geometry: + Bottom Camera (Depth): + collapsed: false + Bottom Camera (RGB): + collapsed: false + Displays: + collapsed: false + Front Camera (Depth): + collapsed: false + Front Camera (RGB): + collapsed: false + Height: 1131 + Hide Left Dock: false + Hide Right Dock: false + QMainWindow State: 000000ff00000000fd0000000400000000000001560000043cfc020000000dfb0000001200530065006c0065006300740069006f006e00000001e10000009b0000005c00fffffffb0000001e0054006f006f006c002000500072006f007000650072007400690065007302000001ed000001df00000185000000a3fb000000120056006900650077007300200054006f006f02000001df000002110000018500000122fb000000200054006f006f006c002000500072006f0070006500720074006900650073003203000002880000011d000002210000017afb000000100044006900730070006c006100790073010000001400000161000000c700fffffffb0000002000730065006c0065006300740069006f006e00200062007500660066006500720200000138000000aa0000023a00000294fb00000014005700690064006500530074006500720065006f02000000e6000000d2000003ee0000030bfb0000000c004b0069006e0065006300740200000186000001060000030c00000261fb0000000a0049006d006100670065000000037e000000d20000000000000000fb0000002400460072006f006e0074002000430061006d006500720061002000280052004700420029010000017b000000b80000002800fffffffb0000002800460072006f006e0074002000430061006d006500720061002000280044006500700074006800290100000239000000a90000002800fffffffb000000260042006f00740074006f006d002000430061006d00650072006100200028005200470042002901000002e8000000af0000002800fffffffb0000002a0042006f00740074006f006d002000430061006d00650072006100200028004400650070007400680029010000039d000000b30000002800ffffff00000001000001340000043cfc0200000009fb00000024005300630065006e0065002000430061006d0065007200610020002800520047004200290100000014000000a50000002800fffffffb00000028005300630065006e0065002000430061006d0065007200610020002800440065007000740068002901000000bf000000a80000002800fffffffb0000002e0052006f0062006f007400200042006100730065002000430061006d006500720061002000280052004700420029010000016d000000ad0000002800fffffffb000000320052006f0062006f007400200042006100730065002000430061006d006500720061002000280044006500700074006800290100000220000000b80000002800fffffffb0000002400570072006900730074002000430061006d00650072006100200028005200470042002901000002de000000b40000002800fffffffb0000002800570072006900730074002000430061006d006500720061002000280044006500700074006800290100000398000000b80000002800fffffffb0000001e0054006f006f006c002000500072006f00700065007200740069006500730100000041000000780000000000000000fb0000000a00560069006500770073000000003b000003d2000000a000fffffffb0000001200530065006c0065006300740069006f006e010000025a000000b200000000000000000000000200000490000000a9fc0100000001fb0000000a00560069006500770073030000004e00000080000002e10000019700000003000007800000003dfc0100000002fb0000000800540069006d00650000000000000007800000025300fffffffb0000000800540069006d00650100000000000004500000000000000000000004ea0000043c00000004000000040000000800000008fc0000000100000002000000010000000a0054006f006f006c00730000000000ffffffff0000000000000000 + Robot Base Camera (Depth): + collapsed: false + Robot Base Camera (RGB): + collapsed: false + Scene Camera (Depth): + collapsed: false + Scene Camera (RGB): + collapsed: false + Selection: + collapsed: false + Time: + collapsed: false + Tool Properties: + collapsed: false + Views: + collapsed: false + Width: 1920 + Wrist Camera (Depth): + collapsed: false + Wrist Camera (RGB): + collapsed: false + X: 0 + Y: 0 diff --git a/crates/space_robotics_bench/Cargo.toml b/crates/space_robotics_bench/Cargo.toml new file mode 100644 index 0000000..43aa2d8 --- /dev/null +++ b/crates/space_robotics_bench/Cargo.toml @@ -0,0 +1,35 @@ +[package] +name = "space_robotics_bench" +description.workspace = true +categories.workspace = true +keywords.workspace = true +readme.workspace = true +license.workspace = true +homepage.workspace = true +repository.workspace = true +documentation.workspace = true +edition.workspace = true +rust-version.workspace = true +version.workspace = true +publish.workspace = true + +[dependencies] +display_json = { workspace = true, optional = true } +fast_poisson = { workspace = true } +figment = { workspace = true } +paste = { workspace = true } +pyo3 = { workspace = true } +rand = { workspace = true } +rand_xoshiro = { workspace = true } +rayon = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true, optional = true } +serde_yaml = { workspace = true, optional = true } +thiserror = { workspace = true } +toml = { workspace = true, optional = true } + +[features] +default = ["yaml"] +json = ["dep:serde_json", "figment/json", "dep:display_json"] +toml = ["dep:toml", "figment/toml"] +yaml = ["dep:serde_yaml", "figment/yaml"] diff --git a/crates/space_robotics_bench/src/envs/cfg/assets.rs b/crates/space_robotics_bench/src/envs/cfg/assets.rs new file mode 100644 index 0000000..d955019 --- /dev/null +++ b/crates/space_robotics_bench/src/envs/cfg/assets.rs @@ -0,0 +1,48 @@ +use super::AssetVariant; +use pyo3::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, Serialize, Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] +#[cfg_attr(feature = "json", derive(display_json::DisplayAsJson))] +#[pyclass(eq, get_all, set_all)] +pub struct Assets { + pub robot: Asset, + pub object: Asset, + pub terrain: Asset, + pub vehicle: Asset, +} + +#[pymethods] +impl Assets { + fn __repr__(&self) -> String { + format!("{self:?}") + } + + #[new] + fn new(robot: Asset, object: Asset, terrain: Asset, vehicle: Asset) -> Self { + Self { + robot, + object, + terrain, + vehicle, + } + } +} + +#[derive(Deserialize, Serialize, Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] +#[pyclass(eq, get_all, set_all)] +pub struct Asset { + pub variant: AssetVariant, +} + +#[pymethods] +impl Asset { + fn __repr__(&self) -> String { + format!("{self:?}") + } + + #[new] + fn new(variant: AssetVariant) -> Self { + Self { variant } + } +} diff --git a/crates/space_robotics_bench/src/envs/cfg/config.rs b/crates/space_robotics_bench/src/envs/cfg/config.rs new file mode 100644 index 0000000..15f4462 --- /dev/null +++ b/crates/space_robotics_bench/src/envs/cfg/config.rs @@ -0,0 +1,241 @@ +use super::{Asset, AssetVariant, Assets, Scenario}; +use crate::Result; +#[cfg(any(feature = "json", feature = "toml", feature = "yaml"))] +use figment::providers::Format; +use figment::Figment; +use pyo3::prelude::*; +use serde::{Deserialize, Serialize}; +#[cfg(any(feature = "json", feature = "toml", feature = "yaml"))] +use std::io::Write; + +#[derive(Deserialize, Serialize, Debug, Clone, Copy, PartialEq)] +#[cfg_attr(feature = "json", derive(display_json::DisplayAsJson))] +#[pyclass(eq, get_all, set_all)] +pub struct EnvironmentConfig { + pub scenario: Scenario, + pub assets: Assets, + pub seed: u64, + pub detail: f32, +} + +impl Default for EnvironmentConfig { + fn default() -> Self { + Self { + scenario: Scenario::Moon, + assets: Assets { + robot: Asset { + variant: AssetVariant::Dataset, + }, + object: Asset { + variant: AssetVariant::Procedural, + }, + terrain: Asset { + variant: AssetVariant::Procedural, + }, + vehicle: Asset { + variant: AssetVariant::Dataset, + }, + }, + seed: 0, + detail: 1.0, + } + } +} + +const SUPPORTED_FILE_EXTENSIONS: &[&str] = &[ + #[cfg(feature = "json")] + "json", + #[cfg(feature = "toml")] + "toml", + #[cfg(feature = "yaml")] + "yaml", + #[cfg(feature = "yaml")] + "yml", +]; + +impl EnvironmentConfig { + pub fn extract( + cfg_path: Option>, + env_prefix: Option<&str>, + other: Option, + ) -> Result { + let mut figment = Figment::new(); + + // 1. (Optional) Load configuration from file + if let Some(path) = cfg_path { + let path = path.as_ref(); + if !path.exists() { + return Err(std::io::Error::new( + std::io::ErrorKind::NotFound, + format!("File not found: {}", path.display()), + ) + .into()); + } + match path.extension().and_then(|ext| ext.to_str()) { + #[cfg(feature = "json")] + Some("json") => figment = figment.merge(figment::providers::Json::file(path)), + #[cfg(feature = "toml")] + Some("toml") => figment = figment.merge(figment::providers::Toml::file(path)), + #[cfg(feature = "yaml")] + Some("yaml" | "yml") => { + figment = figment.merge(figment::providers::Yaml::file(path)); + } + Some(_) => { + return Err(figment::Error::from(format!( + "Unsupported file extension: {} (supported=[{}])", + path.display(), + SUPPORTED_FILE_EXTENSIONS.join(", "), + )) + .into()) + } + None => { + return Err(figment::Error::from(format!( + "Missing file extension: {}", + path.display() + )) + .into()) + } + } + } + + // 2. (Optional) Load configuration from environment variables + if let Some(env_prefix) = env_prefix { + figment = figment.merge(figment::providers::Env::prefixed(env_prefix).split('_')); + } + + // 3. (Optional) Load configuration from other sources + if let Some(other) = other { + figment = figment.merge(figment::providers::Serialized::defaults(other)); + } + + // Finally, apply default values for missing fields + figment = figment.join(figment::providers::Serialized::defaults( + EnvironmentConfig::default(), + )); + + Ok(figment.extract()?) + } + + pub fn write(&self, path: impl AsRef) -> Result<()> { + let path = path.as_ref(); + + // Return error if none or if not supported extension + match path.extension().and_then(|ext| ext.to_str()) { + #[cfg(feature = "json")] + Some("json") => { + let file = std::fs::File::create(path)?; + let mut writer = std::io::BufWriter::new(file); + serde_json::to_writer_pretty(&mut writer, self).map_err(|e| { + std::io::Error::new(std::io::ErrorKind::InvalidInput, e.to_string()) + })?; + writer.flush()?; + Ok(()) + } + #[cfg(feature = "toml")] + Some("toml") => { + let mut file = std::fs::File::create(path)?; + let content = toml::to_string_pretty(self).map_err(|e| { + std::io::Error::new(std::io::ErrorKind::InvalidInput, e.to_string()) + })?; + file.write_all(content.as_bytes())?; + Ok(()) + } + #[cfg(feature = "yaml")] + Some("yaml" | "yml") => { + let file = std::fs::File::create(path)?; + let mut writer = std::io::BufWriter::new(file); + serde_yaml::to_writer(&mut writer, self).map_err(|e| { + std::io::Error::new(std::io::ErrorKind::InvalidInput, e.to_string()) + })?; + writer.flush()?; + Ok(()) + } + Some(_) | None => Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "Unsupported file extension: {} (supported=[{}])", + path.display(), + SUPPORTED_FILE_EXTENSIONS.join(", "), + ), + ) + .into()), + } + } +} + +#[pymethods] +impl EnvironmentConfig { + fn __repr__(&self) -> String { + format!("{self:?}") + } + + #[new] + fn new(scenario: Scenario, assets: Assets, seed: u64, detail: f32) -> Self { + Self { + scenario, + assets, + seed, + detail, + } + } + + #[staticmethod] + #[pyo3( + name = "extract", + signature = (cfg_path=None, env_prefix=Some("SRB_"), other=None), + )] + fn py_extract( + cfg_path: Option<&str>, + env_prefix: Option<&str>, + other: Option, + ) -> PyResult { + Ok(Self::extract(cfg_path, env_prefix, other)?) + } + + #[pyo3(name = "write")] + fn py_write(&self, path: &str) -> PyResult<()> { + Ok(self.write(path)?) + } + + #[cfg(any(feature = "json", feature = "toml", feature = "yaml"))] + fn __reduce__(&self) -> PyResult<(PyObject, PyObject)> { + Python::with_gil(|py| { + py.run_bound("import space_robotics_bench", None, None) + .unwrap(); + let deserialize = py + .eval_bound( + "space_robotics_bench._rs.envs.EnvironmentConfig._deserialize", + None, + None, + ) + .unwrap(); + + #[cfg(feature = "json")] + let data = serde_json::to_vec(self).unwrap(); + #[cfg(all(feature = "yaml", not(feature = "json")))] + let data = serde_yaml::to_string(self).unwrap().as_bytes().to_vec(); + #[cfg(all(feature = "toml", not(any(feature = "json", feature = "yaml"))))] + let data = toml::to_string(self).unwrap().as_bytes().to_vec(); + + Ok((deserialize.to_object(py), (data,).to_object(py))) + }) + } + + #[staticmethod] + #[pyo3(name = "_deserialize")] + #[cfg(any(feature = "json", feature = "toml", feature = "yaml"))] + fn deserialize(data: Vec) -> PyResult { + #[cfg(feature = "json")] + { + Ok(serde_json::from_slice(&data).unwrap()) + } + #[cfg(all(feature = "yaml", not(feature = "json")))] + { + Ok(serde_yaml::from_slice(&data).unwrap()) + } + #[cfg(all(feature = "toml", not(any(feature = "json", feature = "yaml"))))] + { + Ok(toml::from_str(std::str::from_utf8(&data).unwrap()).unwrap()) + } + } +} diff --git a/crates/space_robotics_bench/src/envs/cfg/enums/asset_variant.rs b/crates/space_robotics_bench/src/envs/cfg/enums/asset_variant.rs new file mode 100644 index 0000000..8eb7eda --- /dev/null +++ b/crates/space_robotics_bench/src/envs/cfg/enums/asset_variant.rs @@ -0,0 +1,24 @@ +use pyo3::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, Serialize, Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] +#[cfg_attr(feature = "json", derive(display_json::DisplayAsJson))] +#[pyclass(frozen, eq, eq_int, hash, rename_all = "SCREAMING_SNAKE_CASE")] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub enum AssetVariant { + #[serde(alias = "none")] + None, + #[default] + #[serde(alias = "primitive", alias = "PRIM", alias = "prim")] + Primitive, + #[serde(alias = "dataset", alias = "DB", alias = "db")] + Dataset, + #[serde( + alias = "procedural", + alias = "PROC", + alias = "proc", + alias = "PROCGEN", + alias = "procgen" + )] + Procedural, +} diff --git a/crates/space_robotics_bench/src/envs/cfg/enums/mod.rs b/crates/space_robotics_bench/src/envs/cfg/enums/mod.rs new file mode 100644 index 0000000..b830c36 --- /dev/null +++ b/crates/space_robotics_bench/src/envs/cfg/enums/mod.rs @@ -0,0 +1,5 @@ +mod asset_variant; +mod scenario; + +pub use asset_variant::*; +pub use scenario::*; diff --git a/crates/space_robotics_bench/src/envs/cfg/enums/scenario.rs b/crates/space_robotics_bench/src/envs/cfg/enums/scenario.rs new file mode 100644 index 0000000..a11d775 --- /dev/null +++ b/crates/space_robotics_bench/src/envs/cfg/enums/scenario.rs @@ -0,0 +1,240 @@ +use pyo3::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, Serialize, Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] +#[cfg_attr(feature = "json", derive(display_json::DisplayAsJson))] +#[pyclass(frozen, eq, eq_int, hash, rename_all = "SCREAMING_SNAKE_CASE")] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub enum Scenario { + #[serde(alias = "asteroid")] + Asteroid, + #[serde(alias = "earth", alias = "TERRESTRIAL", alias = "terrestrial")] + Earth, + #[serde(alias = "mars", alias = "MARTIAN", alias = "martian")] + Mars, + #[default] + #[serde(alias = "moon", alias = "LUNAR", alias = "lunar")] + Moon, + #[serde(alias = "orbit", alias = "ORBITAL", alias = "orbital")] + Orbit, +} + +impl Scenario { + /// Magnitude of gravitational acceleration in m/s². + /// + /// # Assumptions + /// + /// - Asteroid: 50% gravitational acceleration of Ceres (largest body in the asteroid belt). + /// - Orbit: No gravitational acceleration. + #[must_use] + pub fn gravity_magnitude(self) -> f64 { + match self { + Self::Asteroid => 0.14219, + Self::Earth => 9.80665, + Self::Mars => 3.72076, + Self::Moon => 1.62496, + Self::Orbit => 0.0, + } + } + + /// Difference between the maximum and minimum of gravitational acceleration in m/s². + /// + /// # Assumptions + /// + /// - Asteroid: Ceres is considered as the maximum (largest body in the asteroid belt). + /// - Orbit: No gravitational acceleration. + #[must_use] + pub fn gravity_variation(self) -> f64 { + match self { + Self::Asteroid => 2.0 * Self::Asteroid.gravity_magnitude(), + Self::Earth => 0.0698, + Self::Mars => 0.0279, + Self::Moon => 0.0253, + Self::Orbit => 0.0, + } + } + + /// Range of gravitational acceleration in m/s² calculated as the magnitude ± variation/2. + #[must_use] + pub fn gravity_range(self) -> (f64, f64) { + let magnitude = self.gravity_magnitude(); + let delta = self.gravity_variation() / 2.0; + (magnitude - delta, magnitude + delta) + } + + /// Intensity of Solar light in W/m². + /// + /// # Notes + /// + /// - Asteroid: Taken at 2.7 AU. + /// - Earth | Mars: Taken at the surface. The peak value (sunny day) is subtracted by half of the variation. + /// - Moon | Orbit: Taken at 1 AU. + #[must_use] + pub fn light_intensity(self) -> f64 { + match self { + Self::Asteroid => 190.0, + Self::Earth => 1000.0 - Self::Earth.light_intensity_variation() / 2.0, + Self::Mars => 842.0 - Self::Mars.light_intensity_variation() / 2.0, + Self::Moon | Self::Orbit => 1361.0, + } + } + + /// Difference between the maximum and minimum of Solar light intensity in W/m². + /// + /// # Notes + /// + /// - Asteroid: Approximate range between 2.55 and 2.97 AU. + /// - Earth | Mars: Guesstimated effect of atmosphere and weather. + /// - Moon | Orbit: Minor variation due to elliptical orbit. + #[must_use] + pub fn light_intensity_variation(self) -> f64 { + match self { + Self::Asteroid => 50.0, + Self::Earth => 450.0, + Self::Mars => 226.0, + Self::Moon | Self::Orbit => 0.5, + } + } + + /// Range of Solar light intensity in W/m² calculated as the intensity ± variation/2. + #[must_use] + pub fn light_intensity_range(self) -> (f64, f64) { + let intensity = self.light_intensity(); + let delta = self.light_intensity_variation() / 2.0; + (intensity - delta, intensity + delta) + } + + /// Angular diameter of the Solar light source in degrees. + /// + /// # Assumptions + /// + /// - Earth | Mars: Taken at their distance from the Sun. + /// - Asteroid | Moon | Orbit: Approximated as a point source due to lack of atmosphere. + #[must_use] + pub fn light_angular_diameter(self) -> f64 { + match self { + Self::Earth => 0.53, + Self::Mars => 0.35, + Self::Asteroid | Self::Moon | Self::Orbit => 0.0, + } + } + + /// Variation of the angular diameter of the Solar light source in degrees. + #[must_use] + pub fn light_angular_diameter_variation(self) -> f64 { + match self { + Self::Earth => 0.021, + Self::Mars => 0.08, + Self::Asteroid | Self::Moon | Self::Orbit => 0.0, + } + } + + /// Range of the angular diameter of the Solar light source in degrees calculated as the diameter ± variation/2. + #[must_use] + pub fn light_angular_diameter_range(self) -> (f64, f64) { + let diameter = self.light_angular_diameter(); + let delta = self.light_angular_diameter_variation() / 2.0; + (diameter - delta, diameter + delta) + } + + /// Temperature of the Solar light source in K. + /// + /// # Assumptions + /// + /// - Earth | Mars: Guesstimated effect atmosphere and weather. + /// - Asteroid | Moon | Orbit: Intrinsic color temperature of the Sun. + #[must_use] + pub fn light_color_temperature(self) -> f64 { + match self { + Self::Earth => 5750.0, + Self::Mars => 6250.0, + Self::Asteroid | Self::Moon | Self::Orbit => 5778.0, + } + } + + /// Variation of the temperature of the Solar light source in K. + /// + /// # Assumptions + /// + /// - Earth | Mars: Guesstimated effect atmosphere and weather. + /// - Asteroid | Moon | Orbit: No significant variation. + #[must_use] + pub fn light_color_temperature_variation(self) -> f64 { + match self { + Self::Earth => 1500.0, + Self::Mars => 500.0, + Self::Asteroid | Self::Moon | Self::Orbit => 0.0, + } + } + + /// Range of the temperature of the Solar light source in K calculated as the temperature ± variation/2. + #[must_use] + pub fn light_color_temperature_range(self) -> (f64, f64) { + let temperature = self.light_color_temperature(); + let delta = self.light_color_temperature_variation() / 2.0; + (temperature - delta, temperature + delta) + } +} + +#[pymethods] +impl Scenario { + #[getter("gravity_magnitude")] + fn py_gravity_magnitude(&self) -> PyResult { + Ok(self.gravity_magnitude()) + } + + #[getter("gravity_variation")] + fn py_gravity_variation(&self) -> PyResult { + Ok(self.gravity_variation()) + } + + #[getter("gravity_range")] + fn py_gravity_range(&self) -> PyResult<(f64, f64)> { + Ok(self.gravity_range()) + } + + #[getter("light_intensity")] + fn py_light_intensity(&self) -> PyResult { + Ok(self.light_intensity()) + } + + #[getter("light_intensity_variation")] + fn py_light_intensity_variation(&self) -> PyResult { + Ok(self.light_intensity_variation()) + } + + #[getter("light_intensity_range")] + fn py_light_intensity_range(&self) -> PyResult<(f64, f64)> { + Ok(self.light_intensity_range()) + } + + #[getter("light_angular_diameter")] + fn py_light_angular_diameter(&self) -> PyResult { + Ok(self.light_angular_diameter()) + } + + #[getter("light_angular_diameter_variation")] + fn py_light_angular_diameter_variation(&self) -> PyResult { + Ok(self.light_angular_diameter_variation()) + } + + #[getter("light_angular_diameter_range")] + fn py_light_angular_diameter_range(&self) -> PyResult<(f64, f64)> { + Ok(self.light_angular_diameter_range()) + } + + #[getter("light_color_temperature")] + fn py_light_color_temperature(&self) -> PyResult { + Ok(self.light_color_temperature()) + } + + #[getter("light_color_temperature_variation")] + fn py_light_color_temperature_variation(&self) -> PyResult { + Ok(self.light_color_temperature_variation()) + } + + #[getter("light_color_temperature_range")] + fn py_light_color_temperature_range(&self) -> PyResult<(f64, f64)> { + Ok(self.light_color_temperature_range()) + } +} diff --git a/crates/space_robotics_bench/src/envs/cfg/mod.rs b/crates/space_robotics_bench/src/envs/cfg/mod.rs new file mode 100644 index 0000000..324e372 --- /dev/null +++ b/crates/space_robotics_bench/src/envs/cfg/mod.rs @@ -0,0 +1,61 @@ +mod assets; +mod config; +mod enums; + +pub use assets::*; +pub use config::*; +pub use enums::*; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + #[cfg(feature = "yaml")] + fn extract() { + figment::Jail::expect_with(|jail| { + // Arrange + let cfg_path = if cfg!(feature = "yaml") { + jail.create_file( + "config.yaml", + r" + seed: 42 + assets: + vehicle: + variant: dataset + ", + )?; + Some("config.yaml") + } else { + None + }; + jail.set_env("SRB_SCENARIO", "ORBIT"); + jail.set_env("SRB_ASSETS_TERRAIN_VARIANT", "PROCEDURAL"); + jail.set_env("SRB_DETAIL", "0.2"); + + // Act + let config = EnvironmentConfig::extract(cfg_path, Some("SRB_"), None)?; + + // Assert + assert_eq!( + config, + EnvironmentConfig { + scenario: Scenario::Orbit, + assets: Assets { + vehicle: Asset { + variant: AssetVariant::Dataset, + }, + terrain: Asset { + variant: AssetVariant::Procedural, + }, + ..EnvironmentConfig::default().assets + }, + seed: 42, + detail: 0.2, + } + ); + + Ok(()) + }); + } +} diff --git a/crates/space_robotics_bench/src/envs/mod.rs b/crates/space_robotics_bench/src/envs/mod.rs new file mode 100644 index 0000000..a96d5f1 --- /dev/null +++ b/crates/space_robotics_bench/src/envs/mod.rs @@ -0,0 +1,3 @@ +mod cfg; + +pub use cfg::*; diff --git a/crates/space_robotics_bench/src/error.rs b/crates/space_robotics_bench/src/error.rs new file mode 100644 index 0000000..d9c2070 --- /dev/null +++ b/crates/space_robotics_bench/src/error.rs @@ -0,0 +1,33 @@ +/// Alias for [`std::result::Result`] that wraps our [Error] type. +pub type Result = std::result::Result; + +/// Error type for this crate. +#[derive(thiserror::Error, Debug)] +pub enum Error { + #[error(transparent)] + Config(#[from] figment::Error), + #[error(transparent)] + Io(#[from] std::io::Error), + #[error(transparent)] + Python(#[from] pyo3::PyErr), +} + +/// Convert [Error] to [`figment::Error`]. +impl From for figment::Error { + fn from(e: Error) -> Self { + match e { + Error::Config(e) => e, + _ => figment::Error::from(e.to_string()), + } + } +} + +/// Convert [Error] to [`pyo3::PyErr`]. +impl From for pyo3::PyErr { + fn from(e: Error) -> Self { + match e { + Error::Python(e) => e, + _ => pyo3::exceptions::PyException::new_err(e.to_string()), + } + } +} diff --git a/crates/space_robotics_bench/src/lib.rs b/crates/space_robotics_bench/src/lib.rs new file mode 100644 index 0000000..66928a5 --- /dev/null +++ b/crates/space_robotics_bench/src/lib.rs @@ -0,0 +1,5 @@ +pub mod envs; +pub mod utils; + +mod error; +pub use error::{Error, Result}; diff --git a/crates/space_robotics_bench/src/utils/mod.rs b/crates/space_robotics_bench/src/utils/mod.rs new file mode 100644 index 0000000..4aee55b --- /dev/null +++ b/crates/space_robotics_bench/src/utils/mod.rs @@ -0,0 +1 @@ +pub mod sampling; diff --git a/crates/space_robotics_bench/src/utils/sampling/mod.rs b/crates/space_robotics_bench/src/utils/sampling/mod.rs new file mode 100644 index 0000000..c01dd83 --- /dev/null +++ b/crates/space_robotics_bench/src/utils/sampling/mod.rs @@ -0,0 +1,3 @@ +mod poisson_disk; + +pub use poisson_disk::*; diff --git a/crates/space_robotics_bench/src/utils/sampling/poisson_disk.rs b/crates/space_robotics_bench/src/utils/sampling/poisson_disk.rs new file mode 100644 index 0000000..8c2b61e --- /dev/null +++ b/crates/space_robotics_bench/src/utils/sampling/poisson_disk.rs @@ -0,0 +1,85 @@ +use pyo3::prelude::*; +use rand::seq::SliceRandom; +use rand_xoshiro::rand_core::SeedableRng; +use rayon::prelude::*; + +const RADIUS_DECAY: f32 = 0.98; + +macro_rules! sample_poisson_disk_impl { + ($dim:expr) => { + paste::paste! { + #[pyfunction] + #[must_use] pub fn [< sample_poisson_disk_ $dim d >]( + num_samples: usize, + bounds: [[f32; $dim]; 2], + radius: f32, + ) -> Vec<[f32; $dim]> { + let dimensions = { + let mut dims = bounds[1]; + for (i, d) in dims.iter_mut().enumerate() { + *d -= bounds[0][i]; + } + dims + }; + + let mut distribution = fast_poisson::Poisson::<$dim>::new().with_dimensions(dimensions, radius); + + let mut decay_factor = RADIUS_DECAY; + loop { + let points = distribution.generate(); + match points.len() { + n if n < num_samples => { + distribution.set_dimensions(dimensions, decay_factor * radius); + decay_factor *= RADIUS_DECAY; + continue; + } + n if n == num_samples => { + break points + .into_iter() + .map(|mut point| { + for (i, p) in point.iter_mut().enumerate() { + *p += bounds[0][i]; + } + point + }) + .collect() + } + _ => { + break points + .choose_multiple( + &mut rand_xoshiro::Xoshiro128StarStar::from_entropy(), + num_samples, + ) + .cloned() + .map(|mut point| { + for (i, p) in point.iter_mut().enumerate() { + *p += bounds[0][i]; + } + point + }) + .collect() + } + } + } + } + + #[pyfunction] + #[must_use] pub fn [< sample_poisson_disk_ $dim d_looped >]( + num_samples: [usize; 2], + bounds: [[f32; $dim]; 2], + radius: f32, + ) -> Vec> { + (0..num_samples[0]) + .into_par_iter() + .map(|_| [< sample_poisson_disk_ $dim d >](num_samples[1], bounds, radius)) + .collect() + } + + } + }; + [$( $dim:expr ),*] => { + $( sample_poisson_disk_impl!($dim); )* + }; +} + +sample_poisson_disk_impl![2, 3]; diff --git a/crates/space_robotics_bench_gui/Cargo.toml b/crates/space_robotics_bench_gui/Cargo.toml new file mode 100644 index 0000000..c3f6251 --- /dev/null +++ b/crates/space_robotics_bench_gui/Cargo.toml @@ -0,0 +1,54 @@ +[package] +name = "space_robotics_bench_gui" +description.workspace = true +categories.workspace = true +keywords.workspace = true +readme.workspace = true +license.workspace = true +homepage.workspace = true +repository.workspace = true +documentation.workspace = true +edition.workspace = true +rust-version.workspace = true +version.workspace = true +publish.workspace = true + +[[bin]] +name = "gui" +path = "src/main.rs" + +[dependencies] +space_robotics_bench = { workspace = true, features = ["json"] } + +eframe = { version = "0.29", default-features = false, features = [ + "accesskit", + "glow", + "persistence", +] } +egui = { version = "0.29", default-features = false } +egui_extras = { version = "0.29", default-features = false, features = [ + "all_loaders", + "syntect", +] } +egui_commonmark = { version = "0.18", default-features = false, features = [ + "better_syntax_highlighting", + "macros", + "pulldown_cmark", +] } +r2r = { version = "0.9" } + +display_json = { workspace = true } +chrono = { workspace = true } +const_format = { workspace = true } +home = { workspace = true } +image = { workspace = true } +itertools = { workspace = true } +nix = { workspace = true } +paste = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +subprocess = { workspace = true } +sysinfo = { workspace = true } +tracing = { workspace = true } +tracing-subscriber = { workspace = true } +typed-builder = { workspace = true } diff --git a/crates/space_robotics_bench_gui/assets/fonts/Inter-Regular.otf b/crates/space_robotics_bench_gui/assets/fonts/Inter-Regular.otf new file mode 100644 index 0000000..2d0bd1d Binary files /dev/null and b/crates/space_robotics_bench_gui/assets/fonts/Inter-Regular.otf differ diff --git a/crates/space_robotics_bench_gui/assets/fonts/MaterialIconsRound-Regular.otf b/crates/space_robotics_bench_gui/assets/fonts/MaterialIconsRound-Regular.otf new file mode 100644 index 0000000..1b4a78b Binary files /dev/null and b/crates/space_robotics_bench_gui/assets/fonts/MaterialIconsRound-Regular.otf differ diff --git a/crates/space_robotics_bench_gui/assets/fonts/MonaspaceNeon-Medium.otf b/crates/space_robotics_bench_gui/assets/fonts/MonaspaceNeon-Medium.otf new file mode 100644 index 0000000..ca4c04c Binary files /dev/null and b/crates/space_robotics_bench_gui/assets/fonts/MonaspaceNeon-Medium.otf differ diff --git a/crates/space_robotics_bench_gui/build.rs b/crates/space_robotics_bench_gui/build.rs new file mode 100644 index 0000000..4ed538b --- /dev/null +++ b/crates/space_robotics_bench_gui/build.rs @@ -0,0 +1,4 @@ +fn main() { + // Monitor the content directory for changes + println!("cargo:rerun-if-changed=content"); +} diff --git a/crates/space_robotics_bench_gui/content/_images b/crates/space_robotics_bench_gui/content/_images new file mode 120000 index 0000000..42bd8ff --- /dev/null +++ b/crates/space_robotics_bench_gui/content/_images @@ -0,0 +1 @@ +../../../docs/src/_images \ No newline at end of file diff --git a/crates/space_robotics_bench_gui/content/about.md b/crates/space_robotics_bench_gui/content/about.md new file mode 100644 index 0000000..1ebd581 --- /dev/null +++ b/crates/space_robotics_bench_gui/content/about.md @@ -0,0 +1,7 @@ +This demo enables participants to experience teleoperation of a robotic arm under various extraterrestrial conditions. + +While performing a task, participants can opt-in for **collection of the demonstrated trajectories** to contribute towards a training dataset usable for **Imitation Learning**, **Offline Reinforcement Learning**, and other related areas of **Robot Learning** research focused on contact-rich manipulation in space. + +The simulation environments are built on top of **NVIDIA Isaac Sim**, and this graphical interface is developed using **egui**. Most communication among the components is facilitated using **ROS 2** while employing its C++ (rclcpp), Python (rclpy), and Rust (r2r) client libraries for the different components. + +The demo is developed by **Andrej Orsula**. diff --git a/crates/space_robotics_bench_gui/content/controls.md b/crates/space_robotics_bench_gui/content/controls.md new file mode 100644 index 0000000..923db5a --- /dev/null +++ b/crates/space_robotics_bench_gui/content/controls.md @@ -0,0 +1 @@ +This demo uses 3D Systems Touch, a haptic device that controls the robot and provides you with a sense of touch through force feedback. **Grab the pen and try it yourself!** diff --git a/crates/space_robotics_bench_gui/src/app.rs b/crates/space_robotics_bench_gui/src/app.rs new file mode 100644 index 0000000..313b8a2 --- /dev/null +++ b/crates/space_robotics_bench_gui/src/app.rs @@ -0,0 +1,1713 @@ +use std::io::{Read, Write}; + +use eframe::epaint::Color32; +use egui_commonmark::{commonmark_str, CommonMarkCache}; +use r2r::{ + std_msgs::msg::{Bool as BoolMsg, Empty as EmptyMsg, Float64 as Float64Msg}, + QosProfile, +}; +use tracing::{error, info, trace, warn}; + +use crate::page::Page; + +const LOGFILE_PATH: &str = "space_robotics_bench_gui.log"; + +#[derive(serde::Deserialize, serde::Serialize)] +#[serde(default)] +pub struct App { + theme: egui::Theme, + #[serde(skip)] + current_page: Page, + #[serde(skip)] + task_config: crate::config::TaskConfig, + #[serde(skip)] + subprocess: Option, + + #[serde(skip)] + show_about: bool, + #[serde(skip)] + show_virtual_keyboard_window: bool, + #[serde(skip)] + show_developer_options: bool, + + #[serde(skip)] + collect_trajectory: bool, + #[serde(skip)] + n_collected_trajectories: usize, + + #[serde(skip)] + gravity: f64, + #[serde(skip)] + latency: f64, + #[serde(skip)] + motion_sensitivity: f64, + #[serde(skip)] + force_feedback_sensitivity: f64, + #[serde(skip)] + max_feedback_force: f64, + + #[serde(skip)] + prev_gravity: f64, + #[serde(skip)] + prev_latency: f64, + #[serde(skip)] + prev_motion_sensitivity: f64, + #[serde(skip)] + prev_force_feedback_sensitivity: f64, + #[serde(skip)] + prev_max_feedback_force: f64, + + #[serde(skip)] + node: r2r::Node, + #[serde(skip)] + pub_gripper_toggle: r2r::Publisher, + #[serde(skip)] + pub_reset_save_dataset: r2r::Publisher, + #[serde(skip)] + pub_reset_discard_dataset: r2r::Publisher, + #[serde(skip)] + pub_gracefully_shutdown_process: r2r::Publisher, + #[serde(skip)] + pub_gravity: r2r::Publisher, + #[serde(skip)] + pub_latency: r2r::Publisher, + #[serde(skip)] + pub_motion_sensitivity: r2r::Publisher, + #[serde(skip)] + pub_force_feedback_sensitivity: r2r::Publisher, + #[serde(skip)] + pub_max_feedback_force: r2r::Publisher, + + #[serde(skip)] + last_message_pub: std::time::Instant, + + #[serde(skip)] + logfile: std::fs::File, + + #[serde(skip)] + hovered_task: Option, + + #[serde(skip)] + commonmark_cache: CommonMarkCache, +} + +impl Default for App { + fn default() -> Self { + let ctx = r2r::Context::create().unwrap(); + let mut node = r2r::Node::create(ctx, "space_robotics_bench_gui_gui", "").unwrap(); + + let pub_gripper_toggle = node + .create_publisher::("touch/event", QosProfile::default()) + .unwrap(); + let pub_reset_save_dataset = node + .create_publisher::("gui/reset_save_dataset", QosProfile::default()) + .unwrap(); + let pub_reset_discard_dataset = node + .create_publisher::("gui/reset_discard_dataset", QosProfile::default()) + .unwrap(); + let pub_gracefully_shutdown_process = node + .create_publisher::("gui/shutdown_process", QosProfile::default()) + .unwrap(); + let pub_gravity = node + .create_publisher::("gui/gravity", QosProfile::default()) + .unwrap(); + let pub_latency = node + .create_publisher::("gui/latency", QosProfile::default()) + .unwrap(); + let pub_motion_sensitivity = node + .create_publisher::("gui/motion_sensitivity", QosProfile::default()) + .unwrap(); + let pub_force_feedback_sensitivity = node + .create_publisher::("gui/force_feedback_sensitivity", QosProfile::default()) + .unwrap(); + let pub_max_feedback_force = node + .create_publisher::("gui/max_feedback_force", QosProfile::default()) + .unwrap(); + + let last_message_pub = std::time::Instant::now(); + + let mut logfile = std::fs::OpenOptions::new() + .create(true) + .read(true) + .append(true) + .open(LOGFILE_PATH) + .unwrap(); + + // Iterate over the lines and count all collected trajectories + let n_collected_trajectories = { + let mut file_str = String::new(); + logfile.read_to_string(&mut file_str).unwrap(); + file_str + .clone() + .lines() + .filter(|x| x.starts_with("DATA")) + .count() + }; + + writeln!( + logfile, + "START, {}, {}", + chrono::Local::now().format("%Y-%m-%d %H:%M:%S"), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(), + ) + .unwrap(); + + Self { + theme: egui::Theme::Dark, + current_page: Page::default(), + task_config: crate::config::TaskConfig::default(), + subprocess: None, + + show_about: false, + show_virtual_keyboard_window: false, + show_developer_options: false, + + collect_trajectory: false, + n_collected_trajectories, + + gravity: crate::config::TaskConfig::default() + .env_cfg + .scenario + .gravity_magnitude(), + latency: 0.0, + motion_sensitivity: 1.0, + force_feedback_sensitivity: 1.0, + max_feedback_force: 2.0, + + prev_gravity: 0.0, + prev_latency: 0.0, + prev_motion_sensitivity: 0.0, + prev_force_feedback_sensitivity: 0.0, + prev_max_feedback_force: 0.0, + + node, + pub_gripper_toggle, + pub_reset_save_dataset, + pub_reset_discard_dataset, + pub_gracefully_shutdown_process, + pub_gravity, + pub_latency, + pub_motion_sensitivity, + pub_force_feedback_sensitivity, + pub_max_feedback_force, + + last_message_pub, + + logfile, + + hovered_task: None, + + commonmark_cache: CommonMarkCache::default(), + } + } +} + +impl App { + #[must_use] + pub fn new(cc: &eframe::CreationContext) -> Self { + // Enable image loading + egui_extras::install_image_loaders(&cc.egui_ctx); + + // Load the fonts + crate::style::load_fonts(&cc.egui_ctx); + + // Enable screen web reader support + #[cfg(target_arch = "wasm32")] + cc.egui_ctx.options_mut(|o| o.screen_reader = true); + + // Construct the app state + let mut app = if let Some(storage) = cc.storage { + // Try to restore previous state + eframe::get_value(storage, eframe::APP_KEY).unwrap_or_default() + } else { + // Otherwise, use default state + Self::default() + }; + + // Set the theme + crate::style::set_theme(&cc.egui_ctx, app.theme); + + // Publish messages + app.publish_messages(); + + app + } +} + +impl eframe::App for App { + fn update(&mut self, ctx: &egui::Context, _frame: &mut eframe::Frame) { + // Synchronize the page URL and content if the URL contains a hash + #[cfg(target_arch = "wasm32")] + if let Some(page) = frame.info().web_info.location.hash.strip_prefix('#') { + if let Some(page) = crate::ENABLED_PAGES + .into_iter() + .find(|x| x.to_string().eq_ignore_ascii_case(page)) + { + // If a known page was requested, update the current page + self.current_page = page; + } else { + // If an unknown page was requested, update the URL to open the default page + crate::utils::egui::open_url_on_page(ctx, Page::default(), true); + } + } else { + // Otherwise, update the URL to match the current page + crate::utils::egui::open_url_on_page(ctx, self.current_page, true); + } + + // Support native fullscreen toggle + #[cfg(not(target_arch = "wasm32"))] + if ctx.input_mut(|i| i.consume_key(egui::Modifiers::NONE, egui::Key::F11)) { + let fullscreen = ctx.input(|i| i.viewport().fullscreen.unwrap_or(false)); + ctx.send_viewport_cmd(egui::ViewportCommand::Fullscreen(!fullscreen)); + } + + // Navigation panel that allows switching between page + egui::TopBottomPanel::top("top_panel").show(ctx, |ui| { + ui.with_layout(egui::Layout::left_to_right(egui::Align::Center), |ui| { + // Navigation + ui.spacing_mut().item_spacing.x = 16.0; + self.navigation_buttons(ui); + + ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| { + self.advanced_opts_button(ui); + + self.show_top_center_bar(ui); + }); + }); + }); + + // Bottom panel + egui::TopBottomPanel::bottom("bottom_panel").show(ctx, |ui| { + ui.with_layout(egui::Layout::left_to_right(egui::Align::Center), |ui| { + ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| { + self.about_button(ui); + + ui.with_layout(egui::Layout::left_to_right(egui::Align::Center), |ui| { + self.dark_mode_toggle_button(ui); + ui.separator(); + self.show_trajectory_collection_checkbox(ui); + self.warn_if_debug_build(ui); + }); + }); + }); + }); + + // Central panel + crate::utils::egui::ScrollableFramedCentralPanel::builder() + .max_content_width(ctx.screen_rect().width()) + .build() + .show(ctx, |ui| { + match self.current_page { + Page::QuickStart => { + self.quickstart_page(ui); + } + Page::Interface => { + self.configuration_page(ui); + } + }; + self.about_window(ui); + self.virtual_keyboard_window(ui); + self.developer_options_window(ui); + }); + + // Publish values + if self.current_page == Page::Interface || self.show_developer_options { + self.publish_messages(); + } + } + + #[cfg(target_arch = "wasm32")] + fn as_any_mut(&mut self) -> Option<&mut dyn std::any::Any> { + Some(&mut *self) + } + + fn save(&mut self, storage: &mut dyn eframe::Storage) { + eframe::set_value(storage, eframe::APP_KEY, self); + } +} + +impl App { + fn navigation_buttons(&mut self, ui: &mut egui::Ui) { + for page in crate::ENABLED_PAGES { + if self.current_page == page { + ui.add(egui::Button::new(page.title())) + .highlight() + .on_hover_text(format!("{} (current page)", page.description())); + } else { + let button = ui + .add(egui::Button::new(page.title())) + .on_hover_text(page.description()); + // If the button is clicked, change the current page + if button.clicked() { + // Change URL to the new page in the same tab + #[cfg(target_arch = "wasm32")] + crate::utils::egui::open_url_on_page(ui.ctx(), page, true); + + // Manually update the current page for non-web platforms + #[cfg(not(target_arch = "wasm32"))] + { + self.current_page = page; + } + + if page == Page::QuickStart { + self.show_about = false; + } + } else { + // Open URL in a new page if the middle mouse button is clicked + #[cfg(target_arch = "wasm32")] + if button.middle_clicked() { + crate::utils::egui::open_url_on_page(ui.ctx(), page, false); + } + } + } + } + } + + fn about_button(&mut self, ui: &mut egui::Ui) { + // About button + let button = ui + .add(egui::Button::new({ + let text = egui::RichText::new("About"); + if self.show_about { + text.strong() + } else { + text + } + })) + .on_hover_text(if self.show_about { + "Close the associated window" + } else { + "Learn more about this demo" + }); + // If the button is clicked, change the current page + if button.clicked() { + self.show_about = !self.show_about; + } + } + + fn advanced_opts_button(&mut self, ui: &mut egui::Ui) { + // Advanced options button + let button = ui + .add(egui::Button::new({ + let text = egui::RichText::new("Advanced Options"); + if self.show_about { + text.strong() + } else { + text + } + })) + .on_hover_text(if self.show_about { + "Close the associated window" + } else { + "Show advanced options" + }); + // If the button is clicked, open the advanced options window + if button.clicked() { + self.show_developer_options = !self.show_developer_options; + } + } + + fn quickstart_page(&mut self, ui: &mut egui::Ui) { + let quick_start_options: [( + egui::Theme, + &str, + crate::utils::Difficulty, + egui::ImageSource, + crate::config::TaskConfig, + ); 9] = [ + ( + egui::Theme::Dark, + "Lunar Sample Collection", + crate::utils::Difficulty::Easy, + crate::macros::include_content_image!("_images/envs/sample_collection_moon.jpg"), + crate::config::TaskConfig { + task: crate::config::Task::SampleCollection, + num_envs: 1, + env_cfg: space_robotics_bench::envs::EnvironmentConfig { + scenario: space_robotics_bench::envs::Scenario::Moon, + assets: space_robotics_bench::envs::Assets { + robot: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + object: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Procedural, + }, + terrain: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Procedural, + }, + vehicle: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + }, + seed: 1, + detail: 1.0, + }, + enable_ui: false, + }, + ), + ( + egui::Theme::Dark, + "Mars Sample Tube Collection", + crate::utils::Difficulty::Easy, + crate::macros::include_content_image!("_images/envs/sample_collection_mars.jpg"), + crate::config::TaskConfig { + task: crate::config::Task::SampleCollection, + num_envs: 1, + env_cfg: space_robotics_bench::envs::EnvironmentConfig { + scenario: space_robotics_bench::envs::Scenario::Mars, + assets: space_robotics_bench::envs::Assets { + robot: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + object: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + terrain: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Procedural, + }, + vehicle: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + }, + seed: 3, + detail: 1.0, + }, + enable_ui: false, + }, + ), + ( + egui::Theme::Light, + "Debris Capture", + crate::utils::Difficulty::Medium, + crate::macros::include_content_image!("_images/envs/debris_capture_orbit.jpg"), + crate::config::TaskConfig { + task: crate::config::Task::DebrisCapture, + num_envs: 1, + env_cfg: space_robotics_bench::envs::EnvironmentConfig { + scenario: space_robotics_bench::envs::Scenario::Orbit, + assets: space_robotics_bench::envs::Assets { + robot: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + object: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + terrain: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::None, + }, + vehicle: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + }, + seed: 2, + detail: 1.0, + }, + enable_ui: false, + }, + ), + ( + egui::Theme::Dark, + "Lunar Peg-in-Hole Assembly", + crate::utils::Difficulty::Medium, + crate::macros::include_content_image!("_images/envs/peg_in_hole_moon.jpg"), + crate::config::TaskConfig { + task: crate::config::Task::PegInHole, + num_envs: 1, + env_cfg: space_robotics_bench::envs::EnvironmentConfig { + scenario: space_robotics_bench::envs::Scenario::Moon, + assets: space_robotics_bench::envs::Assets { + robot: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + object: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + terrain: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Procedural, + }, + vehicle: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + }, + seed: 7, + detail: 1.0, + }, + enable_ui: false, + }, + ), + ( + egui::Theme::Dark, + "Mars Perseverance Rover", + crate::utils::Difficulty::Demo, + crate::macros::include_content_image!("_images/envs/perseverance.jpg"), + crate::config::TaskConfig { + task: crate::config::Task::Perseverance, + num_envs: 1, + env_cfg: space_robotics_bench::envs::EnvironmentConfig { + scenario: space_robotics_bench::envs::Scenario::Mars, + assets: space_robotics_bench::envs::Assets { + robot: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + object: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + terrain: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Procedural, + }, + vehicle: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + }, + seed: 4, + detail: 1.0, + }, + enable_ui: false, + }, + ), + ( + egui::Theme::Light, + "Peg-in-Hole Assembly", + crate::utils::Difficulty::Challenging, + crate::macros::include_content_image!("_images/envs/peg_in_hole_orbit.jpg"), + crate::config::TaskConfig { + task: crate::config::Task::PegInHole, + num_envs: 1, + env_cfg: space_robotics_bench::envs::EnvironmentConfig { + scenario: space_robotics_bench::envs::Scenario::Orbit, + assets: space_robotics_bench::envs::Assets { + robot: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + object: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + terrain: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::None, + }, + vehicle: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + }, + seed: 9, + detail: 1.0, + }, + enable_ui: false, + }, + ), + ( + egui::Theme::Dark, + "Solar Panel Assembly", + crate::utils::Difficulty::Challenging, + crate::macros::include_content_image!("_images/envs/solar_panel_assembly_moon.jpg"), + crate::config::TaskConfig { + task: crate::config::Task::SolarPanelAssembly, + num_envs: 1, + env_cfg: space_robotics_bench::envs::EnvironmentConfig { + scenario: space_robotics_bench::envs::Scenario::Moon, + assets: space_robotics_bench::envs::Assets { + robot: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + object: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + terrain: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Procedural, + }, + vehicle: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + }, + seed: 12, + detail: 1.0, + }, + enable_ui: false, + }, + ), + ( + egui::Theme::Dark, + "Ingenuity Mars Helicopter", + crate::utils::Difficulty::Demo, + crate::macros::include_content_image!("_images/envs/ingenuity.jpg"), + crate::config::TaskConfig { + task: crate::config::Task::Ingenuity, + num_envs: 1, + env_cfg: space_robotics_bench::envs::EnvironmentConfig { + scenario: space_robotics_bench::envs::Scenario::Mars, + assets: space_robotics_bench::envs::Assets { + robot: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + object: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + terrain: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Procedural, + }, + vehicle: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + }, + seed: 7, + detail: 1.0, + }, + enable_ui: false, + }, + ), + ( + egui::Theme::Light, + "Canadarm3", + crate::utils::Difficulty::Demo, + crate::macros::include_content_image!("_images/envs/gateway.jpg"), + crate::config::TaskConfig { + task: crate::config::Task::Gateway, + num_envs: 1, + env_cfg: space_robotics_bench::envs::EnvironmentConfig { + scenario: space_robotics_bench::envs::Scenario::Orbit, + assets: space_robotics_bench::envs::Assets { + robot: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + object: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + terrain: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Procedural, + }, + vehicle: space_robotics_bench::envs::Asset { + variant: space_robotics_bench::envs::AssetVariant::Dataset, + }, + }, + seed: 8, + detail: 1.0, + }, + enable_ui: false, + }, + ), + ]; + + ui.with_layout(egui::Layout::top_down(egui::Align::Center), |ui| { + egui::Grid::new("extra_nav_buttons") + .spacing(egui::vec2(8.0, 8.0)) + .show(ui, |ui| { + const N_COLS: usize = 3; + const N_ROWS: usize = 3; + let target_button_width = ui.ctx().available_rect().width() / N_COLS as f32 + - 1.25 * ui.spacing().item_spacing.x; + let target_button_height = (ui.ctx().available_rect().height() - 18.0) + / N_ROWS as f32 + - 1.5 * ui.spacing().item_spacing.y; + + let mut hovered_task = None; + + quick_start_options.into_iter().enumerate().for_each( + |(i, (theme, task, difficulty, thumbnail, config))| { + ui.with_layout(egui::Layout::top_down(egui::Align::Center), |ui| { + difficulty.set_theme(ui, self.theme); + + let button = ui.add_sized( + egui::Vec2::new(target_button_width, target_button_height), + egui::ImageButton::new(thumbnail).rounding( + 0.01 * (target_button_height + target_button_width), + ), + ); + if button.clicked() { + self.task_config = config; + self.gravity = + self.task_config.env_cfg.scenario.gravity_magnitude(); + self.start_subprocess(); + self.current_page = Page::Interface; + } + if button.hovered() { + hovered_task = Some(i); + } + + let task_text_size = 0.125 * button.rect.height(); + ui.allocate_new_ui( + egui::UiBuilder::new().max_rect(egui::Rect { + min: egui::Pos2 { + x: button.rect.min.x, + y: button.rect.max.y - 1.3 * task_text_size, + }, + max: egui::Pos2 { + x: button.rect.max.x, + y: button.rect.max.y - 0.3 * task_text_size, + }, + }), + |ui| { + ui.with_layout( + egui::Layout::bottom_up(egui::Align::Center), + |ui| { + ui.add( + egui::Label::new( + egui::RichText::new(task) + .color( + // if button.hovered() { + match theme { + egui::Theme::Light => { + Color32::from_rgb( + // 11, 12, 16, + 0, 0, 0, + ) + } + egui::Theme::Dark => { + Color32::from_rgb( + 205, 214, 244, + ) + } + }, + ) + .size(task_text_size), + ) + .selectable(false), + ) + }, + ) + }, + ); + + let difficulty_text_size = 0.1125 * button.rect.height(); + ui.allocate_new_ui( + egui::UiBuilder::new().max_rect(egui::Rect { + min: egui::Pos2 { + x: button.rect.min.x + 0.5 * difficulty_text_size, + y: button.rect.min.y + 0.4 * difficulty_text_size, + }, + max: egui::Pos2 { + x: button.rect.max.x, + y: button.rect.min.y + 1.4 * difficulty_text_size, + }, + }), + |ui| { + ui.with_layout( + egui::Layout::left_to_right(egui::Align::TOP), + |ui| { + ui.add( + egui::Label::new( + egui::RichText::new(difficulty.to_string()) + .color(difficulty.get_text_color( + if self.theme == egui::Theme::Dark { + button.hovered() + } else { + !button.hovered() + }, + )) + .size(difficulty_text_size), + ) + .selectable(false), + ) + }, + ) + }, + ); + }); + + if (i + 1) % N_COLS == 0 { + ui.end_row(); + } + }, + ); + + self.hovered_task = hovered_task; + }); + }); + } + + fn configuration_page(&mut self, ui: &mut egui::Ui) { + egui::ScrollArea::vertical() + .show(ui, |ui| { + let margin_x = (ui.ctx().screen_rect().width() - 768.0).max(0.0) / 2.0; + let inner_margin = egui::Margin { + left: margin_x.max(0.0), + right: margin_x.max(0.0), + ..egui::Margin::default() + }; + egui::Frame::default() + .inner_margin(inner_margin) + .show(ui, |ui| + { + ui.with_layout(egui::Layout::top_down(egui::Align::Center), |ui| { + ui.with_layout( + egui::Layout::left_to_right(egui::Align::TOP) + , |ui| { + ui.spacing_mut().item_spacing.x = 0.0; + + let task_focus=match self.task_config.task { + crate::config::Task::SampleCollection + | crate::config::Task::DebrisCapture + | crate::config::Task::PegInHole + | crate::config::Task::SolarPanelAssembly => + { + "Objective" + } + _ => { + "Demo" + } + }; + ui.add( + egui::Label::new( + egui::RichText::new(format!("{task_focus}\n")).strong().size(24.0), + ) + .selectable(false), + ); + + match self.task_config.task { + crate::config::Task::SampleCollection => { + let sample_type = match self.task_config.env_cfg.scenario { + space_robotics_bench::envs::Scenario::Moon => "Lunar ", + space_robotics_bench::envs::Scenario::Mars => "Martian ", + _ => "", + }; + + ui.add( + egui::Label::new(egui::RichText::new(&format!("\t1) Grasp the {sample_type}sample.\n\t2) [Optional] Place it in the rover cargo bay.")).monospace()).selectable(false) + ); + } + crate::config::Task::DebrisCapture => { + ui.add( + egui::Label::new(egui::RichText::new("\t— Detumble and grasp the \"misplaced\" object.").monospace()).selectable(false) + ); + } + crate::config::Task::PegInHole => { + ui.add( + egui::Label::new(egui::RichText::new("\t1) Grasp the profile.\n\t2) Align the profile with the hole.\n\t3) Insert the profile into the hole.").monospace()).selectable(false) + ); + } + crate::config::Task::SolarPanelAssembly => { + ui.add( + egui::Label::new(egui::RichText::new("\t1) Grasp, align, and insert all profiles into their holes.\n\t2) Place the solar panel on top of the profiles.").monospace()).selectable(false) + ); + } + crate::config::Task::Perseverance => { + ui.add( + egui::Label::new(egui::RichText::new("\t— Explore the Martian terrain with the Perseverance rover.").monospace()).selectable(false) + ); + } + crate::config::Task::Ingenuity => { + ui.add( + egui::Label::new(egui::RichText::new("\t— Explore the Martian terrain with the Ingenuity helicopter.").monospace()).selectable(false) + ); + } + crate::config::Task::Gateway => { + ui.add( + egui::Label::new(egui::RichText::new("\t— Explore the Gateway space station with the Canadarm3 robotic arm.").monospace()).selectable(false) + ); + } + } + }); + + + ui.add_space(25.0); + + egui::Grid::new("real_time_env_config").show(ui, |ui| { + ui.style_mut().spacing.slider_width = 170.0; + + ui.add(egui::Label::new(egui::RichText::new("\u{e80b} Gravity").size(22.0)).selectable(false)).on_hover_text( + "Acceleration due to gravity" + ); + ui.add(egui::Slider::new(&mut self.gravity, 0.0..=25.0) + .trailing_fill(true).custom_formatter(|x, _| format!("{x:.2}")).suffix(" m/s²").custom_parser( + |x| if let Ok(x) = x.parse::() { + Some(x.max(0.0_f64)) + } else { + None + } + )); + + ui.with_layout(egui::Layout::left_to_right(egui::Align::Center), |ui| { + if ui.add( + egui::Button::new(egui::RichText::new("Zero") + ).frame(true)).clicked() { + self.gravity = space_robotics_bench::envs::Scenario::Orbit.gravity_magnitude(); + } + if ui.add( + egui::Button::new(egui::RichText::new("Moon") + ).frame(true)).clicked() { + self.gravity = space_robotics_bench::envs::Scenario::Moon.gravity_magnitude(); + } + if ui.add( + egui::Button::new(egui::RichText::new("Mars") + ).frame(true)).clicked() { + self.gravity = space_robotics_bench::envs::Scenario::Mars.gravity_magnitude(); + } + if ui.add( + egui::Button::new(egui::RichText::new("Earth") + ).frame(true)).clicked() { + self.gravity = 9.81; + } + if ui.add( + egui::Button::new(egui::RichText::new("Jupiter") + ).frame(true)).clicked() { + self.gravity = 24.79; + } + }); + ui.end_row(); + + ui.add(egui::Label::new(egui::RichText::new("\u{e8b5} Latency").size(22.0)).selectable(false)).on_hover_text( + "One-way communication delay between the operator and the robot" + ); + ui.add(egui::Slider::new(&mut self.latency, 0.0..=2000.0) + .trailing_fill(true).custom_formatter(|x, _| format!("{x:.0}")).suffix(" ms")); + ui.with_layout(egui::Layout::left_to_right(egui::Align::Center), |ui| { + if ui.add( + egui::Button::new(egui::RichText::new("None") + ).frame(true)).clicked() { + self.latency = 0.0; + } + if ui.add( + egui::Button::new(egui::RichText::new("LEO") + ).frame(true)).clicked() { + self.latency = 50.0; + } + if ui.add( + egui::Button::new(egui::RichText::new("MEO") + ).frame(true)).clicked() { + self.latency = 225.0; + } + if ui.add( + egui::Button::new(egui::RichText::new("GEO") + ).frame(true)).clicked() { + self.latency = 280.0; + } + if ui.add( + egui::Button::new(egui::RichText::new("Moon") + ).frame(true)).clicked() { + self.latency = 1250.0; + } + ui.end_row(); + }); + + ui.end_row(); + }); + + ui.add_space(20.0); + + ui.with_layout( + egui::Layout::bottom_up(egui::Align::Center), + |ui| { + let max_button_width = ui.ctx().available_rect().width().min(768.0); + + if ui.add(egui::Button::new(egui::RichText::new("\u{f23a} Exit").size(32.0)).min_size( + egui::Vec2::new(max_button_width, 75.0), + ).frame(true)).clicked() { + self.restart_episode(); + self.current_page = Page::QuickStart; + self.show_about = false; + self.gravity = self.task_config.env_cfg.scenario.gravity_magnitude(); + self.latency = 0.0; + self.motion_sensitivity = 1.0; + self.force_feedback_sensitivity = 1.0; + self.max_feedback_force = 2.0; + self.collect_trajectory = false; + self.task_config = crate::config::TaskConfig::default(); + self.stop_subprocess(); + } + + ui.add_space(10.0); + + if ui.add(egui::Button::new(egui::RichText::new("\u{e5d5} Restart").size(32.0)).min_size( + egui::Vec2::new(max_button_width, 75.0), + ).frame(true)).clicked() { + self.restart_episode(); + } + }); + }); + }); + }); + } + + fn restart_episode(&mut self) { + if self.subprocess.is_some() { + info!("Restarting episode"); + if self.collect_trajectory { + info!("Saving trajectory"); + self.pub_reset_save_dataset.publish(&EmptyMsg {}).unwrap(); + self.n_collected_trajectories += 1; + + writeln!( + self.logfile, + "DATA, {}, {}, {}", + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(), + chrono::Local::now().format("%Y-%m-%d %H:%M:%S"), + self.task_config + ) + .unwrap(); + } else { + self.pub_reset_discard_dataset + .publish(&EmptyMsg {}) + .unwrap(); + } + } else { + error!("Cannot restart episode: subprocess is not running"); + } + } + + fn about_window(&mut self, ui: &mut egui::Ui) { + if self.show_about { + let available_rect = ui.ctx().available_rect(); + let center_point = available_rect.center(); + + egui::containers::Window::new(egui::RichText::new("About").size(18.0)) + .interactable(true) + .open(&mut self.show_about) + .collapsible(false) + .resizable(false) + .fixed_size([512.0, 1024.0]) + .default_rect(egui::Rect { + min: egui::Pos2::new( + center_point.x - 512.0 / 2.0, + center_point.y - 1024.0 / 2.0, + ), + max: egui::Pos2::new( + center_point.x + 512.0 / 2.0, + center_point.y + 1024.0 / 2.0, + ), + }) + .show(ui.ctx(), |ui| { + ui.with_layout(egui::Layout::top_down(egui::Align::Center), |ui| { + commonmark_str!( + ui, + &mut self.commonmark_cache, + "crates/space_robotics_bench_gui/content/about.md" + ); + }); + }); + } + } + + fn virtual_keyboard_window(&mut self, ui: &mut egui::Ui) { + if self.show_virtual_keyboard_window && self.current_page == Page::Interface { + let available_rect = ui.ctx().available_rect(); + let available_size = available_rect.size(); + + egui::containers::Window::new(egui::RichText::new("Gripper").size(16.0)) + .interactable(true) + .collapsible(true) + .resizable(false) + .max_size([0.61 * available_size.x, 0.61 * available_size.y]) + .default_rect(egui::Rect { + min: egui::Pos2::new(available_rect.max.x, available_rect.min.y), + max: egui::Pos2::new(available_rect.max.x, available_rect.min.y), + }) + .show(ui.ctx(), |ui| { + ui.with_layout(egui::Layout::left_to_right(egui::Align::TOP), |ui| { + if ui + .add( + egui::Button::new( + egui::RichText::new("Toggle\nGripper") + .heading() + .size(32.0), + ) + .frame(true), + ) + .on_hover_text("\u{e811} The button stopped working \u{e811}\n (this is a workaround)") + .clicked() + { + self.pub_gripper_toggle + .publish(&BoolMsg { data: true }) + .unwrap(); + } + }); + }); + } + } + + fn developer_options_window(&mut self, ui: &mut egui::Ui) { + if self.show_developer_options { + let available_rect = ui.ctx().available_rect(); + egui::containers::Window::new(egui::RichText::new("Customization").size(22.0)) + .interactable(true) + .collapsible(false) + .resizable(false) + .default_rect(egui::Rect { + min: available_rect.max, + max: available_rect.max, + }) + .fixed_size([384.0, 256.0]) + .show(ui.ctx(), |ui| { + ui.with_layout(egui::Layout::top_down(egui::Align::LEFT), |ui| { + egui::Grid::new("extra_developer_options_grid").show(ui, |ui| { + ui.add( + egui::Label::new( + egui::RichText::new("\u{ef75} Scenario").size(20.0), + ) + .selectable(false), + ); + egui::ComboBox::new("dev_scenario_combo_box", "") + .width(235.0) + .selected_text( + egui::RichText::new(match self.task_config.env_cfg.scenario { + space_robotics_bench::envs::Scenario::Asteroid => { + "Asteroid" + } + space_robotics_bench::envs::Scenario::Earth => "Earth", + space_robotics_bench::envs::Scenario::Mars => "Mars", + space_robotics_bench::envs::Scenario::Moon => "Moon", + space_robotics_bench::envs::Scenario::Orbit => "Orbit", + }) + .size(20.0), + ) + .show_ui(ui, |ui| { + ui.selectable_value( + &mut self.task_config.env_cfg.scenario, + space_robotics_bench::envs::Scenario::Moon, + "Moon", + ); + ui.selectable_value( + &mut self.task_config.env_cfg.scenario, + space_robotics_bench::envs::Scenario::Mars, + "Mars", + ); + ui.selectable_value( + &mut self.task_config.env_cfg.scenario, + space_robotics_bench::envs::Scenario::Orbit, + "Orbit", + ); + }); + + ui.end_row(); + + ui.add( + egui::Label::new( + egui::RichText::new(format!( + "\u{eb9b} {}", + match self.task_config.task { + crate::config::Task::SampleCollection + | crate::config::Task::DebrisCapture + | crate::config::Task::PegInHole + | crate::config::Task::SolarPanelAssembly => "Task", + crate::config::Task::Perseverance + | crate::config::Task::Ingenuity + | crate::config::Task::Gateway => "Demo", + } + )) + .size(20.0), + ) + .selectable(false), + ); + egui::ComboBox::new("dev_task_combo_box", "") + .width(235.0) + .selected_text( + egui::RichText::new(match self.task_config.task { + crate::config::Task::SampleCollection => { + "Sample Collection" + } + crate::config::Task::PegInHole => "Peg-in-Hole", + crate::config::Task::SolarPanelAssembly => { + "Solar Panel Assembly" + } + crate::config::Task::DebrisCapture => "Debris Capture", + crate::config::Task::Perseverance => "Perseverance", + crate::config::Task::Ingenuity => "Ingenuity", + crate::config::Task::Gateway => "Gateway", + }) + .size(20.0), + ) + .show_ui(ui, |ui| { + ui.selectable_value( + &mut self.task_config.task, + crate::config::Task::SampleCollection, + "Sample Collection", + ); + ui.selectable_value( + &mut self.task_config.task, + crate::config::Task::PegInHole, + "Peg-in-Hole", + ); + ui.selectable_value( + &mut self.task_config.task, + crate::config::Task::SolarPanelAssembly, + "Solar Panel Assembly", + ); + ui.selectable_value( + &mut self.task_config.task, + crate::config::Task::DebrisCapture, + "Debris Capture", + ); + ui.separator(); + ui.selectable_value( + &mut self.task_config.task, + crate::config::Task::Perseverance, + "Perseverance", + ); + ui.selectable_value( + &mut self.task_config.task, + crate::config::Task::Ingenuity, + "Ingenuity", + ); + ui.selectable_value( + &mut self.task_config.task, + crate::config::Task::Gateway, + "Gateway", + ); + }); + + ui.end_row(); + + ui.add( + egui::Label::new( + egui::RichText::new("\u{ea3c} Object Assets").size(20.0), + ) + .selectable(false), + ); + egui::ComboBox::new("dev_object_asset_combo_box", "") + .width(235.0) + .selected_text( + egui::RichText::new( + match self.task_config.env_cfg.assets.object.variant { + space_robotics_bench::envs::AssetVariant::None => { + "None" + }, + space_robotics_bench::envs::AssetVariant::Primitive => { + "Primitive" + }, + space_robotics_bench::envs::AssetVariant::Dataset => { + "Dataset" + }, + space_robotics_bench::envs::AssetVariant::Procedural => { + "Procedural" + } + }, + ) + .size(20.0), + ) + .show_ui(ui, |ui| { + ui.selectable_value( + &mut self.task_config.env_cfg.assets.object.variant, + space_robotics_bench::envs::AssetVariant::Primitive, + "Primitive", + ); + ui.selectable_value( + &mut self.task_config.env_cfg.assets.object.variant, + space_robotics_bench::envs::AssetVariant::Dataset, + "Dataset", + ); + ui.selectable_value( + &mut self.task_config.env_cfg.assets.object.variant, + space_robotics_bench::envs::AssetVariant::Procedural, + "Procedural", + ); + }); + + ui.end_row(); + + ui.style_mut().spacing.slider_width = 185.0; + + ui.add( + egui::Label::new( + egui::RichText::new("\u{e9e1} Random seed").size(20.0), + ) + .selectable(false), + ); + ui.add( + egui::Slider::new(&mut self.task_config.env_cfg.seed, 0..=42) + + .trailing_fill(true) + .integer() + .custom_formatter(|x, _| format!("{x}")), + ); + + ui.end_row(); + + ui.add( + egui::Label::new( + egui::RichText::new("\u{e3ea} Parallel envs").size(20.0), + ) + .selectable(false), + ); + ui.add( + egui::Slider::new(&mut self.task_config.num_envs, 1..=16) + .clamping(egui::SliderClamping::Always) + .trailing_fill(true) + .integer() + .custom_formatter(|x, _| format!("{x}")), + ); + + ui.end_row(); + + ui.add( + egui::Label::new( + egui::RichText::new("\u{e839} Level of detail").size(20.0), + ) + .selectable(false), + ); + ui.add( + egui::Slider::new(&mut self.task_config.env_cfg.detail, 0.01..=2.0) + .clamping(egui::SliderClamping::Always) + .trailing_fill(true) + .custom_formatter(|x, _| format!("{x:.1}")), + ); + + ui.end_row(); + + ui.add( + egui::Label::new( + egui::RichText::new("\u{e1bd} Enable UI").size(20.0), + ) + .selectable(false), + ); + let is_enabled = self.task_config.enable_ui; + ui.add(egui::Checkbox::new( + &mut self.task_config.enable_ui, + egui::RichText::new(if is_enabled { + " Enabled" + } else { + " Disabled" + }) + .size(20.0), + )); + + ui.end_row(); + ui.end_row(); + + ui.add( + egui::Label::new( + egui::RichText::new("\u{e766} Feedback sensitivity").size(20.0), + ) + .selectable(false), + ); + ui.add( + egui::Slider::new(&mut self.force_feedback_sensitivity, 0.0..=20.0) + .clamping(egui::SliderClamping::Always) + .trailing_fill(true) + .custom_formatter(|x, _| format!("{x:.1}")), + ); + + ui.end_row(); + + ui.add( + egui::Label::new( + egui::RichText::new("\u{e9e4} Feedback force limit").size(20.0), + ) + .selectable(false), + ); + ui.add( + egui::Slider::new(&mut self.max_feedback_force, 0.0..=5.0) + .clamping(egui::SliderClamping::Always) + .trailing_fill(true) + .custom_formatter(|x, _| format!("{x:.1}")), + ); + + ui.end_row(); + + ui.add( + egui::Label::new( + egui::RichText::new("\u{f049} Virtual gripper").size(20.0), + ) + .selectable(false), + ); + let is_enabled = self.show_virtual_keyboard_window; + ui.add(egui::Checkbox::new( + &mut self.show_virtual_keyboard_window, + egui::RichText::new(if is_enabled { + " Enabled" + } else { + " Disabled" + }) + .size(20.0), + )); + }); + + // ui.add_space(10.0); + + ui.separator(); + + ui.with_layout( + egui::Layout::centered_and_justified(egui::Direction::TopDown), + |ui| { + if ui + .add( + egui::Button::new( + egui::RichText::new("\u{e1c4}").size(40.0), + ) + .frame(false), + ) + .clicked() + { + self.gravity = + self.task_config.env_cfg.scenario.gravity_magnitude(); + self.start_subprocess(); + self.current_page = Page::Interface; + } + }, + ); + }); + }); + } + } + + fn start_subprocess(&mut self) { + const PYTHON_SCRIPT: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../../scripts/teleop.py"); + + self.stop_subprocess(); + if self.subprocess.is_some() { + warn!("Subprocess already running."); + return; + } + + // Construct the subprocess + println!("python: {}", Self::get_isaacsim_python_exe()); + let exec = subprocess::Exec::cmd(Self::get_isaacsim_python_exe()).arg(PYTHON_SCRIPT); + let exec = self.task_config.set_exec_env(exec); + + // Start the subprocess + info!("Starting subprocess..."); + let mut popen = exec.popen().unwrap(); + + popen + .wait_timeout(std::time::Duration::from_millis(2000)) + .unwrap(); + + if popen.poll().is_none() { + info!("Subprocess started."); + self.subprocess = Some(popen); + } else { + error!("Failed to start subprocess."); + } + } + + /// Note: The `python.sh` interface of isaac sim launches a new process `python3` but does not propagate the signal to the child process. + /// Therefore, we need kill the `python3` process directly. + fn stop_subprocess(&mut self) { + self.pub_gracefully_shutdown_process + .publish(&EmptyMsg {}) + .unwrap(); + + const SUBPROCESS_NAME: &str = "python3"; + const SLEEP_DURATION: std::time::Duration = std::time::Duration::from_millis(10); + + if let Some(p) = &mut self.subprocess { + info!("Stopping subprocess..."); + + if p.wait_timeout(std::time::Duration::from_millis(1000)) + .unwrap() + .is_some() + { + info!("Subprocess stopped."); + self.subprocess = None; + return; + } + + // Try to terminate the process gracefully + for signal in [ + nix::sys::signal::Signal::SIGINT, + nix::sys::signal::Signal::SIGTERM, + nix::sys::signal::Signal::SIGKILL, + ] { + loop { + let _ = Self::kill_all_processes_by_name(SUBPROCESS_NAME, signal); + + if p.wait_timeout(SLEEP_DURATION).unwrap().is_some() { + info!("Subprocess stopped."); + self.subprocess = None; + return; + } + + if signal != nix::sys::signal::Signal::SIGKILL { + break; + } + } + } + } else { + warn!("No known subprocess to stop."); + // Directly kill the process if the subprocess might have been spawned by other means + let _ = Self::kill_all_processes_by_name( + SUBPROCESS_NAME, + nix::sys::signal::Signal::SIGKILL, + ); + } + } + + fn get_isaacsim_python_exe() -> String { + if let Ok(python_exe) = std::env::var("ISAAC_SIM_EXE") { + trace!("ISAAC_SIM_EXE: {}", python_exe); + python_exe.trim().to_owned() + } else { + let home_dir = home::home_dir().unwrap_or("/root".into()); + let isaac_sim_python_sh = home_dir.join("isaac-sim/python.sh"); + if std::path::Path::new(&isaac_sim_python_sh).exists() { + trace!("ISAAC_SIM_EXE: {}", isaac_sim_python_sh.display()); + isaac_sim_python_sh.display().to_string() + } else { + trace!("ISAAC_SIM_EXE: which python3"); + subprocess::Exec::cmd("which") + .arg("python3") + .stdout(subprocess::Redirection::Pipe) + .stderr(subprocess::Redirection::Merge) + .capture() + .expect("No Python interpreter was found.") + .stdout_str() + .trim() + .to_owned() + } + } + } + + fn kill_all_processes_by_name(name: &str, signal: nix::sys::signal::Signal) -> nix::Result<()> { + // Create a System object to gather information + let mut system = sysinfo::System::new_all(); + // Refresh the process list + system.refresh_all(); + + // Iterate through all processes + for process in system.processes() { + // Check if the process name matches the given name + if process.1.name() == name { + let pid = process.1.pid(); + // Send a signal to the process + trace!("Killing process {} with signal {:?}", pid, signal); + nix::sys::signal::kill(nix::unistd::Pid::from_raw(pid.as_u32() as i32), signal)?; + } + } + + Ok(()) + } + + pub fn dark_mode_toggle_button(&mut self, ui: &mut egui::Ui) { + let (icon, tooltip, target_visuals) = match self.theme { + egui::Theme::Dark => ( + "\u{e51c}", + "Switch to light mode", + crate::style::light_visuals(), + ), + egui::Theme::Light => ( + "\u{e518}", + "Switch to dark mode", + crate::style::dark_visuals(), + ), + }; + + if ui + .add(egui::Button::new(icon)) + .on_hover_text(tooltip) + .clicked() + { + ui.ctx().set_visuals(target_visuals.to_owned()); + self.theme = match self.theme { + egui::Theme::Dark => egui::Theme::Light, + egui::Theme::Light => egui::Theme::Dark, + }; + } + } + + fn publish_messages(&mut self) { + self.gravity = self.gravity.max(0.0); + self.latency = self.latency.max(0.0); + self.motion_sensitivity = self.motion_sensitivity.max(0.0); + self.force_feedback_sensitivity = self.force_feedback_sensitivity.max(0.0); + self.max_feedback_force = self.max_feedback_force.clamp(0.0, 5.0); + + if self.prev_gravity != self.gravity { + self.prev_gravity = self.gravity; + self.pub_gravity + .publish(&Float64Msg { data: self.gravity }) + .unwrap(); + } + + if self.prev_latency != self.latency { + self.prev_latency = self.latency; + self.pub_latency + .publish(&Float64Msg { + data: self.latency / 1000.0, + }) + .unwrap(); + } + + if self.prev_motion_sensitivity != self.motion_sensitivity { + self.prev_motion_sensitivity = self.motion_sensitivity; + self.pub_motion_sensitivity + .publish(&Float64Msg { + data: self.motion_sensitivity, + }) + .unwrap(); + } + + if self.prev_force_feedback_sensitivity != self.force_feedback_sensitivity { + self.prev_force_feedback_sensitivity = self.force_feedback_sensitivity; + self.pub_force_feedback_sensitivity + .publish(&Float64Msg { + data: self.force_feedback_sensitivity, + }) + .unwrap(); + } + + if self.prev_max_feedback_force != self.max_feedback_force { + self.prev_max_feedback_force = self.max_feedback_force; + self.pub_max_feedback_force + .publish(&Float64Msg { + data: self.max_feedback_force, + }) + .unwrap(); + } + } + + fn show_top_center_bar(&mut self, ui: &mut egui::Ui) { + ui.with_layout( + egui::Layout::centered_and_justified(egui::Direction::LeftToRight), + |ui| match self.current_page { + Page::QuickStart => { + ui.add( + egui::Label::new( + egui::RichText::new("\u{eb9b} Select your Experience \u{ef75} ") + .family(egui::FontFamily::Proportional) + .strong(), + ) + .selectable(false), + ); + } + Page::Interface => { + ui.add( + egui::Label::new( + egui::RichText::new("\u{f049} Complete the Task \u{ea3c} ") + .family(egui::FontFamily::Proportional) + .strong(), + ) + .selectable(false), + ); + } + }, + ); + } + + fn show_trajectory_collection_checkbox(&mut self, ui: &mut egui::Ui) { + ui.add(egui::Checkbox::new( + &mut self.collect_trajectory, + egui::RichText::new(format!( + " Collect Trajectory ( {} samples )", + self.n_collected_trajectories + )) + .heading() + .size(20.0), + )) + .on_hover_text("Your participation makes our robots more intelligent!"); + } + + fn warn_if_debug_build(&mut self, ui: &mut egui::Ui) { + if cfg!(debug_assertions) { + ui.separator(); + ui.add(egui::Button::new("⚠ Debug build ⚠")) + .on_hover_ui(|ui| { + ui.label( + egui::RichText::new(format!( + "Current page: {:?}\n\ + Screen size: {:?}\n\ + ", + self.current_page, + ui.ctx().screen_rect().size(), + )) + .font(egui::FontId::monospace( + ui.style().text_styles[&egui::TextStyle::Button].size, + )), + ); + }); + ui.separator(); + } + } +} diff --git a/crates/space_robotics_bench_gui/src/config.rs b/crates/space_robotics_bench_gui/src/config.rs new file mode 100644 index 0000000..8150e18 --- /dev/null +++ b/crates/space_robotics_bench_gui/src/config.rs @@ -0,0 +1,185 @@ +use serde::{Deserialize, Serialize}; +use space_robotics_bench::envs::EnvironmentConfig; + +const ENVIRON_PREFIX: &str = "SRB_"; + +#[derive( + Deserialize, + Serialize, + Debug, + display_json::DisplayAsJson, + Clone, + Copy, + PartialEq, + Eq, + Hash, + Default, +)] +#[serde(rename_all = "snake_case")] +pub enum Task { + #[default] + SampleCollection, + DebrisCapture, + PegInHole, + SolarPanelAssembly, + Perseverance, + Ingenuity, + Gateway, +} + +#[derive( + Deserialize, + Serialize, + display_json::DebugAsJson, + display_json::DisplayAsJson, + Clone, + Copy, + PartialEq, +)] +pub struct TaskConfig { + pub task: Task, + pub num_envs: u64, + pub env_cfg: EnvironmentConfig, + pub enable_ui: bool, +} + +impl Default for TaskConfig { + fn default() -> Self { + Self { + task: Task::SampleCollection, + num_envs: 1, + env_cfg: EnvironmentConfig::default(), + enable_ui: false, + } + } +} + +impl TaskConfig { + pub fn set_exec_env(mut self, mut exec: subprocess::Exec) -> subprocess::Exec { + // Arguments + if self.enable_ui { + exec = exec.args(&[ + "--task", + &format!("{}_visual", self.task.to_string().trim_matches('"')), + "--with_ros2", + ]); + } else { + exec = exec.args(&["--task", self.task.to_string().trim_matches('"')]); + } + self.num_envs = self.num_envs.max(1); + exec = exec.args(&["--num_envs", self.num_envs.to_string().as_str()]); + exec = exec.args(&["--teleop_device", "keyboard", "spacemouse", "touch", "ros2"]); + if !self.enable_ui { + exec = exec.arg("--disable_ui"); + } + + // Environment variables - Environment + exec = exec.env( + const_format::concatcp!(ENVIRON_PREFIX, "SEED"), + std::env::var(const_format::concatcp!(ENVIRON_PREFIX, "SEED")) + .unwrap_or(self.env_cfg.seed.to_string().trim_matches('"').to_owned()), + ); + exec = exec.env( + const_format::concatcp!(ENVIRON_PREFIX, "SCENARIO"), + std::env::var(const_format::concatcp!(ENVIRON_PREFIX, "SCENARIO")).unwrap_or( + self.env_cfg + .scenario + .to_string() + .trim_matches('"') + .to_owned(), + ), + ); + exec = exec.env( + const_format::concatcp!(ENVIRON_PREFIX, "DETAIL"), + std::env::var(const_format::concatcp!(ENVIRON_PREFIX, "DETAIL")) + .unwrap_or(self.env_cfg.detail.to_string().trim_matches('"').to_owned()), + ); + exec = exec.env( + const_format::concatcp!(ENVIRON_PREFIX, "ASSETS_ROBOT_VARIANT"), + std::env::var(const_format::concatcp!( + ENVIRON_PREFIX, + "ASSETS_ROBOT_VARIANT" + )) + .unwrap_or( + self.env_cfg + .assets + .robot + .variant + .to_string() + .trim_matches('"') + .to_owned(), + ), + ); + exec = exec.env( + const_format::concatcp!(ENVIRON_PREFIX, "ASSETS_OBJECT_VARIANT"), + std::env::var(const_format::concatcp!( + ENVIRON_PREFIX, + "ASSETS_OBJECT_VARIANT" + )) + .unwrap_or( + self.env_cfg + .assets + .object + .variant + .to_string() + .trim_matches('"') + .to_owned(), + ), + ); + exec = exec.env( + const_format::concatcp!(ENVIRON_PREFIX, "ASSETS_TERRAIN_VARIANT"), + std::env::var(const_format::concatcp!( + ENVIRON_PREFIX, + "ASSETS_TERRAIN_VARIANT" + )) + .unwrap_or( + self.env_cfg + .assets + .terrain + .variant + .to_string() + .trim_matches('"') + .to_owned(), + ), + ); + exec = exec.env( + const_format::concatcp!(ENVIRON_PREFIX, "ASSETS_VEHICLE_VARIANT"), + std::env::var(const_format::concatcp!( + ENVIRON_PREFIX, + "ASSETS_VEHICLE_VARIANT" + )) + .unwrap_or( + self.env_cfg + .assets + .vehicle + .variant + .to_string() + .trim_matches('"') + .to_owned(), + ), + ); + + // Environment variables - GUI + exec = exec.env( + "DISPLAY", + std::env::var(const_format::concatcp!(ENVIRON_PREFIX, "DISPLAY")) + .unwrap_or(":0".to_string()), + ); + + // Environment variables - ROS + exec = exec.env( + "ROS_DOMAIN_ID", + std::env::var("ROS_DOMAIN_ID").unwrap_or("0".to_string()), + ); + exec = exec.env( + "RMW_IMPLEMENTATION", + std::env::var(const_format::concatcp!( + ENVIRON_PREFIX, + "RMW_IMPLEMENTATION" + )) + .unwrap_or("rmw_cyclonedds_cpp".to_string()), + ); + + exec + } +} diff --git a/crates/space_robotics_bench_gui/src/consts.rs b/crates/space_robotics_bench_gui/src/consts.rs new file mode 100644 index 0000000..fbc2134 --- /dev/null +++ b/crates/space_robotics_bench_gui/src/consts.rs @@ -0,0 +1,3 @@ +use crate::page::Page; + +pub const ENABLED_PAGES: [Page; 2] = [Page::QuickStart, Page::Interface]; diff --git a/crates/space_robotics_bench_gui/src/lib.rs b/crates/space_robotics_bench_gui/src/lib.rs new file mode 100644 index 0000000..2bb3376 --- /dev/null +++ b/crates/space_robotics_bench_gui/src/lib.rs @@ -0,0 +1,10 @@ +pub use app::App; +pub use consts::*; + +mod app; +mod config; +mod consts; +mod macros; +mod page; +mod style; +mod utils; diff --git a/crates/space_robotics_bench_gui/src/macros/fonts.rs b/crates/space_robotics_bench_gui/src/macros/fonts.rs new file mode 100644 index 0000000..48e6df1 --- /dev/null +++ b/crates/space_robotics_bench_gui/src/macros/fonts.rs @@ -0,0 +1,47 @@ +macro_rules! generate_font_definitions { + {$($family_name:literal $(as $font_family:ident)? [$($name:literal),+ $(,)?]),* $(,)?} => {{ + let mut font_data = ::std::collections::BTreeMap::new(); + let mut families = ::std::collections::BTreeMap::new(); + let mut monospace_fonts = ::std::vec::Vec::new(); + let mut proportional_fonts = ::std::vec::Vec::new(); + $( + $crate::macros::generate_font_definitions!(@insert_family |font_data, families, monospace_fonts, proportional_fonts| $family_name $(as $font_family)? [$($name),+]); + )* + families.insert( + ::egui::FontFamily::Monospace, + monospace_fonts, + ); + families.insert( + ::egui::FontFamily::Proportional, + proportional_fonts, + ); + ::egui::FontDefinitions { + font_data, + families, + } + }}; + (@insert_family |$font_data:ident, $families:ident, $monospace_fonts:ident, $proportional_fonts:ident| $family_name:literal as Monospace [$($name:literal),+ $(,)?]) => { + $crate::macros::generate_font_definitions!(@insert_family |$font_data, $families, $monospace_fonts, $proportional_fonts| $family_name [$($name),+]); + $monospace_fonts.extend_from_slice(&[$($name.to_owned()),+]); + }; + (@insert_family |$font_data:ident, $families:ident, $monospace_fonts:ident, $proportional_fonts:ident| $family_name:literal as Proportional [$($name:literal),+ $(,)?]) => { + $crate::macros::generate_font_definitions!(@insert_family |$font_data, $families, $monospace_fonts, $proportional_fonts| $family_name [$($name),+]); + $proportional_fonts.extend_from_slice(&[$($name.to_owned()),+]); + }; + (@insert_family |$font_data:ident, $families:ident, $monospace_fonts:ident, $proportional_fonts:ident| $family_name:literal [$($name:literal),+ $(,)?]) => { + $( + $crate::macros::generate_font_definitions!(@insert_data |$font_data| $name); + )* + $families.insert( + ::egui::FontFamily::Name($family_name.into()), + vec![$($name.to_owned()),+], + ); + }; + (@insert_data |$font_data:ident| $name:literal) => { + $font_data.insert( + $name.to_owned(), + $crate::macros::include_assets_font!($name), + ); + }; +} +pub(crate) use generate_font_definitions; diff --git a/crates/space_robotics_bench_gui/src/macros/include.rs b/crates/space_robotics_bench_gui/src/macros/include.rs new file mode 100644 index 0000000..f932112 --- /dev/null +++ b/crates/space_robotics_bench_gui/src/macros/include.rs @@ -0,0 +1,24 @@ +macro_rules! include_content_image { + ($file:expr $(,)?) => { + ::egui::ImageSource::Bytes { + uri: ::std::borrow::Cow::Borrowed(concat!("bytes://", concat!("content/", $file))), + bytes: ::egui::load::Bytes::Static(include_bytes!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/content/", + $file + ))), + } + }; +} + +macro_rules! include_assets_font { + ($file:expr $(,)?) => { + ::egui::FontData::from_static(include_bytes!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/assets/fonts/", + $file + ))) + }; +} + +pub(crate) use {include_assets_font, include_content_image}; diff --git a/crates/space_robotics_bench_gui/src/macros/mod.rs b/crates/space_robotics_bench_gui/src/macros/mod.rs new file mode 100644 index 0000000..a513815 --- /dev/null +++ b/crates/space_robotics_bench_gui/src/macros/mod.rs @@ -0,0 +1,5 @@ +mod fonts; +mod include; + +pub(crate) use fonts::*; +pub(crate) use include::*; diff --git a/crates/space_robotics_bench_gui/src/main.rs b/crates/space_robotics_bench_gui/src/main.rs new file mode 100644 index 0000000..2895f9b --- /dev/null +++ b/crates/space_robotics_bench_gui/src/main.rs @@ -0,0 +1,28 @@ +fn main() -> eframe::Result<()> { + tracing_subscriber::fmt::init(); + + let icon = image::load_from_memory_with_format( + include_bytes!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../docs/theme/favicon.png" + )), + image::ImageFormat::Png, + ) + .unwrap() + .to_rgba8(); + let (icon_width, icon_height) = icon.dimensions(); + let native_options = eframe::NativeOptions { + viewport: egui::ViewportBuilder::default().with_icon(egui::IconData { + rgba: icon.into_raw(), + width: icon_width, + height: icon_height, + }), + ..Default::default() + }; + + eframe::run_native( + "Space Robotics Bench", + native_options, + Box::new(|cc| Ok(Box::new(space_robotics_bench_gui::App::new(cc)))), + ) +} diff --git a/crates/space_robotics_bench_gui/src/page.rs b/crates/space_robotics_bench_gui/src/page.rs new file mode 100644 index 0000000..d1cebbc --- /dev/null +++ b/crates/space_robotics_bench_gui/src/page.rs @@ -0,0 +1,33 @@ +#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, serde::Deserialize, serde::Serialize)] +pub enum Page { + QuickStart, + Interface, +} + +impl Default for Page { + fn default() -> Self { + Self::QuickStart + } +} + +impl std::fmt::Display for Page { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{self:?}") + } +} + +impl Page { + pub fn title(&self) -> &str { + match self { + Self::QuickStart => "Quick Start", + Self::Interface => "Interface", + } + } + + pub fn description(&self) -> &str { + match self { + Self::QuickStart => "Select your experience", + Self::Interface => "Complete the task", + } + } +} diff --git a/crates/space_robotics_bench_gui/src/style/fonts.rs b/crates/space_robotics_bench_gui/src/style/fonts.rs new file mode 100644 index 0000000..6b310f8 --- /dev/null +++ b/crates/space_robotics_bench_gui/src/style/fonts.rs @@ -0,0 +1,20 @@ +use egui::Context; + +pub fn set(ctx: &Context) { + let font_definitions = crate::macros::generate_font_definitions! { + // Monospace + "MonaspaceNeon" as Monospace [ + "MonaspaceNeon-Medium.otf", + ], + + // Proportional + "Inter" as Proportional [ + "Inter-Regular.otf", + ], + "MaterialIcons" as Proportional [ + "MaterialIconsRound-Regular.otf", + ], + }; + + ctx.set_fonts(font_definitions); +} diff --git a/crates/space_robotics_bench_gui/src/style/mod.rs b/crates/space_robotics_bench_gui/src/style/mod.rs new file mode 100644 index 0000000..a13fd95 --- /dev/null +++ b/crates/space_robotics_bench_gui/src/style/mod.rs @@ -0,0 +1,23 @@ +pub use visuals::{dark as dark_visuals, light as light_visuals}; + +mod fonts; +mod text; +mod visuals; + +pub fn load_fonts(ctx: &egui::Context) { + // Load and set the fonts + fonts::set(ctx); + + // Set the text styles + ctx.style_mut(|style| { + style.text_styles = text::styles(); + }); +} + +pub fn set_theme(ctx: &egui::Context, theme: egui::Theme) { + // Set the style + ctx.set_visuals(match theme { + egui::Theme::Dark => dark_visuals().clone(), + egui::Theme::Light => light_visuals().clone(), + }); +} diff --git a/crates/space_robotics_bench_gui/src/style/text.rs b/crates/space_robotics_bench_gui/src/style/text.rs new file mode 100644 index 0000000..74b6c20 --- /dev/null +++ b/crates/space_robotics_bench_gui/src/style/text.rs @@ -0,0 +1,24 @@ +use egui::{FontFamily, FontId, TextStyle}; + +pub fn styles() -> std::collections::BTreeMap { + [ + ( + TextStyle::Monospace, + FontId::new(16.0, FontFamily::Monospace), + ), + ( + TextStyle::Small, + FontId::new(12.0, FontFamily::Proportional), + ), + (TextStyle::Body, FontId::new(18.0, FontFamily::Proportional)), + ( + TextStyle::Button, + FontId::new(20.0, FontFamily::Proportional), + ), + ( + TextStyle::Heading, + FontId::new(48.0, FontFamily::Proportional), + ), + ] + .into() +} diff --git a/crates/space_robotics_bench_gui/src/style/visuals.rs b/crates/space_robotics_bench_gui/src/style/visuals.rs new file mode 100644 index 0000000..ca54467 --- /dev/null +++ b/crates/space_robotics_bench_gui/src/style/visuals.rs @@ -0,0 +1,207 @@ +//! Global visual settings for the UI. +//! +//! The color themes are inspired by [Catppuccin](https://catppuccin.com). +//! - Mocha palette for dark mode +//! - Latte palette for light mode + +use std::sync::OnceLock; + +use egui::{ + epaint::Shadow, + style::{HandleShape, NumericColorSpace, Selection, TextCursorStyle, WidgetVisuals, Widgets}, + Color32, Rounding, Stroke, Visuals, +}; + +pub fn dark() -> &'static Visuals { + static VISUALS: OnceLock = OnceLock::new(); + VISUALS.get_or_init(dark_theme) +} + +fn dark_theme() -> Visuals { + Visuals { + dark_mode: true, + override_text_color: None, + widgets: Widgets { + noninteractive: WidgetVisuals { + weak_bg_fill: Color32::from_rgb(30, 30, 46).gamma_multiply(0.75), + bg_fill: Color32::from_rgb(30, 30, 46), + bg_stroke: Stroke::new(1.0, Color32::from_rgb(127, 132, 156)), + fg_stroke: Stroke::new(1.0, Color32::from_rgb(138, 143, 163)), + rounding: Rounding::same(2.0), + expansion: 0.0, + }, + inactive: WidgetVisuals { + weak_bg_fill: Color32::from_rgb(49, 50, 68).gamma_multiply(0.75), + bg_fill: Color32::from_rgb(49, 50, 68), + bg_stroke: Stroke::new(1.0, Color32::from_rgb(127, 132, 156)), + fg_stroke: Stroke::new(1.0, Color32::from_rgb(138, 143, 163)), + rounding: Rounding::same(2.0), + expansion: 0.0, + }, + hovered: WidgetVisuals { + weak_bg_fill: Color32::from_rgb(88, 91, 112).gamma_multiply(0.75), + bg_fill: Color32::from_rgb(88, 91, 112), + bg_stroke: Stroke::new(1.0, Color32::from_rgb(127, 132, 156)), + fg_stroke: Stroke::new(1.5, Color32::from_rgb(205, 214, 244)), + rounding: Rounding::same(3.0), + expansion: 1.0, + }, + active: WidgetVisuals { + weak_bg_fill: Color32::from_rgb(69, 71, 90).gamma_multiply(0.75), + bg_fill: Color32::from_rgb(69, 71, 90), + bg_stroke: Stroke::new(1.0, Color32::from_rgb(127, 132, 156)), + fg_stroke: Stroke::new(2.0, Color32::from_rgb(205, 214, 244)), + rounding: Rounding::same(2.0), + expansion: 1.0, + }, + open: WidgetVisuals { + weak_bg_fill: Color32::from_rgb(49, 50, 68).gamma_multiply(0.75), + bg_fill: Color32::from_rgb(49, 50, 68), + bg_stroke: Stroke::new(1.0, Color32::from_rgb(127, 132, 156)), + fg_stroke: Stroke::new(1.0, Color32::from_rgb(205, 214, 244)), + rounding: Rounding::same(2.0), + expansion: 0.0, + }, + }, + selection: Selection { + bg_fill: Color32::from_rgb(137, 180, 250).gamma_multiply(0.25), + stroke: Stroke::new(1.0, Color32::from_rgb(127, 132, 156)), + }, + hyperlink_color: Color32::from_rgb(245, 224, 220), + faint_bg_color: Color32::from_rgb(49, 50, 68), + extreme_bg_color: Color32::from_rgb(17, 17, 27), + code_bg_color: Color32::from_rgb(24, 24, 37), + warn_fg_color: Color32::from_rgb(250, 179, 135), + error_fg_color: Color32::from_rgb(235, 160, 172), + window_rounding: Rounding::same(6.0), + window_shadow: Shadow { + offset: egui::Vec2 { x: 10.0, y: 20.0 }, + blur: 15.0, + spread: 0.0, + color: Color32::from_rgb(30, 30, 46), + }, + window_fill: Color32::from_rgb(30, 30, 46), + window_stroke: Stroke::new(1.0, Color32::from_rgb(127, 132, 156)), + window_highlight_topmost: true, + menu_rounding: Rounding::same(6.0), + panel_fill: Color32::from_rgb(30, 30, 46), + popup_shadow: Shadow { + offset: egui::Vec2 { x: 6.0, y: 10.0 }, + blur: 8.0, + spread: 0.0, + color: Color32::from_rgb(30, 30, 46), + }, + resize_corner_size: 12.0, + text_cursor: TextCursorStyle { + stroke: Stroke::new(2.0, Color32::from_rgb(210, 219, 250)), + ..Default::default() + }, + clip_rect_margin: 3.0, + button_frame: false, + collapsing_header_frame: false, + indent_has_left_vline: true, + striped: false, + slider_trailing_fill: false, + handle_shape: HandleShape::Circle, + interact_cursor: Some(egui::CursorIcon::PointingHand), + image_loading_spinners: true, + numeric_color_space: NumericColorSpace::GammaByte, + } +} + +pub fn light() -> &'static Visuals { + static VISUALS: OnceLock = OnceLock::new(); + VISUALS.get_or_init(light_theme) +} + +fn light_theme() -> Visuals { + Visuals { + dark_mode: false, + override_text_color: None, + widgets: Widgets { + noninteractive: WidgetVisuals { + weak_bg_fill: Color32::from_rgb(239, 241, 245).gamma_multiply(0.75), + bg_fill: Color32::from_rgb(239, 241, 245), + bg_stroke: Stroke::new(1.0, Color32::from_rgb(140, 143, 161)), + fg_stroke: Stroke::new(1.0, Color32::from_rgb(76, 79, 105)), + rounding: Rounding::same(2.0), + expansion: 0.0, + }, + inactive: WidgetVisuals { + weak_bg_fill: Color32::from_rgb(204, 208, 218).gamma_multiply(0.75), + bg_fill: Color32::from_rgb(204, 208, 218), + bg_stroke: Stroke::new(1.0, Color32::from_rgb(140, 143, 161)), + fg_stroke: Stroke::new(1.0, Color32::from_rgb(76, 79, 105)), + rounding: Rounding::same(2.0), + expansion: 0.0, + }, + hovered: WidgetVisuals { + weak_bg_fill: Color32::from_rgb(172, 176, 190).gamma_multiply(0.75), + bg_fill: Color32::from_rgb(172, 176, 190), + bg_stroke: Stroke::new(1.0, Color32::from_rgb(140, 143, 161)), + fg_stroke: Stroke::new(1.5, Color32::from_rgb(11, 12, 16)), + rounding: Rounding::same(3.0), + expansion: 1.0, + }, + active: WidgetVisuals { + weak_bg_fill: Color32::from_rgb(188, 192, 204).gamma_multiply(0.75), + bg_fill: Color32::from_rgb(188, 192, 204), + bg_stroke: Stroke::new(1.0, Color32::from_rgb(140, 143, 161)), + fg_stroke: Stroke::new(2.0, Color32::from_rgb(11, 12, 16)), + rounding: Rounding::same(2.0), + expansion: 1.0, + }, + open: WidgetVisuals { + weak_bg_fill: Color32::from_rgb(204, 208, 218).gamma_multiply(0.75), + bg_fill: Color32::from_rgb(204, 208, 218), + bg_stroke: Stroke::new(1.0, Color32::from_rgb(140, 143, 161)), + fg_stroke: Stroke::new(1.0, Color32::from_rgb(11, 12, 16)), + rounding: Rounding::same(2.0), + expansion: 0.0, + }, + }, + selection: Selection { + bg_fill: Color32::from_rgb(30, 102, 245).gamma_multiply(0.25), + stroke: Stroke::new(1.0, Color32::from_rgb(140, 143, 128)), + }, + hyperlink_color: Color32::from_rgb(220, 138, 120), + faint_bg_color: Color32::from_rgb(204, 208, 218), + extreme_bg_color: Color32::from_rgb(220, 224, 232), + code_bg_color: Color32::from_rgb(230, 233, 239), + warn_fg_color: Color32::from_rgb(254, 100, 11), + error_fg_color: Color32::from_rgb(230, 69, 83), + window_rounding: Rounding::same(6.0), + window_highlight_topmost: true, + menu_rounding: Rounding::same(6.0), + window_shadow: Shadow { + offset: egui::Vec2 { x: 10.0, y: 20.0 }, + blur: 15.0, + spread: 0.0, + color: Color32::from_rgb(239, 241, 245), + }, + window_fill: Color32::from_rgb(239, 241, 245), + window_stroke: Stroke::new(1.0, Color32::from_rgb(140, 143, 161)), + panel_fill: Color32::from_rgb(239, 241, 245), + popup_shadow: Shadow { + offset: egui::Vec2 { x: 6.0, y: 10.0 }, + blur: 8.0, + spread: 0.0, + color: Color32::from_rgb(239, 241, 245), + }, + resize_corner_size: 12.0, + text_cursor: TextCursorStyle { + stroke: Stroke::new(2.0, Color32::from_rgb(63, 65, 87)), + ..Default::default() + }, + clip_rect_margin: 3.0, + button_frame: false, + collapsing_header_frame: false, + indent_has_left_vline: true, + striped: false, + slider_trailing_fill: false, + handle_shape: HandleShape::Circle, + interact_cursor: Some(egui::CursorIcon::PointingHand), + image_loading_spinners: true, + numeric_color_space: NumericColorSpace::GammaByte, + } +} diff --git a/crates/space_robotics_bench_gui/src/utils/egui.rs b/crates/space_robotics_bench_gui/src/utils/egui.rs new file mode 100644 index 0000000..564d1dd --- /dev/null +++ b/crates/space_robotics_bench_gui/src/utils/egui.rs @@ -0,0 +1,104 @@ +#[must_use = "You should call .show()"] +#[derive(Debug, Clone, Copy, PartialEq, typed_builder::TypedBuilder)] +pub struct ScrollableFramedCentralPanel { + #[builder(default = 1024.0)] + pub max_content_width: f32, + #[builder(default = egui::Margin {left: 0.0, right: 8.0, top: 0.0, bottom: 0.0})] + pub min_inner_margin: egui::Margin, +} + +impl Default for ScrollableFramedCentralPanel { + fn default() -> Self { + Self::builder().build() + } +} + +impl ScrollableFramedCentralPanel { + pub fn show( + self, + ctx: &egui::Context, + add_contents: impl FnOnce(&mut egui::Ui) -> R, + ) -> egui::InnerResponse { + debug_assert!(self.max_content_width.is_sign_positive()); + debug_assert!(self.min_inner_margin.left.is_sign_positive()); + debug_assert!(self.min_inner_margin.right.is_sign_positive()); + debug_assert!(self.min_inner_margin.top.is_sign_positive()); + debug_assert!(self.min_inner_margin.bottom.is_sign_positive()); + + egui::CentralPanel::default().show(ctx, |ui| { + // egui::ScrollArea::vertical() + // .show(ui, |ui| { + let margin_x = (ctx.screen_rect().width() - self.max_content_width).max(0.0) / 2.0; + let inner_margin = egui::Margin { + left: margin_x.max(self.min_inner_margin.left), + right: margin_x.max(self.min_inner_margin.right), + ..self.min_inner_margin + }; + egui::Frame::default() + .inner_margin(inner_margin) + .show(ui, |ui| add_contents(ui)) + .inner + // }) + // .inner + }) + } +} + +#[cfg(target_arch = "wasm32")] +pub fn open_url_on_page(ctx: &egui::Context, page: crate::page::Page, same_tab: bool) { + let target_url = format!("#{}", page.to_string().to_lowercase()); + ctx.open_url(if same_tab { + egui::OpenUrl::same_tab(target_url) + } else { + egui::OpenUrl::new_tab(target_url) + }); +} + +pub fn clickable_url(response: egui::Response, url: impl ToString) -> egui::Response { + debug_assert!(response.sense.click); + + if response.clicked() { + response.ctx.open_url(egui::OpenUrl::same_tab(url)); + } else { + #[cfg(target_arch = "wasm32")] + if response.middle_clicked() { + response.ctx.open_url(egui::OpenUrl::new_tab(url)); + } + } + response +} + +pub fn strong_heading(ui: &mut egui::Ui, text: impl Into) -> egui::Response { + ui.label(egui::RichText::new(text).heading().strong()) +} + +pub fn centered_strong_heading( + ui: &mut egui::Ui, + text: impl Into, +) -> egui::InnerResponse { + ui.with_layout(egui::Layout::top_down(egui::Align::Center), |ui| { + strong_heading(ui, text) + }) +} + +pub fn heading_sized(ui: &mut egui::Ui, text: impl Into, size: f32) -> egui::Response { + ui.label(egui::RichText::new(text).heading().size(size)) +} + +pub fn strong_heading_sized( + ui: &mut egui::Ui, + text: impl Into, + size: f32, +) -> egui::Response { + ui.label(egui::RichText::new(text).heading().strong().size(size)) +} + +pub fn centered_strong_heading_sized( + ui: &mut egui::Ui, + text: impl Into, + size: f32, +) -> egui::InnerResponse { + ui.with_layout(egui::Layout::top_down(egui::Align::Center), |ui| { + crate::utils::egui::strong_heading_sized(ui, text, size) + }) +} diff --git a/crates/space_robotics_bench_gui/src/utils/enums.rs b/crates/space_robotics_bench_gui/src/utils/enums.rs new file mode 100644 index 0000000..21ee27d --- /dev/null +++ b/crates/space_robotics_bench_gui/src/utils/enums.rs @@ -0,0 +1,116 @@ +use eframe::epaint::Color32; + +#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)] +pub enum Difficulty { + Demo, + Easy, + Medium, + Challenging, +} + +impl std::fmt::Display for Difficulty { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Demo => write!(f, "Demo"), + Self::Easy => write!(f, "Easy"), + Self::Medium => write!(f, "Medium"), + Self::Challenging => write!(f, "Challenging"), + } + } +} + +impl Difficulty { + pub fn get_text_color(self, strong: bool) -> Color32 { + match self { + Self::Demo => { + if strong { + Color32::from_rgb(245, 245, 245) + } else { + Color32::from_rgb(170, 170, 170) + } + } + Self::Easy => { + if strong { + Color32::from_rgb(46, 220, 117) + } else { + Color32::from_rgb(77, 171, 35) + } + } + Self::Medium => { + if strong { + Color32::from_rgb(211, 218, 60) + } else { + Color32::from_rgb(211, 100, 0) + } + } + Self::Challenging => { + if strong { + Color32::from_rgb(219, 39, 40) + } else { + Color32::from_rgb(132, 0, 32) + } + } + } + } + + pub fn set_theme(self, ui: &mut egui::Ui, theme: egui::Theme) { + match self { + Self::Demo => match theme { + egui::Theme::Light => { + ui.style_mut().visuals.widgets.inactive.weak_bg_fill = + Color32::from_rgb(245, 245, 245); + ui.style_mut().visuals.widgets.hovered.weak_bg_fill = + Color32::from_rgb(170, 170, 170); + } + egui::Theme::Dark => { + ui.style_mut().visuals.widgets.inactive.weak_bg_fill = + Color32::from_rgb(170, 170, 170); + ui.style_mut().visuals.widgets.hovered.weak_bg_fill = + Color32::from_rgb(245, 245, 245); + } + }, + Self::Easy => match theme { + egui::Theme::Light => { + ui.style_mut().visuals.widgets.inactive.weak_bg_fill = + Color32::from_rgb(46, 220, 117); + ui.style_mut().visuals.widgets.hovered.weak_bg_fill = + Color32::from_rgb(77, 171, 35); + } + egui::Theme::Dark => { + ui.style_mut().visuals.widgets.inactive.weak_bg_fill = + Color32::from_rgb(77, 171, 35); + ui.style_mut().visuals.widgets.hovered.weak_bg_fill = + Color32::from_rgb(46, 220, 117); + } + }, + Self::Medium => match theme { + egui::Theme::Light => { + ui.style_mut().visuals.widgets.inactive.weak_bg_fill = + Color32::from_rgb(211, 218, 60); + ui.style_mut().visuals.widgets.hovered.weak_bg_fill = + Color32::from_rgb(211, 100, 0); + } + egui::Theme::Dark => { + ui.style_mut().visuals.widgets.inactive.weak_bg_fill = + Color32::from_rgb(211, 100, 0); + ui.style_mut().visuals.widgets.hovered.weak_bg_fill = + Color32::from_rgb(211, 218, 60); + } + }, + Self::Challenging => match theme { + egui::Theme::Light => { + ui.style_mut().visuals.widgets.inactive.weak_bg_fill = + Color32::from_rgb(219, 39, 40); + ui.style_mut().visuals.widgets.hovered.weak_bg_fill = + Color32::from_rgb(132, 0, 32); + } + egui::Theme::Dark => { + ui.style_mut().visuals.widgets.inactive.weak_bg_fill = + Color32::from_rgb(132, 0, 32); + ui.style_mut().visuals.widgets.hovered.weak_bg_fill = + Color32::from_rgb(219, 39, 40); + } + }, + } + } +} diff --git a/crates/space_robotics_bench_gui/src/utils/mod.rs b/crates/space_robotics_bench_gui/src/utils/mod.rs new file mode 100644 index 0000000..adf34b0 --- /dev/null +++ b/crates/space_robotics_bench_gui/src/utils/mod.rs @@ -0,0 +1,6 @@ +#![allow(unused)] + +pub use enums::Difficulty; + +pub mod egui; +mod enums; diff --git a/crates/space_robotics_bench_py/Cargo.toml b/crates/space_robotics_bench_py/Cargo.toml new file mode 100644 index 0000000..e9dcad6 --- /dev/null +++ b/crates/space_robotics_bench_py/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "space_robotics_bench_py" +description.workspace = true +categories.workspace = true +keywords.workspace = true +readme.workspace = true +license.workspace = true +homepage.workspace = true +repository.workspace = true +documentation.workspace = true +edition.workspace = true +rust-version.workspace = true +version.workspace = true +publish.workspace = true + +[lib] +crate-type = ["cdylib"] + +[dependencies] +space_robotics_bench = { workspace = true } + +const_format = { workspace = true } +pyo3 = { workspace = true } diff --git a/crates/space_robotics_bench_py/src/envs/mod.rs b/crates/space_robotics_bench_py/src/envs/mod.rs new file mode 100644 index 0000000..4585735 --- /dev/null +++ b/crates/space_robotics_bench_py/src/envs/mod.rs @@ -0,0 +1,14 @@ +use pyo3::prelude::*; +use space_robotics_bench::envs::{Asset, AssetVariant, Assets, EnvironmentConfig, Scenario}; + +pub(crate) fn register(parent: &Bound<'_, PyModule>) -> PyResult<()> { + let m = PyModule::new_bound(parent.py(), crate::macros::python_module_name!())?; + + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + + crate::macros::python_add_submodule!(parent, m) +} diff --git a/crates/space_robotics_bench_py/src/lib.rs b/crates/space_robotics_bench_py/src/lib.rs new file mode 100644 index 0000000..c95263e --- /dev/null +++ b/crates/space_robotics_bench_py/src/lib.rs @@ -0,0 +1,15 @@ +use pyo3::prelude::*; +mod macros; + +mod envs; +mod utils; + +/// Name of the root Python module generated by this Rust extension module. +const PYTHON_MODULE_NAME: &str = "space_robotics_bench._rs"; + +#[pymodule] +fn _rs(m: &Bound<'_, PyModule>) -> PyResult<()> { + envs::register(m)?; + utils::register(m)?; + Ok(()) +} diff --git a/crates/space_robotics_bench_py/src/macros.rs b/crates/space_robotics_bench_py/src/macros.rs new file mode 100644 index 0000000..9051419 --- /dev/null +++ b/crates/space_robotics_bench_py/src/macros.rs @@ -0,0 +1,29 @@ +macro_rules! python_module_name { + () => { + ::const_format::str_split!(::std::module_path!(), "::") + .last() + .unwrap() + }; +} + +macro_rules! python_module_name_full { + () => { + ::const_format::str_replace!( + ::const_format::str_replace!(::std::module_path!(), "::", "."), + ::std::env!("CARGO_PKG_NAME"), + $crate::PYTHON_MODULE_NAME + ) + }; +} + +macro_rules! python_add_submodule { + ($parent_module:ident, $module:ident) => {{ + let py = $parent_module.py(); + $parent_module.add_submodule(&$module)?; + py.import_bound(::pyo3::intern!(py, "sys"))? + .getattr(::pyo3::intern!(py, "modules"))? + .set_item($crate::macros::python_module_name_full!(), $module) + }}; +} + +pub(crate) use {python_add_submodule, python_module_name, python_module_name_full}; diff --git a/crates/space_robotics_bench_py/src/utils/mod.rs b/crates/space_robotics_bench_py/src/utils/mod.rs new file mode 100644 index 0000000..97b5ce1 --- /dev/null +++ b/crates/space_robotics_bench_py/src/utils/mod.rs @@ -0,0 +1,10 @@ +use pyo3::prelude::*; +mod sampling; + +pub(crate) fn register(parent: &Bound<'_, PyModule>) -> PyResult<()> { + let m = PyModule::new_bound(parent.py(), crate::macros::python_module_name!())?; + + sampling::register(&m)?; + + crate::macros::python_add_submodule!(parent, m) +} diff --git a/crates/space_robotics_bench_py/src/utils/sampling/mod.rs b/crates/space_robotics_bench_py/src/utils/sampling/mod.rs new file mode 100644 index 0000000..79b7067 --- /dev/null +++ b/crates/space_robotics_bench_py/src/utils/sampling/mod.rs @@ -0,0 +1,16 @@ +use pyo3::prelude::*; +use space_robotics_bench::utils::sampling::{ + sample_poisson_disk_2d, sample_poisson_disk_2d_looped, sample_poisson_disk_3d, + sample_poisson_disk_3d_looped, +}; + +pub(crate) fn register(parent: &Bound<'_, PyModule>) -> PyResult<()> { + let m = PyModule::new_bound(parent.py(), crate::macros::python_module_name!())?; + + m.add_function(wrap_pyfunction!(sample_poisson_disk_2d, &m)?)?; + m.add_function(wrap_pyfunction!(sample_poisson_disk_2d_looped, &m)?)?; + m.add_function(wrap_pyfunction!(sample_poisson_disk_3d, &m)?)?; + m.add_function(wrap_pyfunction!(sample_poisson_disk_3d_looped, &m)?)?; + + crate::macros::python_add_submodule!(parent, m) +} diff --git a/crates/space_robotics_bench_py/tests/import_test.py b/crates/space_robotics_bench_py/tests/import_test.py new file mode 100644 index 0000000..3a7286d --- /dev/null +++ b/crates/space_robotics_bench_py/tests/import_test.py @@ -0,0 +1,2 @@ +def test_import(): + import space_robotics_bench # noqa: F401 diff --git a/crates/space_robotics_bench_py/tests/pytest.rs b/crates/space_robotics_bench_py/tests/pytest.rs new file mode 100644 index 0000000..6734592 --- /dev/null +++ b/crates/space_robotics_bench_py/tests/pytest.rs @@ -0,0 +1,19 @@ +//! Run all python tests with pytest. + +#[test] +fn pytest() { + // Arrange + let python_exe = std::env::var("PYO3_PYTHON") + .or_else(|_| std::env::var("ISAAC_SIM_PYTHON")) + .unwrap_or("python3".to_string()); + let mut command = std::process::Command::new(python_exe); + let command = command.arg("-m").arg("pytest"); + + // Act + let output = command.output().unwrap(); + println!("{}", std::str::from_utf8(&output.stdout).unwrap()); + eprintln!("{}", std::str::from_utf8(&output.stderr).unwrap()); + + // Assert + assert!(output.status.success()); +} diff --git a/deny.toml b/deny.toml new file mode 100644 index 0000000..96e7b1c --- /dev/null +++ b/deny.toml @@ -0,0 +1,62 @@ +# `cargo deny` is only intended to run these targets for this project +targets = [ + { triple = "aarch64-unknown-linux-gnu" }, + { triple = "x86_64-unknown-linux-gnu" }, + { triple = "x86_64-unknown-linux-musl" }, +] + +# Considered when running `cargo deny check advisories` +# https://embarkstudios.github.io/cargo-deny/checks/advisories/cfg.html +[advisories] +notice = "deny" +unmaintained = "warn" +unsound = "deny" +vulnerability = "deny" +yanked = "deny" +ignore = [] + +# Considered when running `cargo deny check licenses` +# https://embarkstudios.github.io/cargo-deny/checks/licenses/cfg.html +[licenses] +allow-osi-fsf-free = "neither" +copyleft = "deny" +unlicensed = "deny" +private = { ignore = true } +confidence-threshold = 0.925 +allow = [ + "Apache-2.0 WITH LLVM-exception", # https://spdx.org/licenses/LLVM-exception.html + "Apache-2.0", # https://spdx.org/licenses/Apache-2.0.html + "BSD-2-Clause", # https://spdx.org/licenses/BSD-2-Clause.html + "BSD-3-Clause", # https://spdx.org/licenses/BSD-3-Clause.html + "ISC", # https://spdx.org/licenses/ISC.html + "LicenseRef-UFL-1.0", # https://spdx.org/licenses/LicenseRef-UFL-1.0.html + "MIT", # https://spdx.org/licenses/MIT.html + "MPL-2.0", # https://spdx.org/licenses/MPL-2.0.html + "OFL-1.1", # https://spdx.org/licenses/OFL-1.1.html + "Unicode-DFS-2016", # https://spdx.org/licenses/Unicode-DFS-2016.html + "Zlib", # https://spdx.org/licenses/Zlib.html +] +exceptions = [] + +[[licenses.clarify]] +crate = "ring" +expression = "ISC" +license-files = [{ path = "LICENSE", hash = 0xbd0eed23 }] + +# Considered when running `cargo deny check bans` +# https://embarkstudios.github.io/cargo-deny/checks/bans/cfg.html +[bans] +multiple-versions = "warn" +wildcards = "allow" +deny = [] +skip = [] +skip-tree = [] + +# Considered when running `cargo deny check sources` +# https://embarkstudios.github.io/cargo-deny/checks/sources/cfg.html +[sources] +unknown-registry = "deny" +unknown-git = "deny" + +[sources.allow-org] +github = ["AndrejOrsula", "lampsitter"] diff --git a/docs/.gitignore b/docs/.gitignore new file mode 100644 index 0000000..baa790a --- /dev/null +++ b/docs/.gitignore @@ -0,0 +1,2 @@ +book/ +index.html diff --git a/docs/README.md b/docs/README.md new file mode 100644 index 0000000..b5da422 --- /dev/null +++ b/docs/README.md @@ -0,0 +1,11 @@ +# Documentation + + HTML + +## Local Preview + +We use [mdBook](https://rust-lang.github.io/mdBook) to generate a static site from Markdown files found in the [src](src) directory. To build and preview the site locally, you can [install mdBook](https://rust-lang.github.io/mdBook/guide/installation.html) and run the following command: + +```bash +mdbook serve --open +``` diff --git a/docs/book.toml b/docs/book.toml new file mode 100644 index 0000000..421c915 --- /dev/null +++ b/docs/book.toml @@ -0,0 +1,9 @@ +[book] +title = "Space Robotics Bench" +description = "Comprehensive benchmark for space robotics" +authors = ["Andrej Orsula"] +language = "en" + +[output.html] +site-url = "/space_robotics_bench/" +git-repository-url = "https://github.com/AndrejOrsula/space_robotics_bench" diff --git a/docs/src/README.md b/docs/src/README.md new file mode 100644 index 0000000..9cd0b68 --- /dev/null +++ b/docs/src/README.md @@ -0,0 +1,31 @@ +# Space Robotics Bench + +![](./_images/srb_multi_env.jpg) + +The **Space Robotics Bench** aims to be a comprehensive collection of environments and tasks for robotics research in the challenging domain of space. The benchmark covers a wide range of applications and scenarios while providing a unified framework for experimenting with new tasks. Although the primary focus is on the application of robot learning techniques, the benchmark is designed to be flexible and extensible to accommodate a variety of research directions. + +
+This documentation is currently incomplete. Inactive pages found in the navigation panel indicate what topics will be covered prior to the first release. Please let us know by opening an issue if something is missing or about a specific topic that you are interested in having documented first. Thank you! :) +
+ +## Key Features + +### On-Demand Procedural Generation with [Blender](https://blender.org) + +Blender is used to generate procedural assets across a wide range of scenarios to provide environments that are representative of the diversity in space. By doing so, this benchmark emphasizes the need for generalization and adapatibility of robots in space due to their safety-critical nature. + +### Highly-Parallelized Simulation with [NVIDIA Isaac Sim](https://developer.nvidia.com/isaac-sim) + +By leveraging the hardware-acceleration capabilities of NVIDIA Isaac Sim, all environments support parallel simulation instances, significantly accelerating workflows such as parameter tuning, verification, synthetic data generation, and online learning. The uniqueness of each procedurally generated instance also contributes towards the diversity that robots experience alongside the included domain randomization. Furthermore, compliance with [Isaac Lab](https://isaac-sim.github.io/IsaacLab) enhances compatibility with a wide array of pre-configured robots and sensors. + +### Compatibility with [Gymnasium API](https://gymnasium.farama.org) + +All tasks are registered with the standardized Gymnasium API, ensuring seamless integration with a broad ecosystem of libraries and tools. This enables developers to leverage popular reinforcement learning and imitation learning algorithms while also simplifying the evaluation and comparison of various solutions across diverse scenarios, giving rise to potential collaboration efforts. + +### Integration with [ROS 2](https://ros.org) & [Space ROS](https://space.ros.org) + +The benchmark can also be installed as a ROS 2 package to bring interoperability to its wide ecosystem, including aspects of Space ROS. This integration provides access to a rich set of tools and libraries that accelerate the development and deployment of robotic systems. At the same time, ROS developers get access to a set of reproducible space environments for evaluating their systems and algorithms while benefiting from the procedural variety and parallel instances via namespaced middleware communication. + +### Agnostic Interfaces + +The interfaces of the benchmark are designed with abstraction layers to ensure flexibility for various applications and systems. By adjusting configuration and changing procedural pipelines, a single task definition can be reused across different robots and domains of space. Moreover, all assets are decoupled from the benchmark into a separate [`srb_assets` repository](https://github.com/AndrejOrsula/srb_assets), enabling their straightforward integration with external frameworks and projects. diff --git a/docs/src/SUMMARY.md b/docs/src/SUMMARY.md new file mode 100644 index 0000000..e626666 --- /dev/null +++ b/docs/src/SUMMARY.md @@ -0,0 +1,62 @@ +[Introduction](README.md) + +______________________________________________________________________ + +# Overview + +- [Environments](overview/envs/README.md) + - [Mobile](overview/envs/mobile.md) + - [Wheeled](overview/envs/mobile_wheeled.md) + - [Aerial](overview/envs/mobile_aerial.md) + - [Spacecraft](<>) + - [Manipulation](overview/envs/manipulation.md) + - [Mobile Manipulation](<>) + - [Planetary](<>) + - [Orbital](<>) +- [Integrations](overview/integrations/README.md) + - [Reinforcement Learning](<>) + - [Imitation Learning & Offline RL](<>) + - [ROS 2](overview/integrations/ros2.md) + +# Getting Started + +- [System Requirements](getting_started/requirements.md) +- [Installation](getting_started/installation/README.md) + - [Local](<>) + - [Docker](getting_started/installation/docker.md) + - [Apptainer/Singularity](<>) +- [Basic Usage](getting_started/usage.md) + +# Instructions + +- [Benchmark](instructions/benchmark/README.md) + - [Parallel Environments](instructions/benchmark/parallel_envs.md) + - [Random and Zero Agents](instructions/benchmark/random_zero_agents.md) + - [Visual Observations](instructions/benchmark/visual_observations.md) + - [Graphical User Interface](instructions/benchmark/gui.md) + - [Configuration](instructions/benchmark/cfg.md) +- [Workflows](instructions/workflows/README.md) + - [ROS 2](instructions/workflows/ros2.md) + - [Stable Baselines3](<>) + - [DreamerV3](<>) + - [Robomimic](<>) + - [Sim-to-Real](<>) +- [Utilities](instructions/utils/README.md) + - [Clean the Assets Cache](instructions/utils/clean_cache.md) + +# Development + +- [Development Environment](development/dev_env/README.md) + - [IDE Configuration](<>) + - [Docker](development/dev_env/docker.md) + - [Dev Container](development/dev_env/dev_container.md) +- [New Assets](development/new_assets/README.md) + - [Datasets](<>) + - [Procedural Assets with Blender](development/new_assets/procgen.md) +- [New Environments](development/new_envs.md) + +______________________________________________________________________ + +[Troubleshooting](misc/troubleshooting.md) +[Attributions](misc/attributions.md) +[Contributors](misc/contributors.md) diff --git a/docs/src/_images/envs/debris_capture_orbit.jpg b/docs/src/_images/envs/debris_capture_orbit.jpg new file mode 100644 index 0000000..987a9b9 Binary files /dev/null and b/docs/src/_images/envs/debris_capture_orbit.jpg differ diff --git a/docs/src/_images/envs/gateway.jpg b/docs/src/_images/envs/gateway.jpg new file mode 100644 index 0000000..24ea618 Binary files /dev/null and b/docs/src/_images/envs/gateway.jpg differ diff --git a/docs/src/_images/envs/ingenuity.jpg b/docs/src/_images/envs/ingenuity.jpg new file mode 100644 index 0000000..2332b38 Binary files /dev/null and b/docs/src/_images/envs/ingenuity.jpg differ diff --git a/docs/src/_images/envs/peg_in_hole_moon.jpg b/docs/src/_images/envs/peg_in_hole_moon.jpg new file mode 100644 index 0000000..c6e6fe4 Binary files /dev/null and b/docs/src/_images/envs/peg_in_hole_moon.jpg differ diff --git a/docs/src/_images/envs/peg_in_hole_orbit.jpg b/docs/src/_images/envs/peg_in_hole_orbit.jpg new file mode 100644 index 0000000..968e91f Binary files /dev/null and b/docs/src/_images/envs/peg_in_hole_orbit.jpg differ diff --git a/docs/src/_images/envs/perseverance.jpg b/docs/src/_images/envs/perseverance.jpg new file mode 100644 index 0000000..d88ea97 Binary files /dev/null and b/docs/src/_images/envs/perseverance.jpg differ diff --git a/docs/src/_images/envs/sample_collection_mars.jpg b/docs/src/_images/envs/sample_collection_mars.jpg new file mode 100644 index 0000000..f2b69f3 Binary files /dev/null and b/docs/src/_images/envs/sample_collection_mars.jpg differ diff --git a/docs/src/_images/envs/sample_collection_moon.jpg b/docs/src/_images/envs/sample_collection_moon.jpg new file mode 100644 index 0000000..53aaf53 Binary files /dev/null and b/docs/src/_images/envs/sample_collection_moon.jpg differ diff --git a/docs/src/_images/envs/solar_panel_assembly_moon.jpg b/docs/src/_images/envs/solar_panel_assembly_moon.jpg new file mode 100644 index 0000000..f5c6e09 Binary files /dev/null and b/docs/src/_images/envs/solar_panel_assembly_moon.jpg differ diff --git a/docs/src/_images/perseverance_ui.jpg b/docs/src/_images/perseverance_ui.jpg new file mode 100644 index 0000000..79d088c Binary files /dev/null and b/docs/src/_images/perseverance_ui.jpg differ diff --git a/docs/src/_images/srb_gui.jpg b/docs/src/_images/srb_gui.jpg new file mode 100644 index 0000000..a540b2e Binary files /dev/null and b/docs/src/_images/srb_gui.jpg differ diff --git a/docs/src/_images/srb_multi_env.jpg b/docs/src/_images/srb_multi_env.jpg new file mode 100644 index 0000000..7b76480 Binary files /dev/null and b/docs/src/_images/srb_multi_env.jpg differ diff --git a/docs/src/development/dev_env/README.md b/docs/src/development/dev_env/README.md new file mode 100644 index 0000000..1364c81 --- /dev/null +++ b/docs/src/development/dev_env/README.md @@ -0,0 +1,6 @@ +# Development Environment + +Whether you are contributing to the benchmark or using the setup in your own projects, this section aims to improve your development experience. + +- [Development inside Docker](./docker.md) +- [Dev Container](./dev_container.md) diff --git a/docs/src/development/dev_env/dev_container.md b/docs/src/development/dev_env/dev_container.md new file mode 100644 index 0000000..9e778d0 --- /dev/null +++ b/docs/src/development/dev_env/dev_container.md @@ -0,0 +1,15 @@ +# Dev Container + +[Dev Containers](https://containers.dev) allow for a fully isolated development environment tailored to specific project needs. This is particularly useful for ensuring all dependencies are installed and consistent across different development machines. + +## Open the Benchmark in a Dev Container + +To simplify the process of building and opening the repository as a Dev Container in Visual Studio Code (VS Code), you can run the [`.devcontainer/open.bash`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/.devcontainer/open.bash) that automates the process. + +```bash +.devcontainer/open.bash +``` + +## Modify the Dev Container + +You can also customize the included [`.devcontainer/devcontainer.json`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/.devcontainer/devcontainer.json) configuration to suit your specific development requirements. diff --git a/docs/src/development/dev_env/docker.md b/docs/src/development/dev_env/docker.md new file mode 100644 index 0000000..674b3c5 --- /dev/null +++ b/docs/src/development/dev_env/docker.md @@ -0,0 +1,11 @@ +# Development inside Docker + +The Space Robotics Bench supports a Docker setup, which in itself provides an isolated development environment. By default, the [`.docker/run.bash`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/.docker/run.bash) script already mounts the source code into the container (can be disabled with `WITH_DEV_VOLUME=false`). In itself, this already makes the standalone Docker setup quite convenient for development. + +## Joint a Running Container + +Once the Docker container is running, you can join the running Docker container with the [`.docker/join.bash`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/.docker/join.bash) script: + +```bash +.docker/join.bash +``` diff --git a/docs/src/development/new_assets/README.md b/docs/src/development/new_assets/README.md new file mode 100644 index 0000000..201d631 --- /dev/null +++ b/docs/src/development/new_assets/README.md @@ -0,0 +1,5 @@ +# New Assets + +All assets used by the Space Robotics Bench are separated into the [`srb_assets`](https://github.com/AndrejOrsula/srb_assets) repository to encourage their reuse. Both static assets (datasets) and procedural pipelines in Blender are supported. + +- [Procedural Assets with Blender](./procgen.md) diff --git a/docs/src/development/new_assets/procgen.md b/docs/src/development/new_assets/procgen.md new file mode 100644 index 0000000..9690583 --- /dev/null +++ b/docs/src/development/new_assets/procgen.md @@ -0,0 +1,19 @@ +# Procedural Assets with Blender + +## Motivation + +Procedural generation is a powerful technique for creating diverse and realistic environments without relying on static, disk-consuming datasets. This approach allows for the generation of an infinite number of unique environments, a feature that has been underutilized in the fields of robotics and space exploration. The Space Robotics Bench seeks to address this gap by offering a versatile framework for procedurally generating 3D assets, which can be combined to create complex environments suitable for the development, training, and validation of space robotic systems. + +## Approach + +The package utilizes [Blender](https://www.blender.org) to procedurally generate both the geometry and materials (PBR textures) of 3D assets. + +Geometry generation is achieved using Blender's [Geometry Nodes](https://docs.blender.org/manual/en/latest/modeling/geometry_nodes/introduction.html), a robust node-based system that allows for the creation, manipulation, and modification of arbitrary geometry and data types. First introduced in Blender 2.92 (2021), Geometry Nodes have evolved significantly, supporting the creation of intricate geometries through a series of interconnected node trees. Each node system can consist of multiple node trees that handle different aspects of the geometry. By applying randomness and variation within these node trees, a wide range of unique assets can be produced simply by adjusting the seed value. + +Blender's [Shader Nodes](https://docs.blender.org/manual/en/latest/render/shader_nodes/introduction.html), which have a longer history, are used to define the appearance of objects through material properties. Like Geometry Nodes, Shader Nodes are also node-based and allow for the creation of complex materials. Blender provides several procedural textures and maps (e.g., Perlin noise, Voronoi, Wave), which can be adjusted and combined to form more sophisticated materials. By integrating randomness into the shader nodes, each procedurally generated asset can have a unique appearance, even with the same underlying geometry. + +## Workflow + +The package includes a `blender/procgen_assets.py` script that automates the entire procedural generation process, including node construction, modifier application, seeding, texture baking, and model export. This script is fully standalone and interacts with Blender's Python API (`bpy`) through its binary executable. Although Blender can be used as a Python module via [bpy](https://pypi.org/project/bpy), it is often linked to a specific Python version and has longer release cycles. The standalone script offers more flexibility, allowing it to be used with any Blender version. + +Node trees can be generated from Python source files provided as input to the script. The [Node To Python addon](https://extensions.blender.org/add-ons/node-to-python) simplifies the creation of such source code. This addon enables users to design node trees in Blender's graphical interface and convert them into Python code that can be integrated into the `procgen_assets.py` script. This method allows users to prototype assets interactively within Blender's GUI and then export them into code. diff --git a/docs/src/development/new_envs.md b/docs/src/development/new_envs.md new file mode 100644 index 0000000..e26358f --- /dev/null +++ b/docs/src/development/new_envs.md @@ -0,0 +1,19 @@ +# New Environments + +The process of introducing a new environment into the Space Robotics Bench is intended to be modular. + +## 1. Duplicate an Existing Environment + +Navigate to the [`tasks`](https://github.com/AndrejOrsula/space_robotics_bench/tree/main/space_robotics_bench/tasks) directory, which houses the existing environments. Then, duplicate one of the existing demos or task directories that resembles your desired task/demo more and rename it to the name of your new environment. + +## 2. Modify the Environment Configuration + +Customize your new environment by altering the configuration files and task implementation code within the folder. This may include asset selection, interaction rules, or specific environmental dynamics. + +## 3. Automatic Registration + +The new environment will be automatically registered with the Gymnasium API. The environment will be registered under the directory name you assigned during the duplication process. + +## 4. Running Your New Environment + +Test your new environment by specifying the name of your new environment via the `--env`/`--task`/`--demo` argument. diff --git a/docs/src/getting_started/installation/README.md b/docs/src/getting_started/installation/README.md new file mode 100644 index 0000000..4cd8592 --- /dev/null +++ b/docs/src/getting_started/installation/README.md @@ -0,0 +1,7 @@ +# Installation + +Currently, only Docker-based setup is fully supported and tested. + +Local installation should be as straightforward as following the Dockerfile instructions. However, it is yet to be explored! + +- [Installation (Docker)](./docker.md) diff --git a/docs/src/getting_started/installation/docker.md b/docs/src/getting_started/installation/docker.md new file mode 100644 index 0000000..c172910 --- /dev/null +++ b/docs/src/getting_started/installation/docker.md @@ -0,0 +1,84 @@ +# Installation (Docker) + +This section provides instructions for running the simulation within a Docker container. Before proceeding, ensure that your system meets the [system requirements](../requirements.md). If you are using a different operating system, you may need to adjust the following steps accordingly or refer to the official documentation for each step. + +## 1. Install [Docker Engine](https://docs.docker.com/engine) + +First, install Docker Engine by following the [official installation instructions](https://docs.docker.com/engine/install). For example: + +```bash +curl -fsSL https://get.docker.com | sh +sudo systemctl enable --now docker + +sudo groupadd docker +sudo usermod -aG docker $USER +newgrp docker +``` + +## 2. Install [NVIDIA Container Toolkit](https://github.com/NVIDIA/nvidia-container-toolkit) + +Next, install the NVIDIA Container Toolkit, which is required to enable GPU support for Docker containers. Follow the [official installation guide](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) or use the following commands: + +```bash +curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg && curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list +sudo apt-get update +sudo apt-get install -y nvidia-container-toolkit +sudo nvidia-ctk runtime configure --runtime=docker +sudo systemctl restart docker +``` + +## 3. Gain Access to the [Isaac Sim Docker Image](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/isaac-sim) + +To run the simulation, you need access to the Isaac Sim Docker image, which requires registering an account and generating an API key from NVIDIA GPU Cloud (NGC). + +### 3.1 Register and Log In to [NVIDIA GPU Cloud (NGC)](https://www.nvidia.com/en-us/gpu-cloud) + +Visit the [NGC portal](https://ngc.nvidia.com/signin) and register or log in to your account. + +### 3.2 Generate Your [NGC API Key](https://docs.nvidia.com/ngc/gpu-cloud/ngc-user-guide/README.md#ngc-api-keys) + +Follow the [official guide](https://docs.nvidia.com/ngc/gpu-cloud/ngc-user-guide/README.md#generating-personal-api-key) to generate your personal NGC API key. + +### 3.3 Log In to NGC via Docker + +Once you have your API key, log in to NGC through Docker: + +```bash +docker login nvcr.io +``` + +When prompted for a username, enter `$oauthtoken` (exactly as shown): + +```bash +Username: $oauthtoken +``` + +When prompted for a password, use the API key you just generated: + +```bash +Password: +``` + +## 4. Clone the Repository + +Next, clone the `space_robotics_bench` repository locally. Make sure to include the `--recurse-submodules` flag to clone also the submodule containing simulation assets. + +```bash +git clone --recurse-submodules https://github.com/AndrejOrsula/space_robotics_bench.git +``` + +## 5. Build the Docker Image + +Now, you can build the Docker image for `space_robotics_bench` by running the provided [`.docker/build.bash`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/.docker/build.bash) script. Note that the first build process may take up to 30 minutes (depending on your network speed and system configuration). + +```bash +space_robotics_bench/.docker/build.bash +``` + +## 6. Verify the Image Build + +To ensure that the image was built successfully, run the following command. You should see the `space_robotics_bench` image listed among recently created Docker images. + +```bash +docker images +``` diff --git a/docs/src/getting_started/requirements.md b/docs/src/getting_started/requirements.md new file mode 100644 index 0000000..dd0d127 --- /dev/null +++ b/docs/src/getting_started/requirements.md @@ -0,0 +1,36 @@ +# System Requirements + +
+This project requires a dedicated NVIDIA GPU with RT Cores (RTX series). Isaac Sim does not support GPUs from other vendors. +
+ +## Hardware Requirements + +The hardware requirements for running this simulation are inherited from the [Isaac Sim requirements](https://docs.omniverse.nvidia.com/isaacsim/latest/installation/requirements.html). While it is possible to run the simulation on lower-spec systems than those recommended, performance will be significantly reduced. + +| Component | Minimum Requirement | +| ------------ | ------------------------------------- | +| Architecture | `x86_64` | +| CPU | Any smart silicon-rich rock | +| RAM | 16 GB | +| GPU | NVIDIA GPU with RT Cores (RTX series) | +| VRAM | 4 GB | +| Disk Space | 30 GB | +| Network | 12 GB (for pulling Docker images) | + +## Software Requirements + +The following software requirements are essential for running the simulation. Other operating systems may work, but significant adjustments may be required. + +| Component | Requirement | +| ------------- | --------------------------------------------------- | +| OS | Linux-based distribution (e.g., Ubuntu 22.04/24.04) | +| NVIDIA Driver | 535.183.01 (tested; other versions may work) | + +### Additional Docker Requirements + +The current setup requires the X11 window manager to enable running GUI from within the Docker container. Other window managers may work, but significant adjustments may be required. + +| Component | Requirement | +| -------------- | ----------- | +| Window Manager | X11 | diff --git a/docs/src/getting_started/usage.md b/docs/src/getting_started/usage.md new file mode 100644 index 0000000..3a85a96 --- /dev/null +++ b/docs/src/getting_started/usage.md @@ -0,0 +1,91 @@ +# Basic Usage + +After successful [installation](./installation/index.html), you are ready to use the Space Robotics Bench. This page will guide you through controlling robots in various scenarios using a simple teleoperation. + +
+When using the Docker setup, it is strongly recommended that you always use the provided .docker/run.bash script. It configures the environment automatically and mounts caching volumes. You can optionally provide a command that will be executed immediately inside the container. If no command is specified, you will be dropped into an interactive shell. Throughout this documentation, if you omit the .docker/run.bash prefix, it assumes that you are already inside the Docker container, or you are using a local installation. + +```bash +# cd space_robotics_bench +.docker/run.bash ${OPTIONAL_CMD} +``` + +
+ +## Verify the Functionality of Isaac Sim + +Let's start by verifying that Isaac Sim is functioning correctly: + +
+The first time Isaac Sim starts, it may take a few minutes to compile shaders. However, subsequent runs will use cached artefacts, which significantly speed up the startup. +
+ +```bash +# Single quotes are required for the tilde (~) to expand correctly inside the container. +.docker/run.bash '~/isaac-sim/isaac-sim.sh' +``` + +If any issues arise, consult the [Troubleshooting](../misc/troubleshooting.md#runtime-errors) section or the [official Isaac Sim documentation](https://docs.omniverse.nvidia.com/isaacsim), as this issue is likely unrelated to this project. + +## Journey into the Unknown + +Once Isaac Sim is confirmed to be working, you can begin exploring the demos and tasks included with the environments. Let's start with a simple teleoperation example with the [`teleop.py`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/scripts/teleop.py) script: + +```bash +# Option 1: Using the script directly +.docker/run.bash scripts/teleop.py --env perseverance +# Option 2: Using ROS 2 package installation +.docker/run.bash ros2 run space_robotics_bench teleop.py --env perseverance +``` + +After a few moments, Isaac Sim should appear. The window will briefly remain inactive as the assets are procedurally generated in the background. The generation time depends on the complexity of the assets and your hardware, particularly the GPU, which will be used to bake PBR textures. However, future runs will use cached assets, as long as the configuration remains unchanged and the cache is not cleared, see [Clean the Assets Cache](../instructions/utils/clean_cache.md). + +Eventually, you will be greeted by the Mars Perseverance Rover on a procedurally generated Martian landscape. + +![](../_images/perseverance_ui.jpg) + +At the same time, the terminal will display the following keyboard scheme: + +``` ++------------------------------------------------+ +| Keyboard Scheme (focus the Isaac Sim window) | ++------------------------------------------------+ ++------------------------------------------------+ +| Reset: [ L ] | ++------------------------------------------------+ +| Planar Motion | +| [ W ] (+X) | +| ↑ | +| | | +| (-Y) [ A ] ← + → [ D ] (+Y) | +| | | +| ↓ | +| [ S ] (-X) | ++------------------------------------------------+ +``` + +While the Isaac Sim window is in focus, you can control the rover using the `W`, `A`, `S`, and `D` keys for motion. Use your mouse to navigate the camera. If the rover gets stuck, pressing `L` will reset its position. + +To close the demo, press `Ctrl+C` in the terminal. This will gracefully shut down the demo, close Isaac Sim, and return you to your host environment. + +### Blurry Textures? + +By default, the textures in the environment might appear blurry due to the configuration setting the baked texture resolution to 50.0% (`default=0.5`). This setting allows procedural generation to be faster on low-end hardware. If your hardware is capable, you can increase the resolution by adjusting the `detail` parameter, see [Benchmark Configuration](../instructions/benchmark/cfg.md): + +```bash +.docker/run.bash -e SRB_DETAIL=1.0 scripts/teleop.py --env perseverance +``` + +## Explore Unknown Domains + +You can explore other environments by using the `--env`, `--task`, or `--demo` arguments interchangeably. A full list of available environments is documented in the [Environment Overview](../overview/envs/index.html), or you can conveniently list them using this command: + +```bash +.docker/run.bash scripts/list_envs.py +``` + +Use this example as a **gateway** into exploring further on your own: + +```bash +.docker/run.bash scripts/teleop.py --env perseverance +``` diff --git a/docs/src/instructions/benchmark/README.md b/docs/src/instructions/benchmark/README.md new file mode 100644 index 0000000..320ff5d --- /dev/null +++ b/docs/src/instructions/benchmark/README.md @@ -0,0 +1,9 @@ +# Instructions for the Benchmark + +This section covers instructions for specific aspects of the Space Robotics Bench. + +- [Simulating Parallel Environments](./parallel_envs.md) +- [Random and Zero Agents](./random_zero_agents.md) +- [Enabling Visual Observations](./visual_observations.md) +- [Graphical User Interface](./gui.md) +- [Benchmark Configuration](./cfg.md) diff --git a/docs/src/instructions/benchmark/cfg.md b/docs/src/instructions/benchmark/cfg.md new file mode 100644 index 0000000..3d19d32 --- /dev/null +++ b/docs/src/instructions/benchmark/cfg.md @@ -0,0 +1,51 @@ +# Benchmark Configuration + +## Environment Configuration + +The environments can be configured in two ways: + +1. **Modifying the [`env.yaml`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/config/env.yaml) file**. +1. **Using environment variables**. + +The default configuration file contains various settings that control the seed, scenario, level of detail, and options for assets (robot, object, terrain, vehicle). + +```yaml +seed: 42 # SRB_SEED [int] +scenario: mars # SRB_SCENARIO [mars, moon, orbit] +detail: 0.5 # SRB_DETAIL [float] +assets: + robot: + variant: dataset # SRB_ASSETS_ROBOT_VARIANT [dataset] + object: + variant: procedural # SRB_ASSETS_OBJECT_VARIANT [primitive, dataset, procedural] + terrain: + variant: procedural # SRB_ASSETS_TERRAIN_VARIANT [none, primitive, dataset, procedural] + vehicle: + variant: dataset # SRB_ASSETS_VEHICLE_VARIANT [none, dataset] +``` + +### Setting Configuration via Environment Variables + +Values from the configuration file can be overridden using environment variables. Furthermore, you can directly pass them into the [`.docker/run.bash`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/.docker/run.bash) script. For instance: + +```bash +.docker/run.bash -e SRB_DETAIL=1.0 -e SRB_SCENARIO=moon ... +``` + +## CLI Arguments + +The following arguments are common across all entrypoint scripts, e.g. [`teleop.py`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/scripts/teleop.py), [`random_agent.py`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/scripts/random_agent.py) and [`ros2.py`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/scripts/ros2.py): + +- `-h`, `--help`: Display the help message and exit. +- `--task TASK`, `--demo TASK`, `--env TASK`: Specify the name of the task or environment. You can list available tasks using `list_envs.py`. +- `--num_envs NUM_ENVS`: Number of parallel environments to simulate. +- `--disable_ui`: Disable the majority of the Isaac Sim UI. +- `--headless`: Force the display to remain off, making the simulation headless. +- `--device DEVICE`: Set the device for simulation (e.g., `"cpu"`, `"cuda"`, or `"cuda:N"` where `N` is the device ID). + +## Additional Environment Variables + +- `SRB_SKIP_REGISTRATION` (default: `false`): When set to `"true"`|`1`, automatic registering of environments with the Gymnasium registry is disabled. This can be useful in specific deployment or testing scenarios. +- `SRB_SKIP_EXT_MOD_UPDATE` (default: `false`): When set to `"true"`|`1`, the Rust extension module will not be automatically recompiled on startup of Python entrypoint scripts. By default, this ensures that the extension module is always up-to-date with the source code. Skipping this step can be useful when the extension module never changes to reduce startup time slightly. +- `SRB_WITH_TRACEBACK` (default: `false`): When set to `"true"`|`1`, rich traceback information is displayed for exceptions. This can be useful for debugging. + - `SRB_WITH_TRACEBACK_LOCALS` (default: `false`): When set to `"true"`|`1` and `SRB_WITH_TRACEBACK` is enabled, local variables are included in the traceback information. This can be useful for debugging, but it can also be overwhelming in some cases. diff --git a/docs/src/instructions/benchmark/gui.md b/docs/src/instructions/benchmark/gui.md new file mode 100644 index 0000000..ab12e65 --- /dev/null +++ b/docs/src/instructions/benchmark/gui.md @@ -0,0 +1,15 @@ +# Graphical User Interface (GUI) + +The Space Robotics Bench comes with a simple GUI application that can serve as a more approachable demonstration of its capabilities than pure CLI. The GUI is built on top of [egui](https://github.com/emilk/egui) and leverages [r2r](https://github.com/sequenceplanner/r2r) ROS 2 Rust bindings to communicate with the rest of the benchmark. The initial screen of the GUI is shown below. + +![](../../_images/srb_gui.jpg) + +## Usage + +To run the GUI application, you can use the included [`gui.bash`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/scripts/gui.bash) script, which internally calls a variant of `cargo run -p space_robotics_bench_gui` command. + +```bash +.docker/run.bash scripts/gui.bash +``` + +The GUI runs the [`teleop.py`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/scripts/teleop.py) script for the selected environments, but the idea is to eventually support multiple workflows. Nine pre-configured tasks/demos are available in the Quick Start window, and a specific scenario can also be defined through the advanced configuration. diff --git a/docs/src/instructions/benchmark/parallel_envs.md b/docs/src/instructions/benchmark/parallel_envs.md new file mode 100644 index 0000000..fa6ed39 --- /dev/null +++ b/docs/src/instructions/benchmark/parallel_envs.md @@ -0,0 +1,13 @@ +# Simulating Parallel Environments + +## The `--num_envs` Argument + +All Python entrypoint scripts, e.g. [`teleop.py`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/scripts/teleop.py), [`random_agent.py`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/scripts/random_agent.py) and [`ros2.py`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/scripts/ros2.py), accept an optional `--num_envs` argument. By default, this is set to `1`, but you can specify more environments for parallel execution. For example, to run four environments, use the following command: + +```bash +.docker/run.bash scripts/teleop.py --task sample_collection --num_envs 4 +``` + +Each environment will generate its own procedural assets, providing unique experiences across different simulations. However, note that the time taken to generate these assets scales linearly with the number of environments. These assets will be cached for future runs unless the cache is cleared (explained later in this document). + +After the environments are initialized, they can be controlled in sync using the same keyboard scheme displayed in the terminal. diff --git a/docs/src/instructions/benchmark/random_zero_agents.md b/docs/src/instructions/benchmark/random_zero_agents.md new file mode 100644 index 0000000..fed0434 --- /dev/null +++ b/docs/src/instructions/benchmark/random_zero_agents.md @@ -0,0 +1,19 @@ +# Random and Zero Agents + +Instead of manually controlling each environment via `teleop.py`, you can use random and zero agents to test and debug certain functionalities. + +## Random Agent + +The [`random_agent.py`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/scripts/random_agent.py) script allows environments to act based on random actions sampled from the action space. This is particularly useful for verifying if environments are running as intended without manual control: + +```bash +.docker/run.bash scripts/random_agent.py --task sample_collection --num_envs 4 +``` + +## Zero Agent + +Alternatively, [`zero_agent.py`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/scripts/zero_agent.py) executes environments where all actions are zero-valued, mimicking a steady-state system. This can be useful for analyzing the idle behaviour of environments: + +```bash +.docker/run.bash scripts/zero_agent.py --task sample_collection --num_envs 4 +``` diff --git a/docs/src/instructions/benchmark/visual_observations.md b/docs/src/instructions/benchmark/visual_observations.md new file mode 100644 index 0000000..8d34c1a --- /dev/null +++ b/docs/src/instructions/benchmark/visual_observations.md @@ -0,0 +1,9 @@ +# Enabling Visual Observations + +## The `*_visual` Environment Variant + +All environments have a `*_visual` variant that differs in the enabled sensors, namely cameras that provide visual observations at the cost of increased computational requirements. + +```bash +.docker/run.bash scripts/teleop.py --task sample_collection_visual +``` diff --git a/docs/src/instructions/utils/README.md b/docs/src/instructions/utils/README.md new file mode 100644 index 0000000..4905509 --- /dev/null +++ b/docs/src/instructions/utils/README.md @@ -0,0 +1,5 @@ +# Utilities + +This section covers instructions for certain utilities of the Space Robotics Bench. + +- [Clean the Assets Cache](./clean_cache.md) diff --git a/docs/src/instructions/utils/clean_cache.md b/docs/src/instructions/utils/clean_cache.md new file mode 100644 index 0000000..bf25a3a --- /dev/null +++ b/docs/src/instructions/utils/clean_cache.md @@ -0,0 +1,7 @@ +# Clean the Assets Cache + +After running several demos or simulations, the procedurally generated assets (such as textures and meshes) can accumulate in the cache. To free up disk space, you can clean this cache: + +```bash +.docker/run.bash scripts/utils/clean_procgen_cache.py +``` diff --git a/docs/src/instructions/workflows/README.md b/docs/src/instructions/workflows/README.md new file mode 100644 index 0000000..f7ae6a8 --- /dev/null +++ b/docs/src/instructions/workflows/README.md @@ -0,0 +1,5 @@ +# Instructions for Integrated Workflows + +This section covers instructions for common workflows that are integrated directly into the Space Robotics Bench. + +- [ROS 2 Workflow](./ros2.md) diff --git a/docs/src/instructions/workflows/ros2.md b/docs/src/instructions/workflows/ros2.md new file mode 100644 index 0000000..4301e79 --- /dev/null +++ b/docs/src/instructions/workflows/ros2.md @@ -0,0 +1,99 @@ +# ROS 2 Workflow + +Environments of the Space Robotics Bench can be integrated with ROS 2 to enable control of the robots and data collection over the middleware communication. + +## Single Environment + +The [`ros2.py`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/scripts/ros2.py) script is the primary entry point for interfacing with the environments through ROS 2. This script spawns a single ROS node that maps inputs and outputs for the environment and provides miscellaneous functionalities such as resetting the simulation. Here is an example using the Ingenuity demo: + +```bash +.docker/run.bash ros2 run space_robotics_bench ros2.py --env ingenuity +``` + +Once the environment is initialized, open a new terminal to inspect the available ROS topics. You can either use your ROS setup or join the running Docker container with the [`.docker/join.bash`](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/.docker/join.bash) script: + +```bash +.docker/join.bash +``` + +Now, list the available ROS topics: + +```bash +ros2 topic list +# Expected output: +# /clock +# /env/info +# /env/reward +# /env/terminated +# /env/truncated +# /parameter_events +# /robot/cmd_vel +# /rosout +# /tf +``` + +To control the robot, publish a `Twist` message to the `/robot/cmd_vel` topic: + +```bash +ros2 topic pub --once /robot/cmd_vel geometry_msgs/msg/Twist '{linear: {x: 1.0}}' +``` + +You can reset the simulation by calling the `/sim/reset` service: + +```bash +ros2 service call /sim/reset std_srvs/srv/Empty +``` + +## Parallel Environments + +You can run multiple environments in parallel by using the `--num_envs` argument. Each environment will map to its own ROS namespace. For example, try running the Ingenuity demo with 4 environments: + +```bash +.docker/run.bash ros2 run space_robotics_bench ros2.py --env ingenuity --num_envs 4 +``` + +List the available ROS topics again: + +```bash +ros2 topic list +# Expected output: +# /clock +# /env0/reward +# /env0/robot/cmd_vel +# /env0/terminated +# /env0/truncated +# /env1/reward +# /env1/robot/cmd_vel +# /env1/terminated +# /env1/truncated +# /env2/reward +# /env2/robot/cmd_vel +# /env2/terminated +# /env2/truncated +# /env3/reward +# /env3/robot/cmd_vel +# /env3/terminated +# /env3/truncated +# /envs/info +# /envs/robot/cmd_vel +# /parameter_events +# /rosout +# /tf +``` + +Each environment has its own namespace, allowing individual control. For example: + +```bash +ros2 topic pub --once /env0/robot/cmd_vel geometry_msgs/msg/Twist '{linear: {x: -1.0}}' +ros2 topic pub --once /env1/robot/cmd_vel geometry_msgs/msg/Twist '{linear: {x: 1.0}}' +ros2 topic pub --once /env2/robot/cmd_vel geometry_msgs/msg/Twist '{linear: {y: -1.0}}' +ros2 topic pub --once /env3/robot/cmd_vel geometry_msgs/msg/Twist '{linear: {y: 1.0}}' +``` + +## Launch with `rviz2` and `teleop_twist_keyboard` + +For convenience, you can launch `rviz2` alongside `ros2.py` and `teleop_twist_keyboard` for visualization and control via keyboard: + +```bash +.docker/run.bash ros2 launch space_robotics_bench demo.launch.py task:=ingenuity_visual num_envs:=4 +``` diff --git a/docs/src/misc/attributions.md b/docs/src/misc/attributions.md new file mode 100644 index 0000000..54f576d --- /dev/null +++ b/docs/src/misc/attributions.md @@ -0,0 +1,10 @@ +# Attributions + +All modifications to the listed assets, unless stated otherwise, involve non-destructive transformations, mesh simplification, conversion to the [Universal Scene Description (USD)](https://openusd.org) format, rigging, and application of USD APIs for integration purposes. + +1. **[Mars Perseverance Rover, 3D Model](https://science.nasa.gov/resource/mars-perseverance-rover-3d-model)** by NASA. +1. **[Mars Ingenuity Helicopter, 3D Model](https://science.nasa.gov/resource/mars-ingenuity-helicopter-3d-model)** by NASA. +1. **[Mars 2020 Sample Tube, 3D Model](https://github.com/nasa/NASA-3D-Resources/blob/8780ccf7afd4e6dcdd9c0bed313354173dcd924f/3D%20Models/Mars%202020%20Sample%20Tube%203D%20print%20files/SAMPLE_TUBE.STL)** by NASA. This mesh was modified to include a cap, and additional materials were added. +1. **[Gateway Core, 3D Model](https://nasa3d.arc.nasa.gov/detail/gateway)** by NASA. The model was separated into individual assets: Canadarm3 (small and large) and the Gateway Core itself. +1. **[Low Lunar Orbit, HDR Image](https://nasa3d.arc.nasa.gov/detail/gateway)** by NASA. This image was rotated by 90 degrees to better align with the implemented environment. +1. **[Lunar Rover from the Movie *Moon*, 3D Model](https://skfb.ly/owELz)** by Watndit, licensed under [Creative Commons Attribution](http://creativecommons.org/licenses/by/4.0). The original model was heavily remodelled and retextured, with significant parts of the geometry removed. The wheels were replaced with a model from NASA's [Curiosity Rover, 3D Model](https://science.nasa.gov/resource/curiosity-rover-3d-model). diff --git a/docs/src/misc/contributors.md b/docs/src/misc/contributors.md new file mode 100644 index 0000000..8204a88 --- /dev/null +++ b/docs/src/misc/contributors.md @@ -0,0 +1,3 @@ +# Contributors + +- [Andrej Orsula](https://github.com/AndrejOrsula) diff --git a/docs/src/misc/troubleshooting.md b/docs/src/misc/troubleshooting.md new file mode 100644 index 0000000..3f6031c --- /dev/null +++ b/docs/src/misc/troubleshooting.md @@ -0,0 +1,19 @@ +# Troubleshooting + +## Runtime Errors + +### Driver Incompatibility + +If you encounter the following error message: + +```log +[Error] [carb.graphics-vulkan.plugin] VkResult: ERROR_INCOMPATIBLE_DRIVER +``` + +This indicates that your NVIDIA driver is incompatible with Omniverse. To resolve the issue, update your NVIDIA driver according to the [Isaac Sim driver requirements](https://docs.omniverse.nvidia.com/isaacsim/latest/installation/requirements.html#isaac-sim-short-driver-requirements). + +## Unexpected Behavior + +### Teleoperation Stuck + +During teleoperation, if you change your window focus, Omniverse may fail to register a button release, causing the robot to continuously move in one direction. To fix this, press the `L` key to reset the environment. diff --git a/docs/src/overview/envs/README.md b/docs/src/overview/envs/README.md new file mode 100644 index 0000000..a82229b --- /dev/null +++ b/docs/src/overview/envs/README.md @@ -0,0 +1,19 @@ +# Environments + +This section provides an overview of the environments currently available in the Space Robotics Bench. + +Before using these environments: + +1. [Ensure you meet the system requirements](../../getting_started/requirements.md) +1. [Install the benchmark](../../getting_started/installation/index.html) +1. [Learn the basic usage](../../getting_started/usage.md) + +Environments are separated into two categories: + +- **Tasks** that come with an objective for an agent to complete. Therefore, each environment instance provides a reward signal that guides the agent towards completing the goal. +- **Demos** that provide a sandbox of the included capabilities without any specific objective. These environments can serve in the initial task design while designing a new scenario or defining the action and observation spaces. + +The environments are grouped based on the robot type: + +- [Mobile Robot Environments](./mobile.md) +- [Robotic Manipulation Environments](./manipulation.md) diff --git a/docs/src/overview/envs/manipulation.md b/docs/src/overview/envs/manipulation.md new file mode 100644 index 0000000..4641454 --- /dev/null +++ b/docs/src/overview/envs/manipulation.md @@ -0,0 +1,118 @@ +# Robotic Manipulation Environments + +## Tasks + +### Sample Collection + +#### Scenario: Moon, Objects: Procedural + +```bash +.docker/run.bash -e SRB_SCENARIO=moon -e SRB_ASSETS_OBJECT_VARIANT=procedural scripts/teleop.py --env sample_collection +``` + +![](../../_images/envs/sample_collection_moon.jpg) + +#### Scenario: Mars, Objects: Dataset + +```bash +.docker/run.bash -e SRB_SCENARIO=mars -e SRB_ASSETS_OBJECT_VARIANT=dataset scripts/teleop.py --env sample_collection +``` + +![](../../_images/envs/sample_collection_mars.jpg) + +#### Other Examples + +```bash +# Scenario: Moon, Objects: Primitive +.docker/run.bash -e SRB_SCENARIO=moon -e SRB_ASSETS_OBJECT_VARIANT=primitive scripts/teleop.py --env sample_collection + +# Scenario: Mars, Objects: Procedural +.docker/run.bash -e SRB_SCENARIO=mars -e SRB_ASSETS_OBJECT_VARIANT=procedural scripts/teleop.py --env sample_collection + +# Scenario: Moon, Objects: Multi + Dataset +.docker/run.bash -e SRB_SCENARIO=orbit -e SRB_ASSETS_OBJECT_VARIANT=dataset scripts/teleop.py --env sample_collection_multi +``` + +### Debris Capture + +#### Scenario: Orbit, Objects: Dataset + +```bash +.docker/run.bash -e SRB_SCENARIO=orbit -e SRB_ASSETS_OBJECT_VARIANT=dataset scripts/teleop.py --env debris_capture +``` + +![](../../_images/envs/debris_capture_orbit.jpg) + +#### Other Examples + +```bash +# Scenario: Orbit, Objects: Procedural +.docker/run.bash -e SRB_SCENARIO=orbit -e SRB_ASSETS_OBJECT_VARIANT=procedural scripts/teleop.py --env debris_capture +``` + +### Peg-in-Hole + +#### Scenario: Moon, Objects: Dataset + +```bash +.docker/run.bash -e SRB_SCENARIO=moon -e SRB_ASSETS_OBJECT_VARIANT=dataset scripts/teleop.py --env peg_in_hole +``` + +![](../../_images/envs/peg_in_hole_moon.jpg) + +#### Scenario: Orbit, Objects: Dataset + +```bash +.docker/run.bash -e SRB_SCENARIO=orbit -e SRB_ASSETS_OBJECT_VARIANT=dataset scripts/teleop.py --env peg_in_hole +``` + +![](../../_images/envs/peg_in_hole_orbit.jpg) + +#### Other Examples + +```bash +# Scenario: Moon, Objects: Prodecural +.docker/run.bash -e SRB_SCENARIO=moon -e SRB_ASSETS_OBJECT_VARIANT=procedural scripts/teleop.py --env peg_in_hole + +# Scenario: Mars, Objects: Dataset +.docker/run.bash -e SRB_SCENARIO=mars -e SRB_ASSETS_OBJECT_VARIANT=dataset scripts/teleop.py --env peg_in_hole + +# Scenario: Moon, Objects: Multi + Dataset +.docker/run.bash -e SRB_SCENARIO=mars -e SRB_ASSETS_OBJECT_VARIANT=dataset scripts/teleop.py --env peg_in_hole_multi + +# Scenario: Mars, Objects: Multi + Procedural +.docker/run.bash -e SRB_SCENARIO=mars -e SRB_ASSETS_OBJECT_VARIANT=procedural scripts/teleop.py --env peg_in_hole_multi + +# Scenario: Orbit, Objects: Multi + Dataset +.docker/run.bash -e SRB_SCENARIO=orbit -e SRB_ASSETS_OBJECT_VARIANT=dataset scripts/teleop.py --env peg_in_hole_multi +``` + +### Solar Panel Assembly + +#### Scenario: Moon + +```bash +.docker/run.bash -e SRB_SCENARIO=moon scripts/teleop.py --env solar_panel_assembly +``` + +![](../../_images/envs/solar_panel_assembly_moon.jpg) + +#### Other Examples + +```bash +# Scenario: Mars +.docker/run.bash -e SRB_SCENARIO=mars scripts/teleop.py --env solar_panel_assembly + +# Scenario: Orbit +.docker/run.bash -e SRB_SCENARIO=orbit scripts/teleop.py --env solar_panel_assembly +``` + +## Demos + +### Gateway + +```bash +.docker/run.bash -e SRB_SCENARIO=orbit scripts/teleop.py --env gateway +``` + +![](../../_images/envs/gateway.jpg) diff --git a/docs/src/overview/envs/mobile.md b/docs/src/overview/envs/mobile.md new file mode 100644 index 0000000..b263768 --- /dev/null +++ b/docs/src/overview/envs/mobile.md @@ -0,0 +1,6 @@ +# Mobile Robot Environments + +Mobile environments are currently limited to simple demos for wheeled and aerial robots. Future plans include integrating spacecrafts in orbital scenarios and defining a set of tasks for each robot type. + +- [Wheeled Robot Environments](./mobile_wheeled.md) +- [Aerial Robot Environments](./mobile_aerial.md) diff --git a/docs/src/overview/envs/mobile_aerial.md b/docs/src/overview/envs/mobile_aerial.md new file mode 100644 index 0000000..d6bbbbb --- /dev/null +++ b/docs/src/overview/envs/mobile_aerial.md @@ -0,0 +1,15 @@ +# Aerial Robot Environments + +## Tasks + +> No tasks are implemented at the moment. + +## Demos + +### Ingenuity + +```bash +.docker/run.bash -e SRB_SCENARIO=mars scripts/teleop.py --env ingenuity +``` + +![](../../_images/envs/ingenuity.jpg) diff --git a/docs/src/overview/envs/mobile_wheeled.md b/docs/src/overview/envs/mobile_wheeled.md new file mode 100644 index 0000000..0665f5a --- /dev/null +++ b/docs/src/overview/envs/mobile_wheeled.md @@ -0,0 +1,15 @@ +# Wheeled Robot Environments + +## Tasks + +> No tasks are implemented at the moment. + +## Demos + +### Perseverance + +```bash +.docker/run.bash -e SRB_SCENARIO=mars scripts/teleop.py --env perseverance +``` + +![](../../_images/envs/perseverance.jpg) diff --git a/docs/src/overview/integrations/README.md b/docs/src/overview/integrations/README.md new file mode 100644 index 0000000..414ccbe --- /dev/null +++ b/docs/src/overview/integrations/README.md @@ -0,0 +1,5 @@ +# Integrations + +The Space Robotics Bench features a number of integrations that simplify common workflows. This section provides an overview of the design with further references for the specific instructions. + +- [Integration with ROS 2](./ros2.md) diff --git a/docs/src/overview/integrations/ros2.md b/docs/src/overview/integrations/ros2.md new file mode 100644 index 0000000..d9105ae --- /dev/null +++ b/docs/src/overview/integrations/ros2.md @@ -0,0 +1,15 @@ +# Integration with ROS 2 + +> Take a look at [ROS 2 Workflow](../../instructions/workflows/ros2.md) if you are directly interested in instructions with concrete examples. + +## Motivation + +ROS has become the de facto standard for developing robotic systems across various environments, including outer space, with the advent of ROS 2. The Space Robotics Bench integrates seamlessly with the ROS 2 ecosystem, facilitating the exposure of relevant simulation data to ROS nodes. This integration aims to accelerate the iterative development and testing of space robotic systems. + +## Approach + +Isaac Sim’s computational graph is primarily offloaded to the system’s dedicated NVIDIA GPU, which presents challenges in directly exposing all internal states to the ROS middleware without compromising performance. Instead, the package focuses on exposing the inputs and outputs of each registered Gymnasium environment alongside a fixed global mapping configuration to maintain modularity and flexibility within the simulation architecture. + +## Workflow + +The Space Robotics Bench provides a `ros2.py` script that spawns a ROS node to interface with the environments. Subscribers, publishers, and services are dynamically created based on the selected environment and global mapping configuration. When running multiple environment instances in parallel, the script automatically assigns different namespaces to inputs and outputs, preventing conflicts. The script also includes additional functionalities, such as simulation reset capabilities. diff --git a/docs/theme/favicon.png b/docs/theme/favicon.png new file mode 100644 index 0000000..36e906c Binary files /dev/null and b/docs/theme/favicon.png differ diff --git a/docs/theme/favicon.svg b/docs/theme/favicon.svg new file mode 100644 index 0000000..e86a07c --- /dev/null +++ b/docs/theme/favicon.svg @@ -0,0 +1,13 @@ + + + + + + diff --git a/hyperparams/dreamerv3.yaml b/hyperparams/dreamerv3.yaml new file mode 100644 index 0000000..340327c --- /dev/null +++ b/hyperparams/dreamerv3.yaml @@ -0,0 +1,224 @@ +seed: 42 +jax.prealloc: True + +## Input +(enc|dec).spaces: "state.*|proprio.*|image.*" +# (enc|dec).spaces: "proprio.*|image.*" +run.log_keys_video: + [ + image_scene, + image_scene_rgb, + image_scene_depth, + image_base, + image_base_rgb, + image_base_depth, + image_wrist, + image_wrist_rgb, + image_wrist_depth, + ] +run.log_video_fps: 25 + +## Run +run.steps: 100000000 +run.train_ratio: 128 + +## Replay +replay.size: 40000 # 32 GB +# replay.size: 100000 # 64 GB +# replay.size: 250000 # 128 GB +run.train_fill: 16384 +replay.chunksize: 4096 +# replay.fracs: { uniform: 0.5, priority: 0.5 } + +## Network +(enc|dec).simple.minres: 4 # 64x64 px +# Note: Continuity predictor can be small for environments with no special termination conditions +conhead.units: 1 + +## Events +run.timer: False +run.log_every: 300 +run.eval_every: 1200 +run.save_every: 1800 + +method: dreamerv3 +################## +### Model size ### +################## +# size12m: &size12m +# dyn.rssm: {deter: 2048, hidden: 256, classes: 16} +# .*\.depth: 16 +# .*\.units: 256 + +# size25m: &size25m +# dyn.rssm: {deter: 3072, hidden: 384, classes: 24} +# .*\.depth: 24 +# .*\.units: 384 + +# size50m: &size50m +# dyn.rssm: {deter: 4096, hidden: 512, classes: 32} +# .*\.depth: 32 +# .*\.units: 512 + +# size100m: &size100m +# dyn.rssm: {deter: 6144, hidden: 768, classes: 48} +# .*\.depth: 48 +# .*\.units: 768 + +# size200m: &size200m +# dyn.rssm: {deter: 8192, hidden: 1024, classes: 64} +# .*\.depth: 64 +# .*\.units: 1024 + +# size400m: &size400m +# dyn.rssm: {deter: 12288, hidden: 1536, classes: 96} +# .*\.depth: 96 +# .*\.units: 1536 + +############### +### Default ### +############### +# seed: 0 +# method: name +# task: dummy_disc +# logdir: /dev/null +# eval_dir: '' +# filter: 'score|length|fps|ratio|train/.*_loss$|train/rand/.*/mean' +# tensorboard_videos: True + +# replay: +# size: 5e6 +# online: True +# fracs: {uniform: 1.0, priority: 0.0, recency: 0.0} +# prio: {exponent: 0.8, maxfrac: 0.5, initial: inf, zero_on_sample: True} +# priosignal: model +# recexp: 1.0 +# chunksize: 1024 +# save_wait: False + +# jax: +# platform: gpu +# jit: True +# compute_dtype: bfloat16 +# param_dtype: float32 +# prealloc: True +# checks: False +# logical_cpus: 0 +# debug: False +# policy_devices: [0] +# train_devices: [0] +# sync_every: 1 +# profiler: False +# transfer_guard: True +# assert_num_devices: -1 +# fetch_policy_carry: False +# nvidia_flags: False +# xla_dump: False + +# run: +# script: train +# steps: 1e10 +# duration: 0 +# num_envs: 16 +# num_envs_eval: 4 +# expl_until: 0 +# log_every: 120 +# save_every: 900 +# eval_every: 180 +# eval_initial: True +# eval_eps: 1 +# train_ratio: 32.0 +# train_fill: 0 +# eval_fill: 0 +# log_zeros: True +# log_keys_video: [image] +# log_keys_sum: '^$' +# log_keys_avg: '^$' +# log_keys_max: '^$' +# log_video_fps: 20 +# log_video_streams: 4 +# log_episode_timeout: 60 +# from_checkpoint: '' +# actor_addr: 'tcp://localhost:{auto}' +# replay_addr: 'tcp://localhost:{auto}' +# logger_addr: 'tcp://localhost:{auto}' +# actor_batch: 8 +# actor_threads: 4 +# env_replica: -1 +# ipv6: False +# usage: {psutil: True, nvsmi: True, gputil: False, malloc: False, gc: False} +# timer: True +# driver_parallel: True +# agent_process: False +# remote_replay: False + +# wrapper: {length: 0, reset: True, discretize: 0, checks: True} +# env: +# atari: {size: [64, 64], repeat: 4, sticky: True, gray: True, actions: all, lives: unused, noops: 0, autostart: False, pooling: 2, aggregate: max, resize: pillow} +# crafter: {size: [64, 64], logs: False, use_logdir: False} +# atari100k: {size: [64, 64], repeat: 4, sticky: False, gray: False, actions: all, lives: unused, noops: 0, autostart: False, resize: pillow, length: 100000} +# dmlab: {size: [64, 64], repeat: 4, episodic: True, actions: popart, use_seed: False} +# minecraft: {size: [64, 64], break_speed: 100.0, logs: False} +# dmc: {size: [64, 64], repeat: 2, image: True, camera: -1} +# procgen: {size: [64, 64]} +# loconav: {size: [64, 64], repeat: 2, camera: -1} + +# # Agent +# report: True +# report_gradnorms: False +# batch_size: 16 +# batch_length: 65 +# batch_length_eval: 33 +# replay_length: 0 +# replay_length_eval: 0 +# replay_context: 1 +# random_agent: False +# loss_scales: {dec_cnn: 1.0, dec_mlp: 1.0, reward: 1.0, cont: 1.0, dyn: 1.0, rep: 0.1, actor: 1.0, critic: 1.0, replay_critic: 0.3} +# opt: {scaler: rms, lr: 4e-5, eps: 1e-20, momentum: True, wd: 0.0, warmup: 1000, globclip: 0.0, agc: 0.3, beta1: 0.9, beta2: 0.999, details: False, pmin: 1e-3, anneal: 0, schedule: constant} +# separate_lrs: False +# lrs: {dec: 1e-4, enc: 1e-4, dyn: 1e-4, rew: 1e-4, con: 1e-4, actor: 3e-5, critic: 3e-5} +# ac_grads: none +# reset_context: 0.0 +# replay_critic_loss: True +# replay_critic_grad: True +# replay_critic_bootstrap: imag +# reward_grad: True +# report_openl_context: 8 + +# # World Model +# dyn: +# typ: rssm +# rssm: {deter: 8192, hidden: 1024, stoch: 32, classes: 64, act: silu, norm: rms, unimix: 0.01, outscale: 1.0, winit: normal, imglayers: 2, obslayers: 1, dynlayers: 1, absolute: False, cell: blockgru, blocks: 8, block_fans: False, block_norm: False} +# enc: +# spaces: '.*' +# typ: simple +# simple: {depth: 64, mults: [1, 2, 3, 4, 4], layers: 3, units: 1024, act: silu, norm: rms, winit: normal, symlog: True, outer: True, kernel: 5, minres: 4} +# dec: +# spaces: '.*' +# typ: simple +# simple: {inputs: [deter, stoch], vecdist: symlog_mse, depth: 64, mults: [1, 2, 3, 4, 4], layers: 3, units: 1024, act: silu, norm: rms, outscale: 1.0, winit: normal, outer: True, kernel: 5, minres: 4, block_space: 8, block_fans: False, block_norm: False, hidden_stoch: True, space_hidden: 0} +# rewhead: {layers: 1, units: 1024, act: silu, norm: rms, dist: symexp_twohot, outscale: 0.0, inputs: [deter, stoch], winit: normal, bins: 255, block_fans: False, block_norm: False} +# conhead: {layers: 1, units: 1024, act: silu, norm: rms, dist: binary, outscale: 1.0, inputs: [deter, stoch], winit: normal, block_fans: False, block_norm: False} +# contdisc: True +# rssm_loss: {free: 1.0} + +# # Actor Critic +# actor: {layers: 3, units: 1024, act: silu, norm: rms, minstd: 0.1, maxstd: 1.0, outscale: 0.01, unimix: 0.01, inputs: [deter, stoch], winit: normal, block_fans: False, block_norm: False} +# critic: {layers: 3, units: 1024, act: silu, norm: rms, dist: symexp_twohot, outscale: 0.0, inputs: [deter, stoch], winit: normal, bins: 255, block_fans: False, block_norm: False} +# actor_dist_disc: onehot +# actor_dist_cont: normal +# imag_start: all +# imag_repeat: 1 +# imag_length: 15 +# imag_unroll: False +# horizon: 333 +# return_lambda: 0.95 +# return_lambda_replay: 0.95 +# slow_critic_update: 1 +# slow_critic_fraction: 0.02 +# retnorm: {impl: perc, rate: 0.01, limit: 1.0, perclo: 5.0, perchi: 95.0} +# valnorm: {impl: off, rate: 0.01, limit: 1e-8} +# advnorm: {impl: off, rate: 0.01, limit: 1e-8} +# actent: 3e-4 +# slowreg: 1.0 +# slowtar: False diff --git a/hyperparams/robomimic/bc.json b/hyperparams/robomimic/bc.json new file mode 100644 index 0000000..b0cf585 --- /dev/null +++ b/hyperparams/robomimic/bc.json @@ -0,0 +1,217 @@ +{ + "algo_name": "bc", + "experiment": { + "name": "bc", + "validate": true, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 50, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": true, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": false, + "n": 50, + "horizon": 400, + "rate": 50, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir": "../bc_trained_models", + "num_data_workers": 0, + "hdf5_cache_mode": "all", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": false, + "hdf5_normalize_obs": false, + "hdf5_filter_key": "train", + "hdf5_validation_filter_key": "valid", + "seq_length": 1, + "pad_seq_length": true, + "frame_stack": 1, + "pad_frame_stack": true, + "dataset_keys": [ + "actions", + "rewards", + "dones" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 100, + "num_epochs": 2000, + "seed": 1 + }, + "algo": { + "optim_params": { + "policy": { + "optimizer_type": "adam", + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.1, + "epoch_schedule": [], + "scheduler_type": "multistep" + }, + "regularization": { + "L2": 0.0 + } + } + }, + "loss": { + "l2_weight": 1.0, + "l1_weight": 0.0, + "cos_weight": 0.0 + }, + "actor_layer_dims": [ + 1024, + 1024 + ], + "gaussian": { + "enabled": false, + "fixed_std": false, + "init_std": 0.1, + "min_std": 0.01, + "std_activation": "softplus", + "low_noise_eval": true + }, + "gmm": { + "enabled": false, + "num_modes": 5, + "min_std": 0.0001, + "std_activation": "softplus", + "low_noise_eval": true + }, + "vae": { + "enabled": false, + "latent_dim": 14, + "latent_clip": null, + "kl_weight": 1.0, + "decoder": { + "is_conditioned": true, + "reconstruction_sum_across_elements": false + }, + "prior": { + "learn": false, + "is_conditioned": false, + "use_gmm": false, + "gmm_num_modes": 10, + "gmm_learn_weights": false, + "use_categorical": false, + "categorical_dim": 10, + "categorical_gumbel_softmax_hard": false, + "categorical_init_temp": 1.0, + "categorical_temp_anneal_step": 0.001, + "categorical_min_temp": 0.3 + }, + "encoder_layer_dims": [ + 300, + 400 + ], + "decoder_layer_dims": [ + 300, + 400 + ], + "prior_layer_dims": [ + 300, + 400 + ] + }, + "rnn": { + "enabled": false, + "horizon": 10, + "hidden_dim": 400, + "rnn_type": "LSTM", + "num_layers": 2, + "open_loop": false, + "kwargs": { + "bidirectional": false + } + }, + "transformer": { + "enabled": false, + "context_length": 10, + "embed_dim": 512, + "num_layers": 6, + "num_heads": 8, + "emb_dropout": 0.1, + "attn_dropout": 0.1, + "block_output_dropout": 0.1, + "sinusoidal_embedding": false, + "activation": "gelu", + "supervise_all_steps": false, + "nn_parameter_for_timesteps": true + } + }, + "observation": { + "modalities": { + "obs": { + "low_dim": [ + "joint_pos", + "robot_ee_pose", + "sample_pose_absolute", + "robot_sample_pose_relative", + "robot_hole_pose_relative", + "sample_hole_pose_relative" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + }, + "meta": { + "hp_base_config_file": null, + "hp_keys": [], + "hp_values": [] + } +} diff --git a/hyperparams/robomimic/bcq.json b/hyperparams/robomimic/bcq.json new file mode 100644 index 0000000..31cb98a --- /dev/null +++ b/hyperparams/robomimic/bcq.json @@ -0,0 +1,237 @@ +{ + "algo_name": "bcq", + "experiment": { + "name": "bcq", + "validate": true, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 50, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": true, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": false, + "n": 50, + "horizon": 400, + "rate": 50, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir": "../bcq_trained_models", + "num_data_workers": 0, + "hdf5_cache_mode": "all", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": true, + "hdf5_normalize_obs": false, + "hdf5_filter_key": "train", + "hdf5_validation_filter_key": "valid", + "seq_length": 1, + "pad_seq_length": true, + "frame_stack": 1, + "pad_frame_stack": true, + "dataset_keys": [ + "actions", + "rewards", + "dones" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 100, + "num_epochs": 2000, + "seed": 1 + }, + "algo": { + "optim_params": { + "critic": { + "learning_rate": { + "initial": 0.001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + }, + "start_epoch": -1, + "end_epoch": -1 + }, + "action_sampler": { + "learning_rate": { + "initial": 0.001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + }, + "start_epoch": -1, + "end_epoch": -1 + }, + "actor": { + "learning_rate": { + "initial": 0.001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + }, + "start_epoch": -1, + "end_epoch": -1 + } + }, + "discount": 0.99, + "n_step": 1, + "target_tau": 0.005, + "infinite_horizon": false, + "critic": { + "use_huber": false, + "max_gradient_norm": null, + "value_bounds": null, + "num_action_samples": 10, + "num_action_samples_rollout": 100, + "ensemble": { + "n": 2, + "weight": 0.75 + }, + "distributional": { + "enabled": false, + "num_atoms": 51 + }, + "layer_dims": [ + 300, + 400 + ] + }, + "action_sampler": { + "actor_layer_dims": [ + 1024, + 1024 + ], + "gmm": { + "enabled": false, + "num_modes": 5, + "min_std": 0.0001, + "std_activation": "softplus", + "low_noise_eval": true + }, + "vae": { + "enabled": false, + "latent_dim": 14, + "latent_clip": null, + "kl_weight": 1.0, + "decoder": { + "is_conditioned": true, + "reconstruction_sum_across_elements": false + }, + "prior": { + "learn": false, + "is_conditioned": false, + "use_gmm": false, + "gmm_num_modes": 10, + "gmm_learn_weights": false, + "use_categorical": false, + "categorical_dim": 10, + "categorical_gumbel_softmax_hard": false, + "categorical_init_temp": 1.0, + "categorical_temp_anneal_step": 0.001, + "categorical_min_temp": 0.3 + }, + "encoder_layer_dims": [ + 300, + 400 + ], + "decoder_layer_dims": [ + 300, + 400 + ], + "prior_layer_dims": [ + 300, + 400 + ] + }, + "freeze_encoder_epoch": -1 + }, + "actor": { + "enabled": false, + "perturbation_scale": 0.05, + "layer_dims": [ + 300, + 400 + ] + } + }, + "observation": { + "modalities": { + "obs": { + "low_dim": [ + "joint_pos", + "robot_ee_pose", + "sample_pose_absolute", + "robot_sample_pose_relative", + "robot_hole_pose_relative", + "sample_hole_pose_relative" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + }, + "meta": { + "hp_base_config_file": null, + "hp_keys": [], + "hp_values": [] + } +} diff --git a/hyperparams/robomimic/cql.json b/hyperparams/robomimic/cql.json new file mode 100644 index 0000000..e206fd5 --- /dev/null +++ b/hyperparams/robomimic/cql.json @@ -0,0 +1,184 @@ +{ + "algo_name": "cql", + "experiment": { + "name": "cql", + "validate": true, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 50, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": true, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": false, + "n": 50, + "horizon": 400, + "rate": 50, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir": "../cql_trained_models", + "num_data_workers": 0, + "hdf5_cache_mode": "all", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": true, + "hdf5_normalize_obs": false, + "hdf5_filter_key": "train", + "hdf5_validation_filter_key": "valid", + "seq_length": 1, + "pad_seq_length": true, + "frame_stack": 1, + "pad_frame_stack": true, + "dataset_keys": [ + "actions", + "rewards", + "dones" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 1024, + "num_epochs": 2000, + "seed": 1 + }, + "algo": { + "optim_params": { + "critic": { + "learning_rate": { + "initial": 0.001, + "decay_factor": 0.0, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + } + }, + "actor": { + "learning_rate": { + "initial": 0.0003, + "decay_factor": 0.0, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + } + } + }, + "discount": 0.99, + "n_step": 1, + "target_tau": 0.005, + "actor": { + "bc_start_steps": 0, + "target_entropy": "default", + "max_gradient_norm": null, + "net": { + "type": "gaussian", + "common": { + "std_activation": "exp", + "use_tanh": true, + "low_noise_eval": true + }, + "gaussian": { + "init_last_fc_weight": 0.001, + "init_std": 0.3, + "fixed_std": false + } + }, + "layer_dims": [ + 300, + 400 + ] + }, + "critic": { + "use_huber": false, + "max_gradient_norm": null, + "value_bounds": null, + "num_action_samples": 1, + "cql_weight": 1.0, + "deterministic_backup": true, + "min_q_weight": 1.0, + "target_q_gap": 5.0, + "num_random_actions": 10, + "ensemble": { + "n": 2 + }, + "layer_dims": [ + 300, + 400 + ] + } + }, + "observation": { + "modalities": { + "obs": { + "low_dim": [ + "joint_pos", + "robot_ee_pose", + "sample_pose_absolute", + "robot_sample_pose_relative", + "robot_hole_pose_relative", + "sample_hole_pose_relative" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + }, + "meta": { + "hp_base_config_file": null, + "hp_keys": [], + "hp_values": [] + } +} diff --git a/hyperparams/robomimic/gl.json b/hyperparams/robomimic/gl.json new file mode 100644 index 0000000..8b716a9 --- /dev/null +++ b/hyperparams/robomimic/gl.json @@ -0,0 +1,186 @@ +{ + "algo_name": "gl", + "experiment": { + "name": "gl", + "validate": true, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 50, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": true, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": false, + "n": 50, + "horizon": 400, + "rate": 50, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir": "../gl_trained_models", + "num_data_workers": 0, + "hdf5_cache_mode": "all", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": true, + "hdf5_normalize_obs": false, + "hdf5_filter_key": "train", + "hdf5_validation_filter_key": "valid", + "seq_length": 1, + "pad_seq_length": true, + "frame_stack": 1, + "pad_frame_stack": true, + "dataset_keys": [ + "actions", + "rewards", + "dones" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 100, + "num_epochs": 2000, + "seed": 1 + }, + "algo": { + "optim_params": { + "goal_network": { + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + } + } + }, + "subgoal_horizon": 10, + "ae": { + "planner_layer_dims": [ + 300, + 400 + ] + }, + "vae": { + "enabled": false, + "latent_dim": 16, + "latent_clip": null, + "kl_weight": 1.0, + "decoder": { + "is_conditioned": true, + "reconstruction_sum_across_elements": false + }, + "prior": { + "learn": false, + "is_conditioned": false, + "use_gmm": false, + "gmm_num_modes": 10, + "gmm_learn_weights": false, + "use_categorical": false, + "categorical_dim": 10, + "categorical_gumbel_softmax_hard": false, + "categorical_init_temp": 1.0, + "categorical_temp_anneal_step": 0.001, + "categorical_min_temp": 0.3 + }, + "encoder_layer_dims": [ + 300, + 400 + ], + "decoder_layer_dims": [ + 300, + 400 + ], + "prior_layer_dims": [ + 300, + 400 + ] + } + }, + "observation": { + "modalities": { + "obs": { + "low_dim": [ + "joint_pos", + "robot_ee_pose", + "sample_pose_absolute", + "robot_sample_pose_relative", + "robot_hole_pose_relative", + "sample_hole_pose_relative" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + }, + "subgoal": { + "low_dim": [ + "joint_pos", + "robot_ee_pose", + "sample_pose_absolute", + "robot_sample_pose_relative", + "robot_hole_pose_relative", + "sample_hole_pose_relative" + ], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + }, + "meta": { + "hp_base_config_file": null, + "hp_keys": [], + "hp_values": [] + } +} diff --git a/hyperparams/robomimic/hbc.json b/hyperparams/robomimic/hbc.json new file mode 100644 index 0000000..3d4f45b --- /dev/null +++ b/hyperparams/robomimic/hbc.json @@ -0,0 +1,293 @@ +{ + "algo_name": "hbc", + "experiment": { + "name": "hbc", + "validate": true, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 50, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": true, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": false, + "n": 50, + "horizon": 400, + "rate": 50, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir": "../hbc_trained_models", + "num_data_workers": 0, + "hdf5_cache_mode": "all", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": true, + "hdf5_normalize_obs": false, + "hdf5_filter_key": "train", + "hdf5_validation_filter_key": "valid", + "seq_length": 10, + "pad_seq_length": true, + "frame_stack": 1, + "pad_frame_stack": true, + "dataset_keys": [ + "actions", + "rewards", + "dones" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 100, + "num_epochs": 2000, + "seed": 1 + }, + "algo": { + "mode": "separate", + "actor_use_random_subgoals": false, + "subgoal_update_interval": 10, + "latent_subgoal": { + "enabled": false, + "prior_correction": { + "enabled": false, + "num_samples": 100 + } + }, + "planner": { + "optim_params": { + "goal_network": { + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + } + } + }, + "subgoal_horizon": 10, + "ae": { + "planner_layer_dims": [ + 300, + 400 + ] + }, + "vae": { + "enabled": false, + "latent_dim": 16, + "latent_clip": null, + "kl_weight": 1.0, + "decoder": { + "is_conditioned": true, + "reconstruction_sum_across_elements": false + }, + "prior": { + "learn": false, + "is_conditioned": false, + "use_gmm": false, + "gmm_num_modes": 10, + "gmm_learn_weights": false, + "use_categorical": false, + "categorical_dim": 10, + "categorical_gumbel_softmax_hard": false, + "categorical_init_temp": 1.0, + "categorical_temp_anneal_step": 0.001, + "categorical_min_temp": 0.3 + }, + "encoder_layer_dims": [ + 300, + 400 + ], + "decoder_layer_dims": [ + 300, + 400 + ], + "prior_layer_dims": [ + 300, + 400 + ] + } + }, + "actor": { + "optim_params": { + "policy": { + "optimizer_type": "adam", + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.1, + "epoch_schedule": [], + "scheduler_type": "multistep" + }, + "regularization": { + "L2": 0.0 + } + } + }, + "loss": { + "l2_weight": 1.0, + "l1_weight": 0.0, + "cos_weight": 0.0 + }, + "actor_layer_dims": [ + 1024, + 1024 + ], + "rnn": { + "enabled": false, + "horizon": 10, + "hidden_dim": 400, + "rnn_type": "LSTM", + "num_layers": 2, + "open_loop": false, + "kwargs": { + "bidirectional": false + } + }, + "transformer": { + "enabled": false, + "context_length": 10, + "embed_dim": 512, + "num_layers": 6, + "num_heads": 8, + "emb_dropout": 0.1, + "attn_dropout": 0.1, + "block_output_dropout": 0.1, + "sinusoidal_embedding": false, + "activation": "gelu", + "supervise_all_steps": false, + "nn_parameter_for_timesteps": true + } + } + }, + "observation": { + "planner": { + "modalities": { + "obs": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + }, + "subgoal": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + }, + "actor": { + "modalities": { + "obs": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + } + }, + "meta": { + "hp_base_config_file": null, + "hp_keys": [], + "hp_values": [] + } +} diff --git a/hyperparams/robomimic/iql.json b/hyperparams/robomimic/iql.json new file mode 100644 index 0000000..4b764de --- /dev/null +++ b/hyperparams/robomimic/iql.json @@ -0,0 +1,194 @@ +{ + "algo_name": "iql", + "experiment": { + "name": "iql", + "validate": true, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 50, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": true, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": false, + "n": 50, + "horizon": 400, + "rate": 50, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir": "../iql_trained_models", + "num_data_workers": 0, + "hdf5_cache_mode": "all", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": true, + "hdf5_normalize_obs": false, + "hdf5_filter_key": "train", + "hdf5_validation_filter_key": "valid", + "seq_length": 1, + "pad_seq_length": true, + "frame_stack": 1, + "pad_frame_stack": true, + "dataset_keys": [ + "actions", + "rewards", + "dones" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 100, + "num_epochs": 2000, + "seed": 1 + }, + "algo": { + "optim_params": { + "critic": { + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.0, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + } + }, + "vf": { + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.0, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + } + }, + "actor": { + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.0, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + } + } + }, + "discount": 0.99, + "target_tau": 0.01, + "actor": { + "net": { + "type": "gaussian", + "common": { + "std_activation": "softplus", + "low_noise_eval": true, + "use_tanh": false + }, + "gaussian": { + "init_last_fc_weight": 0.001, + "init_std": 0.3, + "fixed_std": false + }, + "gmm": { + "num_modes": 5, + "min_std": 0.0001 + } + }, + "layer_dims": [ + 300, + 400 + ], + "max_gradient_norm": null + }, + "critic": { + "ensemble": { + "n": 2 + }, + "layer_dims": [ + 300, + 400 + ], + "use_huber": false, + "max_gradient_norm": null + }, + "adv": { + "clip_adv_value": null, + "beta": 1.0, + "use_final_clip": true + }, + "vf_quantile": 0.9 + }, + "observation": { + "modalities": { + "obs": { + "low_dim": [ + "joint_pos", + "robot_ee_pose", + "sample_pose_absolute", + "robot_sample_pose_relative", + "robot_hole_pose_relative", + "sample_hole_pose_relative" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + }, + "meta": { + "hp_base_config_file": null, + "hp_keys": [], + "hp_values": [] + } +} diff --git a/hyperparams/robomimic/iris.json b/hyperparams/robomimic/iris.json new file mode 100644 index 0000000..cb8ffdd --- /dev/null +++ b/hyperparams/robomimic/iris.json @@ -0,0 +1,465 @@ +{ + "algo_name": "iris", + "experiment": { + "name": "iris", + "validate": true, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 50, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": true, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": false, + "n": 50, + "horizon": 400, + "rate": 50, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir": "../iris_trained_models", + "num_data_workers": 0, + "hdf5_cache_mode": "all", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": true, + "hdf5_normalize_obs": false, + "hdf5_filter_key": "train", + "hdf5_validation_filter_key": "valid", + "seq_length": 10, + "pad_seq_length": true, + "frame_stack": 1, + "pad_frame_stack": true, + "dataset_keys": [ + "actions", + "rewards", + "dones" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 100, + "num_epochs": 2000, + "seed": 1 + }, + "algo": { + "mode": "separate", + "actor_use_random_subgoals": false, + "subgoal_update_interval": 10, + "latent_subgoal": { + "enabled": false, + "prior_correction": { + "enabled": false, + "num_samples": 100 + } + }, + "value_planner": { + "planner": { + "optim_params": { + "goal_network": { + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + } + } + }, + "subgoal_horizon": 10, + "ae": { + "planner_layer_dims": [ + 300, + 400 + ] + }, + "vae": { + "enabled": false, + "latent_dim": 16, + "latent_clip": null, + "kl_weight": 1.0, + "decoder": { + "is_conditioned": true, + "reconstruction_sum_across_elements": false + }, + "prior": { + "learn": false, + "is_conditioned": false, + "use_gmm": false, + "gmm_num_modes": 10, + "gmm_learn_weights": false, + "use_categorical": false, + "categorical_dim": 10, + "categorical_gumbel_softmax_hard": false, + "categorical_init_temp": 1.0, + "categorical_temp_anneal_step": 0.001, + "categorical_min_temp": 0.3 + }, + "encoder_layer_dims": [ + 300, + 400 + ], + "decoder_layer_dims": [ + 300, + 400 + ], + "prior_layer_dims": [ + 300, + 400 + ] + } + }, + "value": { + "optim_params": { + "critic": { + "learning_rate": { + "initial": 0.001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + }, + "start_epoch": -1, + "end_epoch": -1 + }, + "action_sampler": { + "learning_rate": { + "initial": 0.001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + }, + "start_epoch": -1, + "end_epoch": -1 + }, + "actor": { + "learning_rate": { + "initial": 0.001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + }, + "start_epoch": -1, + "end_epoch": -1 + } + }, + "discount": 0.99, + "n_step": 1, + "target_tau": 0.005, + "infinite_horizon": false, + "critic": { + "use_huber": false, + "max_gradient_norm": null, + "value_bounds": null, + "num_action_samples": 10, + "num_action_samples_rollout": 100, + "ensemble": { + "n": 2, + "weight": 0.75 + }, + "distributional": { + "enabled": false, + "num_atoms": 51 + }, + "layer_dims": [ + 300, + 400 + ] + }, + "action_sampler": { + "actor_layer_dims": [ + 1024, + 1024 + ], + "gmm": { + "enabled": false, + "num_modes": 5, + "min_std": 0.0001, + "std_activation": "softplus", + "low_noise_eval": true + }, + "vae": { + "enabled": false, + "latent_dim": 14, + "latent_clip": null, + "kl_weight": 1.0, + "decoder": { + "is_conditioned": true, + "reconstruction_sum_across_elements": false + }, + "prior": { + "learn": false, + "is_conditioned": false, + "use_gmm": false, + "gmm_num_modes": 10, + "gmm_learn_weights": false, + "use_categorical": false, + "categorical_dim": 10, + "categorical_gumbel_softmax_hard": false, + "categorical_init_temp": 1.0, + "categorical_temp_anneal_step": 0.001, + "categorical_min_temp": 0.3 + }, + "encoder_layer_dims": [ + 300, + 400 + ], + "decoder_layer_dims": [ + 300, + 400 + ], + "prior_layer_dims": [ + 300, + 400 + ] + }, + "freeze_encoder_epoch": -1 + }, + "actor": { + "enabled": false, + "perturbation_scale": 0.05, + "layer_dims": [ + 300, + 400 + ] + } + }, + "num_samples": 100 + }, + "actor": { + "optim_params": { + "policy": { + "optimizer_type": "adam", + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.1, + "epoch_schedule": [], + "scheduler_type": "multistep" + }, + "regularization": { + "L2": 0.0 + } + } + }, + "loss": { + "l2_weight": 1.0, + "l1_weight": 0.0, + "cos_weight": 0.0 + }, + "actor_layer_dims": [ + 1024, + 1024 + ], + "rnn": { + "enabled": false, + "horizon": 10, + "hidden_dim": 400, + "rnn_type": "LSTM", + "num_layers": 2, + "open_loop": false, + "kwargs": { + "bidirectional": false + } + }, + "transformer": { + "enabled": false, + "context_length": 10, + "embed_dim": 512, + "num_layers": 6, + "num_heads": 8, + "emb_dropout": 0.1, + "attn_dropout": 0.1, + "block_output_dropout": 0.1, + "sinusoidal_embedding": false, + "activation": "gelu", + "supervise_all_steps": false, + "nn_parameter_for_timesteps": true + } + } + }, + "observation": { + "value_planner": { + "planner": { + "modalities": { + "obs": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + }, + "subgoal": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + }, + "value": { + "modalities": { + "obs": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + } + }, + "actor": { + "modalities": { + "obs": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + } + }, + "meta": { + "hp_base_config_file": null, + "hp_keys": [], + "hp_values": [] + } +} diff --git a/hyperparams/robomimic/td3_bc.json b/hyperparams/robomimic/td3_bc.json new file mode 100644 index 0000000..c0dbb63 --- /dev/null +++ b/hyperparams/robomimic/td3_bc.json @@ -0,0 +1,167 @@ +{ + "algo_name": "td3_bc", + "experiment": { + "name": "td3_bc", + "validate": true, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 20, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": true, + "on_best_rollout_success_rate": false + }, + "epoch_every_n_steps": 5000, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": false, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": false, + "n": 50, + "horizon": 1000, + "rate": 1, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir": "../td3_bc_trained_models", + "num_data_workers": 0, + "hdf5_cache_mode": "all", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": true, + "hdf5_normalize_obs": true, + "hdf5_filter_key": "train", + "hdf5_validation_filter_key": "valid", + "seq_length": 1, + "pad_seq_length": true, + "frame_stack": 1, + "pad_frame_stack": true, + "dataset_keys": [ + "actions", + "rewards", + "dones" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 256, + "num_epochs": 200, + "seed": 1 + }, + "algo": { + "optim_params": { + "critic": { + "learning_rate": { + "initial": 0.0003, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + }, + "start_epoch": -1, + "end_epoch": -1 + }, + "actor": { + "learning_rate": { + "initial": 0.0003, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + }, + "start_epoch": -1, + "end_epoch": -1 + } + }, + "alpha": 2.5, + "discount": 0.99, + "n_step": 1, + "target_tau": 0.005, + "infinite_horizon": false, + "critic": { + "use_huber": false, + "max_gradient_norm": null, + "value_bounds": null, + "ensemble": { + "n": 2, + "weight": 1.0 + }, + "layer_dims": [ + 256, + 256 + ] + }, + "actor": { + "update_freq": 2, + "noise_std": 0.2, + "noise_clip": 0.5, + "layer_dims": [ + 256, + 256 + ] + } + }, + "observation": { + "modalities": { + "obs": { + "low_dim": [ + "flat" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + }, + "meta": { + "hp_base_config_file": null, + "hp_keys": [], + "hp_values": [] + } +} diff --git a/hyperparams/sb3/ppo.yaml b/hyperparams/sb3/ppo.yaml new file mode 100644 index 0000000..666d82f --- /dev/null +++ b/hyperparams/sb3/ppo.yaml @@ -0,0 +1,22 @@ +# Reference: https://github.com/DLR-RM/rl-baselines3-zoo/blob/27e081eb24419ee843ae1c329b0482db823c9fc1/hyperparams/ppo.yml +seed: 42 + +policy: "MlpPolicy" +n_timesteps: 200000000 +batch_size: 256 +n_steps: 512 +gamma: 0.99 +learning_rate: lin_5e-5 +ent_coef: 0.002 +clip_range: 0.3 +n_epochs: 8 +gae_lambda: 0.9 +max_grad_norm: 2 +vf_coef: 0.4 +policy_kwargs: "dict( + log_std_init=-2, + ortho_init=False, + activation_fn=nn.ReLU, + share_features_extractor=False, + net_arch=dict(pi=[2048, 2048], vf=[2048, 2048]) + )" diff --git a/hyperparams/sb3/ppo_lstm.yaml b/hyperparams/sb3/ppo_lstm.yaml new file mode 100644 index 0000000..d63b4a5 --- /dev/null +++ b/hyperparams/sb3/ppo_lstm.yaml @@ -0,0 +1,22 @@ +# Reference: https://github.com/DLR-RM/rl-baselines3-zoo/blob/27e081eb24419ee843ae1c329b0482db823c9fc1/hyperparams/ppo.yml +seed: 42 + +policy: "MlpLstmPolicy" +n_timesteps: 200000000 +batch_size: 256 +n_steps: 512 +gamma: 0.99 +learning_rate: lin_5e-5 +ent_coef: 0.002 +clip_range: 0.3 +n_epochs: 8 +gae_lambda: 0.9 +max_grad_norm: 2 +vf_coef: 0.4 +policy_kwargs: "dict( + log_std_init=-2, + ortho_init=False, + activation_fn=nn.ReLU, + share_features_extractor=False, + net_arch=dict(pi=[2048, 2048], vf=[2048, 2048]) + )" diff --git a/hyperparams/sb3/sac.yaml b/hyperparams/sb3/sac.yaml new file mode 100644 index 0000000..733c4ea --- /dev/null +++ b/hyperparams/sb3/sac.yaml @@ -0,0 +1,20 @@ +# Reference: https://github.com/DLR-RM/rl-baselines3-zoo/blob/27e081eb24419ee843ae1c329b0482db823c9fc1/hyperparams/sac.yml +seed: 42 + +policy: "MlpPolicy" +n_timesteps: 200000000 +learning_starts: 10000 +buffer_size: 2000000 +train_freq: 1 +gradient_steps: 16 +batch_size: 4096 +gamma: 0.998 +tau: 0.005 +learning_rate: lin_3e-4 +ent_coef: "auto" +target_entropy: "auto" +policy_kwargs: "dict( + log_std_init=-2, + share_features_extractor=False, + net_arch=dict(pi=[2048, 2048], qf=[2048, 2048]) + )" diff --git a/launch/demo.launch.py b/launch/demo.launch.py new file mode 100755 index 0000000..bec0bf1 --- /dev/null +++ b/launch/demo.launch.py @@ -0,0 +1,141 @@ +#!/usr/bin/env -S ros2 launch +""" +Interactive demo showcasing the Space Robotics Bench + +This script launches a procedural simulation environment via `../scripts/ros2.py` +while also conveniently initializing a teleoperation node for controlling the robot +via keyboard, together with RViz2 for basic visualization. + +A basic set of arguments can be passed to the script to select the task and customize +the simulation environment. For more advanced configurations, please refer to the script itself. + +Examples: + ros2 launch space_robotics_bench demo.launch.py + ros2 launch space_robotics_bench demo.launch.py task:=sample_collection num_envs:=16 +""" + +from os import path +from typing import List + +from launch_ros.actions import Node + +from launch import LaunchDescription +from launch.actions import DeclareLaunchArgument, ExecuteProcess +from launch.conditions import IfCondition +from launch.substitutions import ( + EnvironmentVariable, + LaunchConfiguration, + PathJoinSubstitution, +) + + +def generate_launch_description() -> LaunchDescription: + # Declare all launch arguments + declared_arguments = generate_declared_arguments() + + # Get substitution for all arguments + task = LaunchConfiguration("task") + num_envs = LaunchConfiguration("num_envs") + enable_rviz = LaunchConfiguration("enable_rviz") + rviz_config = LaunchConfiguration("rviz_config") + use_sim_time = LaunchConfiguration("use_sim_time") + log_level = LaunchConfiguration("log_level") + + process = [ + ExecuteProcess( + cmd=[ + PathJoinSubstitution( + [EnvironmentVariable("ISAAC_SIM_PYTHON", default_value="python3")] + ), + path.join( + path.dirname(path.dirname(path.realpath(__file__))), + "scripts", + "ros2.py", + ), + "--task", + task, + "--num_envs", + num_envs, + "--disable_ui", + ], + output="screen", + shell=True, + emulate_tty=True, + ) + ] + nodes = [ + Node( + package="teleop_twist_keyboard", + executable="teleop_twist_keyboard", + output="screen", + arguments=[ + "--ros-args", + "--log-level", + log_level, + ], + remappings=[("cmd_vel", "/robot/cmd_vel")], + parameters=[{"use_sim_time": use_sim_time}], + prefix="xterm -e", + ), + Node( + package="rviz2", + executable="rviz2", + output="log", + arguments=[ + "--display-config", + rviz_config, + "--ros-args", + "--log-level", + log_level, + ], + parameters=[ + {"use_sim_time": use_sim_time}, + ], + condition=IfCondition(enable_rviz), + ), + ] + + return LaunchDescription(declared_arguments + process + nodes) + + +def generate_declared_arguments() -> List[DeclareLaunchArgument]: + """ + Generate list of all launch arguments that are declared for this launch script. + """ + + return [ + DeclareLaunchArgument( + "task", + default_value="srb/perseverance", + description="Name of the demo/task.", + ), + DeclareLaunchArgument( + "num_envs", + default_value="1", + description="Number of environments to simulate.", + ), + # Miscellaneous + DeclareLaunchArgument( + "enable_rviz", default_value="true", description="Flag to enable RViz2." + ), + DeclareLaunchArgument( + "rviz_config", + default_value=path.join( + path.dirname(path.dirname(path.realpath(__file__))), + "config", + "rviz", + "default.rviz", + ), + description="Path to configuration for RViz2.", + ), + DeclareLaunchArgument( + "use_sim_time", + default_value="true", + description="If true, use simulated clock.", + ), + DeclareLaunchArgument( + "log_level", + default_value="warn", + description="The level of logging that is applied to all ROS 2 nodes launched by this script.", + ), + ] diff --git a/package.xml b/package.xml new file mode 100644 index 0000000..91ce276 --- /dev/null +++ b/package.xml @@ -0,0 +1,24 @@ + + + + space_robotics_bench + 0.0.1 + Comprehensive benchmark for space robotics + Andrej Orsula + MIT OR Apache-2.0 + + https://github.com/AndrejOrsula/space_robotics_bench + + Andrej Orsula + + ament_cmake + ament_cmake_python + + rviz2 + teleop_twist_keyboard + xterm + + + ament_cmake + + diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..54b7f58 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,57 @@ +[build-system] +requires = ["maturin>=1.7,<2"] +build-backend = "maturin" + +[tool.maturin] +features = ["pyo3/extension-module"] +module-name = "space_robotics_bench._rs" +manifest-path = "crates/space_robotics_bench_py/Cargo.toml" + +[tool.pytest.ini_options] +env = ["SRB_SKIP_REGISTRATION=true"] + +[project] +name = "space_robotics_bench" +description = "Comprehensive benchmark for space robotics" +authors = [{ name = "Andrej Orsula", email = "orsula.andrej@gmail.com" }] +maintainers = [{ name = "Andrej Orsula", email = "orsula.andrej@gmail.com" }] +urls = { Repository = "https://github.com/AndrejOrsula/space_robotics_bench", Documentation = "https://AndrejOrsula.github.io/space_robotics_bench" } +classifiers = [ + "Programming Language :: Python", + "Programming Language :: Rust", + "License :: OSI Approved :: MIT License", + "License :: OSI Approved :: Apache Software License", + "Framework :: Robot Framework", +] +keywords = ["benchmark", "robotics", "simulation", "space"] +readme = "README.md" +license = { text = "MIT OR Apache-2.0" } +requires-python = ">=3.10" +dynamic = ["version"] +dependencies = ["platformdirs>=4.2,<5", "pydantic>=2.7,<3"] + +[project.optional-dependencies] +all = [ + ## Algorithms + "space_robotics_bench[dreamerv3]", + "space_robotics_bench[robomimic]", + "space_robotics_bench[sb3]", + ## Hardware + "space_robotics_bench[spacemouse]", + ## Utils + "space_robotics_bench[rich]", + "space_robotics_bench[tests]", +] +## Algorithms +dreamerv3 = [ + "dreamerv3@git+https://github.com/AndrejOrsula/dreamerv3.git@58bd8ff875077a859b68f9ad94d89129d4f8dd68", +] +robomimic = [ + "robomimic@git+https://github.com/ARISE-Initiative/robomimic.git@29d6ca229dec3327f87b54cf1688a94177f92af6", +] +sb3 = ["stable-baselines3>=2.3,<3", "sb3-contrib>=2.3,<3"] +## Hardware +spacemouse = ["pyspacemouse>=1.1,<2"] +## Utils +rich = ["rich>=13.9,<14"] +tests = ["pytest>=8.2,<9", "pytest-env>=1.1,<2"] diff --git a/scripts/_cli_utils.py b/scripts/_cli_utils.py new file mode 100644 index 0000000..914fc39 --- /dev/null +++ b/scripts/_cli_utils.py @@ -0,0 +1,163 @@ +import argparse +import subprocess +import sys +from os import environ, path + +from omni.isaac.lab.app import AppLauncher + + +def add_default_cli_args(parser: argparse.Namespace): + # Environment + class _AutoNamespaceTaskAction(argparse.Action): + def __call__(self, parser, namespace, values, option_string=None): + if "/" not in values: + DEFAULT_TASK_NAMESPACE: str = "srb" + values = f"{DEFAULT_TASK_NAMESPACE}/{values}" + setattr(namespace, self.dest, values) + + environment_group = parser.add_argument_group( + "environment arguments", + description="Arguments for environment.", + ) + environment_group.add_argument( + "--task", + "--demo", + "--env", + type=str, + default="srb/sample_collection", + action=_AutoNamespaceTaskAction, + help="Name of the task/demo/env. You can run the `list_envs.py` script to get a list of all registered tasks.", + ) + environment_group.add_argument( + "--seed", type=int, default=None, help="Seed used for the environment" + ) + environment_group.add_argument( + "--num_envs", + type=int, + default=1, + help="Number of environments to simulate in parallel.", + ) + + # Compute + compute_group = parser.add_argument_group( + "compute arguments", + description="Arguments for compute.", + ) + compute_group.add_argument( + "--disable_fabric", + action="store_true", + default=False, + help="Disable fabric and use USD I/O operations.", + ) + + # Video recording + video_recording_group = parser.add_argument_group( + "video_recording arguments", + description="Arguments for video_recording.", + ) + video_recording_group.add_argument( + "--video", + action="store_true", + default=False, + help="Record videos.", + ) + video_recording_group.add_argument( + "--video_length", + type=int, + default=1000, + help="Length of the recorded video (in steps).", + ) + video_recording_group.add_argument( + "--video_interval", + type=int, + default=10000, + help="Interval between video recordings (in steps).", + ) + + # Experience + experience_group = parser.add_argument_group( + "experience arguments", + description="Arguments for experience.", + ) + experience_group.add_argument( + "--disable_ui", + action="store_true", + default=False, + help="Disable most of the Isaac Sim UI and set it to fullscreen.", + ) + + # Append app launcher arguments + AppLauncher.add_app_launcher_args(parser) + + +def launch_app(args: argparse.Namespace) -> AppLauncher: + _update_extension_module() + _autoenable_cameras(args) + _autoselect_experience(args) + + launcher = AppLauncher(launcher_args=args) + + if args.disable_ui: + _disable_ui() + + return launcher + + +def shutdown_app(launcher: AppLauncher): + launcher.app.close() + + +def _update_extension_module(): + if environ.get("SRB_SKIP_EXT_MOD_UPDATE", "false").lower() not in ["true", "1"]: + print("Updating Rust extension module...") + result = subprocess.run( + [ + environ.get("ISAAC_SIM_PYTHON", "python3"), + "-m", + "pip", + "install", + "--no-input", + "--no-clean", + "--no-compile", + "--no-deps", + "--no-color", + "--disable-pip-version-check", + "--no-python-version-warning", + "--editable", + path.dirname(path.dirname(path.realpath(__file__))), + ], + capture_output=True, + text=True, + ) + if result.returncode != 0: + print(result.stderr, file=sys.stderr) + + +def _autoenable_cameras(args: argparse.Namespace): + if not args.enable_cameras and (args.video or "visual" in args.task): + args.enable_cameras = True + + +def _autoselect_experience(args: argparse.Namespace): + ## Get relative path to the experience + project_dir = path.dirname(path.dirname(path.realpath(__file__))) + experience_dir = path.join(project_dir, "apps") + + ## Select the experience based on args + experience = "srb" + if args.headless: + experience += ".headless" + if args.enable_cameras: + experience += ".rendering" + experience += ".kit" + + ## Set the experience + args.experience = path.join(experience_dir, experience) + + +def _disable_ui(): + import carb.settings + + settings = carb.settings.get_settings() + settings.set("/app/window/hideUi", True) + settings.set("/app/window/fullscreen", True) diff --git a/scripts/algo/dreamerv3/play.py b/scripts/algo/dreamerv3/play.py new file mode 100755 index 0000000..c87701d --- /dev/null +++ b/scripts/algo/dreamerv3/play.py @@ -0,0 +1,262 @@ +#!/root/isaac-sim/python.sh + +import os +import sys + +from omni.isaac.lab.app import AppLauncher + +sys.path.append( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))) +) +from _cli_utils import add_default_cli_args, argparse, launch_app, shutdown_app + +ALGO_NAME = "dreamerv3" +ALGO_CFG_ENTRYPOINT_KEY = f"{ALGO_NAME}_cfg" + + +def main(launcher: AppLauncher, args: argparse.Namespace): + import re + from collections import defaultdict + + import dreamerv3 + import embodied + import gymnasium + import numpy as np + from omni.isaac.kit import SimulationApp + from omni.isaac.lab.utils.dict import print_dict + from omni.isaac.lab.utils.io import dump_pickle, dump_yaml + + import space_robotics_bench # noqa: F401 + from space_robotics_bench.core.sim import SimulationContext + from space_robotics_bench.core.wrappers.dreamerv3 import ( + DriverParallelEnv, + EmbodiedEnvWrapper, + process_dreamerv3_cfg, + ) + from space_robotics_bench.utils.parsing import ( + create_eval_logdir_path, + get_checkpoint_path, + load_cfg_from_registry, + parse_task_cfg, + ) + + ## Extract simulation app + _sim_app: SimulationApp = launcher.app + + ## Configuration + task_cfg = parse_task_cfg( + task_name=args.task, + device=args.device, + num_envs=args.num_envs, + use_fabric=not args.disable_fabric, + ) + agent_cfg = load_cfg_from_registry(args.task, ALGO_CFG_ENTRYPOINT_KEY) + + # Allow overriding the seed + if args.seed is not None: + agent_cfg["seed"] = args.seed + + # Get path to checkpoint + checkpoint_path = get_checkpoint_path( + ALGO_NAME, + args.task, + checkpoint=args.checkpoint, + ) + + # Directory for logging + logdir = create_eval_logdir_path(checkpoint_path) + + dump_yaml(os.path.join(logdir, "params", "task.yaml"), task_cfg) + dump_pickle(os.path.join(logdir, "params", "task.pkl"), task_cfg) + + dump_yaml(os.path.join(logdir, "params", "agent.yaml"), agent_cfg) + dump_pickle(os.path.join(logdir, "params", "agent.pkl"), agent_cfg) + + # Post-process agent configuration + agent_cfg = process_dreamerv3_cfg( + agent_cfg, + logdir=logdir, + num_envs=args.num_envs, + task_name=args.task, + model_size=args.model_size, + ) + + def make_env(config): + ## Create the environment + env = gymnasium.make(id=args.task, cfg=task_cfg, render_mode="rgb_array") + + ## Initialize the environment + _observation, _info = env.reset() + # Render a couple of frames to help with stabilization of rendered images + sim_context: SimulationContext = SimulationContext.instance() + for _ in range(64): + sim_context.render() + + ## Add wrapper for video recording (if enabled) + if args.video: + video_kwargs = { + "video_folder": os.path.join(logdir, "videos"), + "step_trigger": lambda step: step % args.video_interval == 0, + "video_length": args.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gymnasium.wrappers.RecordVideo(env, **video_kwargs) + + ## Add wrapper for DreamerV3 + env = EmbodiedEnvWrapper(env) + + ## Seed the environment + env.seed(seed=config["seed"] if config["seed"] is not None else 42) + + return env + + def make_agent(config, env): + if config.random_agent: + agent = embodied.RandomAgent(env.obs_space, env.act_space) + else: + agent = dreamerv3.Agent(env.obs_space, env.act_space, config) + return agent + + def make_logger(config): + logdir = embodied.Path(config.logdir) + multiplier = config.env.get(config.task.split("_")[0], {}).get("repeat", 1) + return embodied.Logger( + embodied.Counter(), + [ + embodied.logger.TerminalOutput(config.filter, "Agent"), + embodied.logger.JSONLOutput(logdir, "metrics.jsonl"), + embodied.logger.JSONLOutput(logdir, "scores.jsonl", "episode/score"), + embodied.logger.TensorBoardOutput( + logdir, + config.run.log_video_fps, + config.tensorboard_videos, + parallel=False, + ), + # embodied.logger.WandbOutput(logdir.name, config=config), + ], + multiplier, + ) + + run_args = embodied.Config( + **agent_cfg.run, + logdir=logdir, + batch_length_eval=agent_cfg.batch_length_eval, + batch_length=agent_cfg.batch_length, + batch_size=agent_cfg.batch_size, + replay_context=agent_cfg.replay_context, + ) + + # Note: Everything below is a modified replacement for `embodied.run.eval_only` + # - 1 env is created + # - DriverParallelEnv is used to run the agent in parallel + + env = make_env(agent_cfg) + agent = make_agent(agent_cfg, env) + logger = make_logger(agent_cfg) + + logdir = embodied.Path(run_args.logdir) + logdir.mkdir() + print("Logdir", logdir) + step = logger.step + usage = embodied.Usage(**run_args.usage) + agg = embodied.Agg() + epstats = embodied.Agg() + episodes = defaultdict(embodied.Agg) + should_log = embodied.when.Clock(run_args.log_every) + policy_fps = embodied.FPS() + + @embodied.timer.section("log_step") + def log_step(tran, worker): + episode = episodes[worker] + episode.add("score", tran["reward"], agg="sum") + episode.add("length", 1, agg="sum") + episode.add("rewards", tran["reward"], agg="stack") + + if tran["is_first"]: + episode.reset() + + if worker < run_args.log_video_streams: + for key in run_args.log_keys_video: + if key in tran: + episode.add(f"policy_{key}", tran[key], agg="stack") + for key, value in tran.items(): + if re.match(run_args.log_keys_sum, key): + episode.add(key, value, agg="sum") + if re.match(run_args.log_keys_avg, key): + episode.add(key, value, agg="avg") + if re.match(run_args.log_keys_max, key): + episode.add(key, value, agg="max") + + if tran["is_last"]: + result = episode.result() + logger.add( + { + "score": result.pop("score"), + "length": result.pop("length") - 1, + }, + prefix="episode", + ) + rew = result.pop("rewards") + if len(rew) > 1: + result["reward_rate"] = (np.abs(rew[1:] - rew[:-1]) >= 0.01).mean() + epstats.add(result) + + driver = DriverParallelEnv( + env, + run_args.num_envs, + ) + driver.on_step(lambda tran, _: step.increment()) + driver.on_step(lambda tran, _: policy_fps.step()) + driver.on_step(log_step) + + checkpoint = embodied.Checkpoint() + checkpoint.agent = agent + checkpoint.load(checkpoint_path, keys=["agent"]) + + print("Start evaluation") + policy = lambda *run_args: agent.policy(*run_args, mode="eval") + driver.reset(agent.init_policy) + while step < run_args.steps: + driver(policy, steps=10) + if should_log(step): + logger.add(agg.result()) + logger.add(epstats.result(), prefix="epstats") + logger.add(embodied.timer.stats(), prefix="timer") + logger.add(usage.stats(), prefix="usage") + logger.add({"fps/policy": policy_fps.result()}) + logger.write() + + logger.close() + + +### Helper functions ### +def parse_cli_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + algorithm_group = parser.add_argument_group( + "algorithm arguments", + description="Arguments for algorithm.", + ) + algorithm_group.add_argument( + "--model_size", + type=str, + default="debug", + help="Size of the model to train\n(debug, size12m, size25m, size50m, size100m, size200m, size400m)", + ) + add_default_cli_args(parser) + return parser.parse_args() + + +if __name__ == "__main__": + # Parse arguments + args = parse_cli_args() + + # Launch the app + launcher = launch_app(args) + + # Run the main function + main(launcher=launcher, args=args) + + # Shutdown the app + shutdown_app(launcher) diff --git a/scripts/algo/dreamerv3/train.py b/scripts/algo/dreamerv3/train.py new file mode 100755 index 0000000..a08286b --- /dev/null +++ b/scripts/algo/dreamerv3/train.py @@ -0,0 +1,385 @@ +#!/root/isaac-sim/python.sh + +import os +import sys + +from omni.isaac.lab.app import AppLauncher + +sys.path.append( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))) +) +from _cli_utils import add_default_cli_args, argparse, launch_app, shutdown_app + +ALGO_NAME = "dreamerv3" +ALGO_CFG_ENTRYPOINT_KEY = f"{ALGO_NAME}_cfg" + + +def main(launcher: AppLauncher, args: argparse.Namespace): + import re + from collections import defaultdict + from functools import partial + + import dreamerv3 + import embodied + import gymnasium + import numpy as np + from omni.isaac.kit import SimulationApp + from omni.isaac.lab.utils.dict import print_dict + from omni.isaac.lab.utils.io import dump_pickle, dump_yaml + + import space_robotics_bench # noqa: F401 + from space_robotics_bench.core.sim import SimulationContext + from space_robotics_bench.core.wrappers.dreamerv3 import ( + DriverParallelEnv, + EmbodiedEnvWrapper, + process_dreamerv3_cfg, + ) + from space_robotics_bench.utils.parsing import ( + create_logdir_path, + get_last_run_logdir_path, + load_cfg_from_registry, + parse_task_cfg, + ) + + ## Extract simulation app + _sim_app: SimulationApp = launcher.app + + ## Configuration + task_cfg = parse_task_cfg( + task_name=args.task, + device=args.device, + num_envs=args.num_envs, + use_fabric=not args.disable_fabric, + ) + agent_cfg = load_cfg_from_registry(args.task, ALGO_CFG_ENTRYPOINT_KEY) + + # Allow overriding the seed + if args.seed is not None: + agent_cfg["seed"] = args.seed + + # Directory for logging + if args.continue_training: + logdir = get_last_run_logdir_path(ALGO_NAME, args.task) + else: + logdir = create_logdir_path(ALGO_NAME, args.task) + + # Write configuration to logdir (pt. 1) + dump_yaml(os.path.join(logdir, "params", "task.yaml"), task_cfg) + dump_yaml(os.path.join(logdir, "params", "agent.yaml"), agent_cfg) + dump_pickle(os.path.join(logdir, "params", "agent.pkl"), agent_cfg) + + # Post-process agent configuration + agent_cfg = process_dreamerv3_cfg( + agent_cfg, + logdir=logdir, + num_envs=args.num_envs, + task_name=args.task, + model_size=args.model_size, + ) + + # Write configuration to logdir (pt. 2) + agent_cfg.save(os.path.join(logdir, "config.yaml")) + print_dict(agent_cfg, nesting=4) + + def make_env(config): + ## Create the environment + env = gymnasium.make(id=args.task, cfg=task_cfg, render_mode="rgb_array") + + ## Initialize the environment + _observation, _info = env.reset() + # Render a couple of frames to help with stabilization of rendered images + sim_context: SimulationContext = SimulationContext.instance() + for _ in range(64): + sim_context.render() + + ## Add wrapper for video recording (if enabled) + if args.video: + video_kwargs = { + "video_folder": os.path.join(logdir, "videos"), + "step_trigger": lambda step: step % args.video_interval == 0, + "video_length": args.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gymnasium.wrappers.RecordVideo(env, **video_kwargs) + + ## Add wrapper for DreamerV3 + env = EmbodiedEnvWrapper(env) + + ## Seed the environment + env.seed(seed=config["seed"] if config["seed"] is not None else 42) + + return env + + def make_agent(config, env): + if config.random_agent: + agent = embodied.RandomAgent(env.obs_space, env.act_space) + else: + agent = dreamerv3.Agent(env.obs_space, env.act_space, config) + return agent + + def make_logger(config): + logdir = embodied.Path(config.logdir) + multiplier = config.env.get(config.task.split("_")[0], {}).get("repeat", 1) + return embodied.Logger( + embodied.Counter(), + [ + embodied.logger.TerminalOutput(config.filter, "Agent"), + embodied.logger.JSONLOutput(logdir, "metrics.jsonl"), + embodied.logger.JSONLOutput(logdir, "scores.jsonl", "episode/score"), + embodied.logger.TensorBoardOutput( + logdir, + config.run.log_video_fps, + config.tensorboard_videos, + parallel=False, + ), + # embodied.logger.WandbOutput(logdir.name, config=config), + ], + multiplier, + ) + + def make_replay(config, directory=None, is_eval=False, rate_limit=False): + directory = directory and embodied.Path(config.logdir) / directory + size = int(config.replay.size / 10 if is_eval else config.replay.size) + length = config.replay_length_eval if is_eval else config.replay_length + kwargs = {} + # if rate_limit and config.run.train_ratio > 0: + # kwargs["samples_per_insert"] = config.run.train_ratio / ( + # length - config.replay_context + # ) + # kwargs["tolerance"] = 5 * config.batch_size + # kwargs["min_size"] = min( + # max(config.batch_size, config.run.train_fill), size + # ) + if config.replay.fracs.uniform < 1.0 and not is_eval: + assert config.jax.compute_dtype in ("bfloat16", "float32"), ( + "Gradient scaling for low-precision training can produce invalid loss " + "outputs that are incompatible with prioritized replay." + ) + # import numpy as np + + kwargs["selector"] = embodied.replay.selectors.Mixture( + selectors={ + "uniform": embodied.replay.selectors.Uniform( + seed=agent_cfg["seed"] + ), + "priority": embodied.replay.selectors.Prioritized( + seed=agent_cfg["seed"], **config.replay.prio + ), + # "recency": embodied.replay.selectors.Recency( + # uprobs=1.0 / np.arange(1, size + 1) ** config.replay.recexp, + # seed=agent_cfg["seed"], + # ), + }, + fractions={ + "uniform": config.replay.fracs.uniform, + "priority": config.replay.fracs.priority, + }, + seed=agent_cfg["seed"], + ) + replay = embodied.replay.Replay( + length=length, + capacity=size, + directory=directory, + chunksize=config.replay.chunksize, + online=config.replay.online, + **kwargs, + ) + return replay + + run_args = embodied.Config( + **agent_cfg.run, + logdir=logdir, + batch_length_eval=agent_cfg.batch_length_eval, + batch_length=agent_cfg.batch_length, + batch_size=agent_cfg.batch_size, + replay_context=agent_cfg.replay_context, + replay_length_eval=agent_cfg.replay_length_eval, + replay_length=agent_cfg.replay_length, + ) + + # Note: Everything below is a modified replacement for `embodied.run.train` + # - 1 env is created + # - DriverParallelEnv is used to run the agent in parallel + ## Train the agent + # embodied.run.train( + # partial(make_agent, agent_cfg), + # partial(make_replay, agent_cfg), + # partial(make_env, agent_cfg), + # partial(make_logger, agent_cfg), + # run_args, + # ) + + env = make_env(agent_cfg) + agent = make_agent(agent_cfg, env) + replay = make_replay(agent_cfg) + logger = make_logger(agent_cfg) + + # if not args.continue_training: + # dump_yaml(os.path.join(logdir, "params", "env.yaml"), env.unwrapped.cfg.env_cfg) + + logdir = embodied.Path(run_args.logdir) + logdir.mkdir() + print("[INFO] Logdir", logdir) + step = logger.step + usage = embodied.Usage(**run_args.usage) + agg = embodied.Agg() + epstats = embodied.Agg() + episodes = defaultdict(embodied.Agg) + policy_fps = embodied.FPS() + train_fps = embodied.FPS() + + batch_steps = run_args.batch_size * ( + run_args.batch_length - run_args.replay_context + ) + # should_expl = embodied.when.Until(run_args.expl_until) + should_train = embodied.when.Ratio(run_args.train_ratio / batch_steps) + should_log = embodied.when.Clock(run_args.log_every) + should_eval = embodied.when.Clock(run_args.eval_every) + should_save = embodied.when.Clock(run_args.save_every) + + @embodied.timer.section("log_step") + def log_step(tran, worker): + episode = episodes[worker] + episode.add("score", tran["reward"], agg="sum") + episode.add("length", 1, agg="sum") + episode.add("rewards", tran["reward"], agg="stack") + + if tran["is_first"]: + episode.reset() + + if worker < run_args.log_video_streams: + for key in run_args.log_keys_video: + if key in tran: + episode.add(f"policy_{key}", tran[key], agg="stack") + for key, value in tran.items(): + if re.match(run_args.log_keys_sum, key): + episode.add(key, value, agg="sum") + if re.match(run_args.log_keys_avg, key): + episode.add(key, value, agg="avg") + if re.match(run_args.log_keys_max, key): + episode.add(key, value, agg="max") + + if tran["is_last"]: + result = episode.result() + logger.add( + { + "score": result.pop("score"), + "length": result.pop("length"), + }, + prefix="episode", + ) + rew = result.pop("rewards") + if len(rew) > 1: + result["reward_rate"] = (np.abs(rew[1:] - rew[:-1]) >= 0.01).mean() + epstats.add(result) + + driver = DriverParallelEnv( + env, + args.num_envs, + ) + driver.on_step(lambda tran, _: step.increment()) + driver.on_step(lambda tran, _: policy_fps.step()) + driver.on_step(replay.add) + driver.on_step(log_step) + + dataset_train = iter( + agent.dataset( + partial(replay.dataset, run_args.batch_size, run_args.batch_length) + ) + ) + dataset_report = iter( + agent.dataset( + partial(replay.dataset, run_args.batch_size, run_args.batch_length_eval) + ) + ) + carry = [agent.init_train(run_args.batch_size)] + carry_report = agent.init_report(run_args.batch_size) + + def train_step(tran, worker): + if len(replay) < run_args.batch_size or step < run_args.train_fill: + return + for _ in range(should_train(step)): + with embodied.timer.section("dataset_next"): + batch = next(dataset_train) + outs, carry[0], mets = agent.train(batch, carry[0]) + train_fps.step(batch_steps) + if "replay" in outs: + replay.update(outs["replay"]) + agg.add(mets, prefix="train") + + driver.on_step(train_step) + + checkpoint = embodied.Checkpoint(logdir / "checkpoint.ckpt") + checkpoint.step = step + checkpoint.agent = agent + checkpoint.replay = replay + if run_args.from_checkpoint: + checkpoint.load(run_args.from_checkpoint) + checkpoint.load_or_save() + should_save(step) + + # policy = lambda *run_args: agent.policy( + # *run_args, mode="explore" if should_expl(step) else "train" + # ) + policy = lambda *run_args: agent.policy(*run_args, mode="train") + driver.reset(agent.init_policy) + while step < run_args.steps: + driver(policy, steps=10) + + if should_eval(step) and len(replay): + mets, _ = agent.report(next(dataset_report), carry_report) + logger.add(mets, prefix="report") + + if should_log(step): + logger.add(agg.result()) + logger.add(epstats.result(), prefix="epstats") + logger.add(embodied.timer.stats(), prefix="timer") + logger.add(replay.stats(), prefix="replay") + logger.add(usage.stats(), prefix="usage") + logger.add({"fps/policy": policy_fps.result()}) + logger.add({"fps/train": train_fps.result()}) + logger.write() + + if should_save(step): + checkpoint.save() + + logger.close() + + +### Helper functions ### +def parse_cli_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + algorithm_group = parser.add_argument_group( + "algorithm arguments", + description="Arguments for algorithm.", + ) + algorithm_group.add_argument( + "--model_size", + type=str, + default="debug", + help="Size of the model to train\n(debug, size12m, size25m, size50m, size100m, size200m, size400m)", + ) + algorithm_group.add_argument( + "--continue_training", + action="store_true", + default=False, + help="Continue training the model from the checkpoint of the last run.", + ) + add_default_cli_args(parser) + return parser.parse_args() + + +if __name__ == "__main__": + # Parse arguments + args = parse_cli_args() + + # Launch the app + launcher = launch_app(args) + + # Run the main function + main(launcher=launcher, args=args) + + # Shutdown the app + shutdown_app(launcher) diff --git a/scripts/algo/robomimic/collect.py b/scripts/algo/robomimic/collect.py new file mode 100755 index 0000000..5c8a4e3 --- /dev/null +++ b/scripts/algo/robomimic/collect.py @@ -0,0 +1,250 @@ +#!/root/isaac-sim/python.sh + +import os +import sys + +from omni.isaac.lab.app import AppLauncher + +sys.path.append( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))) +) +from _cli_utils import add_default_cli_args, argparse, launch_app, shutdown_app + +FRAMEWORK_NAME = "robomimic" + + +def main(launcher: AppLauncher, args: argparse.Namespace): + ## Note: Importing modules here due to delayed Omniverse Kit extension loading + import contextlib + from os import path + + import gymnasium + import numpy as np + import torch + from omni.isaac.kit import SimulationApp + from omni.isaac.lab.utils.dict import print_dict + from omni.isaac.lab.utils.io import dump_pickle, dump_yaml + + import space_robotics_bench # Noqa: F401 + from space_robotics_bench.core import mdp + from space_robotics_bench.core.actions import ( + ManipulatorTaskSpaceActionCfg, + MultiCopterActionGroupCfg, + WheeledRoverActionGroupCfg, + ) + from space_robotics_bench.core.managers import SceneEntityCfg + from space_robotics_bench.core.teleop_devices import CombinedInterface + from space_robotics_bench.core.wrappers.robomimic import RobomimicDataCollector + from space_robotics_bench.utils.parsing import create_logdir_path, parse_task_cfg + + if args.headless and "keyboard" in args.teleop_device: + raise ValueError("Native teleoperation is only supported in GUI mode.") + + ## Extract simulation app + sim_app: SimulationApp = launcher.app + + # Parse configuration + task_cfg = parse_task_cfg( + task_name=args.task, + device=args.device, + num_envs=args.num_envs, + use_fabric=not args.disable_fabric, + ) + # Disable truncation + if hasattr(task_cfg, "enable_truncation"): + task_cfg.enable_truncation = False + + ## Create the environment + env = gymnasium.make( + id=args.task, cfg=task_cfg, render_mode="rgb_array" if args.video else None + ) + + ## Create controller + teleop_interface = CombinedInterface( + devices=args.teleop_device, + pos_sensitivity=args.pos_sensitivity, + rot_sensitivity=args.rot_sensitivity, + action_cfg=env.unwrapped.cfg.actions, + ) + teleop_interface.reset() + teleop_interface.add_callback("L", env.reset) + print(teleop_interface) + + ## Initialize the environment + observation, info = env.reset() + + ## Add wrapper for video recording (if enabled) + if args.video: + logdir = create_logdir_path("srb", args.task) + video_kwargs = { + "video_folder": path.join(logdir, "videos"), + "step_trigger": lambda step: step % args.video_interval == 0, + "video_length": args.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gymnasium.wrappers.RecordVideo(env, **video_kwargs) + + def process_actions(twist: np.ndarray, gripper_cmd: bool) -> torch.Tensor: + twist = torch.tensor( + twist, dtype=torch.float32, device=env.unwrapped.device + ).repeat(env.unwrapped.num_envs, 1) + if isinstance(env.unwrapped.cfg.actions, ManipulatorTaskSpaceActionCfg): + if not args.disable_control_scheme_inversion: + twist[:, :2] *= -1.0 + gripper_action = torch.zeros(twist.shape[0], 1, device=twist.device) + gripper_action[:] = -1.0 if gripper_cmd else 1.0 + return torch.concat([twist, gripper_action], dim=1) + elif isinstance(env.unwrapped.cfg.actions, MultiCopterActionGroupCfg): + return torch.concat( + [ + twist[:, :3], + twist[:, 5].unsqueeze(1), + ], + dim=1, + ) + + elif isinstance(env.unwrapped.cfg.actions, WheeledRoverActionGroupCfg): + return twist[:, :2] + + # Specify directory for logging experiments + logdir = create_logdir_path(FRAMEWORK_NAME, f"{args.task}/dataset/") + + # Dump the configuration into log-directory + dump_yaml(os.path.join(logdir, "params", "task.yaml"), task_cfg) + dump_pickle(os.path.join(logdir, "params", "task.pkl"), task_cfg) + + # Create data-collector + collector_interface = RobomimicDataCollector( + env_name=args.task, + directory_path=logdir, + filename=args.filename, + num_demos=args.num_demos, + flush_freq=env.num_envs, + ) + collector_interface.reset() + + ## Run the environment + with contextlib.suppress(KeyboardInterrupt) and torch.inference_mode(): + while sim_app.is_running() and not collector_interface.is_stopped(): + # Get actions from the teleoperation interface + actions = process_actions(*teleop_interface.advance()) + + # Store actions and observations before stepping the environment + collector_interface.add("actions", actions) + for key, value in observation.items(): + if key == "policy": + for inner_key, inner_value in value.items(): + collector_interface.add(f"obs/{inner_key}", inner_value) + else: + collector_interface.add(f"obs/{key}", value) + + # Step the environment + observation, reward, terminated, truncated, info = env.step(actions) + dones = terminated | truncated + + # Note: Each environment is automatically reset (independently) when terminated or truncated + + # Store observations, rewards and dones after stepping the environment + for key, value in observation.items(): + if key == "policy": + for inner_key, inner_value in value.items(): + collector_interface.add(f"next_obs/{inner_key}", inner_value) + else: + collector_interface.add(f"next_obs/{key}", value) + collector_interface.add("rewards", reward) + collector_interface.add("dones", dones) + collector_interface.add("success", torch.zeros_like(dones)) + + # Flush data from collector for successful environments + reset_env_ids = dones.nonzero(as_tuple=False).squeeze(-1) + collector_interface.flush(reset_env_ids) + + # Provide force feedback for teleop devices + if isinstance(env.unwrapped.cfg.actions, ManipulatorTaskSpaceActionCfg): + FT_FEEDBACK_SCALE = torch.tensor([0.16, 0.16, 0.16, 0.0, 0.0, 0.0]) + ft_feedback_asset_cfg = SceneEntityCfg( + "robot", + body_names=task_cfg.robot_cfg.regex_links_hand, + ) + ft_feedback_asset_cfg.resolve(env.unwrapped.scene) + ft_feedback = ( + FT_FEEDBACK_SCALE + * mdp.body_incoming_wrench_mean( + env=env.unwrapped, + asset_cfg=ft_feedback_asset_cfg, + )[0, ...].cpu() + ) + teleop_interface.set_ft_feedback(ft_feedback) + + # Check if the data collection is stopped + if collector_interface.is_stopped(): + break + + # Close the simulator + collector_interface.close() + env.close() + + +### Helper functions ### +def parse_cli_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + algorithm_group = parser.add_argument_group( + "teleop", + description="Arguments for teleoperation", + ) + algorithm_group.add_argument( + "--teleop_device", + type=str, + nargs="+", + default=["keyboard"], + help="Device for interacting with environment", + ) + algorithm_group.add_argument( + "--pos_sensitivity", + type=float, + default=10.0, + help="Sensitivity factor for translation.", + ) + algorithm_group.add_argument( + "--rot_sensitivity", + type=float, + default=40.0, + help="Sensitivity factor for rotation.", + ) + algorithm_group.add_argument( + "--disable_control_scheme_inversion", + action="store_true", + default=False, + help="Flag to disable inverting the control scheme due to view for manipulation-based tasks.", + ) + algorithm_group = parser.add_argument_group( + "data_collection", + description="Arguments for data collection", + ) + algorithm_group.add_argument( + "--num_demos", + type=int, + default=1000, + help="Number of episodes to store in the dataset.", + ) + algorithm_group.add_argument( + "--filename", type=str, default="hdf_dataset", help="Basename of output file." + ) + add_default_cli_args(parser) + return parser.parse_args() + + +if __name__ == "__main__": + # Parse arguments + args = parse_cli_args() + + # Launch the app + launcher = launch_app(args) + + # Run the main function + main(launcher=launcher, args=args) + + # Shutdown the app + shutdown_app(launcher) diff --git a/scripts/algo/robomimic/play.py b/scripts/algo/robomimic/play.py new file mode 100755 index 0000000..19ec498 --- /dev/null +++ b/scripts/algo/robomimic/play.py @@ -0,0 +1,88 @@ +#!/root/isaac-sim/python.sh + + +import os +import sys + +from omni.isaac.lab.app import AppLauncher + +sys.path.append( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))) +) +from _cli_utils import add_default_cli_args, argparse, launch_app, shutdown_app + + +def main(launcher: AppLauncher, args: argparse.Namespace): + import gymnasium + import robomimic.utils.fileutils as FileUtils + import robomimic.utils.torchutils as TorchUtils + import torch + from omni.isaac.kit import SimulationApp + + import space_robotics_bench # noqa: F401 + from space_robotics_bench.utils.parsing import parse_task_cfg + + ## Extract simulation app + sim_app: SimulationApp = launcher.app + + # parse configuration + task_cfg = parse_task_cfg( + args.task, device=args.device, num_envs=1, use_fabric=not args.disable_fabric + ) + # we want to have the terms in the observations returned as a dictionary + # rather than a concatenated tensor + task_cfg.observations.policy.concatenate_terms = False + + # create environment + env = gymnasium.make(args.task, cfg=task_cfg) + + # acquire device + device = TorchUtils.get_torch_device(try_to_use_cuda=True) + # restore policy + policy, _ = FileUtils.policy_from_checkpoint( + ckpt_path=args.checkpoint, device=device, verbose=True + ) + + # reset environment + obs_dict, _ = env.reset() + # robomimic only cares about policy observations + obs = obs_dict["policy"] + # simulate environment + while sim_app.is_running(): + # run everything in inference mode + with torch.inference_mode(): + # compute actions + actions = policy(obs) + actions = ( + torch.from_numpy(actions) + .to(device=device) + .view(1, env.action_space.shape[1]) + ) + # apply actions + obs_dict = env.step(actions)[0] + # robomimic only cares about policy observations + obs = obs_dict["policy"] + + # close the simulator + env.close() + + +### Helper functions ### +def parse_cli_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + add_default_cli_args(parser) + return parser.parse_args() + + +if __name__ == "__main__": + # Parse arguments + args = parse_cli_args() + + # Launch the app + launcher = launch_app(args) + + # Run the main function + main(launcher=launcher, args=args) + + # Shutdown the app + shutdown_app(launcher) diff --git a/scripts/algo/robomimic/tools/create_default_cfg.py b/scripts/algo/robomimic/tools/create_default_cfg.py new file mode 100755 index 0000000..4bfb3ca --- /dev/null +++ b/scripts/algo/robomimic/tools/create_default_cfg.py @@ -0,0 +1,28 @@ +#!/root/isaac-sim/python.sh + +import argparse +import json + +from robomimic.config import config_factory + + +def main(args: argparse.Namespace): + cfg = config_factory(args.algo) + with open(f"{args.algo}.json", "w") as f: + json.dump(cfg, f, indent=4) + + +def parse_cli_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--algo", + type=str, + default="", + help="Name of the algorithm", + ) + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_cli_args() + main(args=args) diff --git a/scripts/algo/robomimic/train.py b/scripts/algo/robomimic/train.py new file mode 100755 index 0000000..f1ac666 --- /dev/null +++ b/scripts/algo/robomimic/train.py @@ -0,0 +1,448 @@ +#!/root/isaac-sim/python.sh + +import os +import sys + +from omni.isaac.lab.app import AppLauncher + +sys.path.append( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))) +) +from _cli_utils import add_default_cli_args, argparse, launch_app, shutdown_app + +FRAMEWORK_NAME = "robomimic" +ALGO_CFG_ENTRYPOINT_KEY = f"{FRAMEWORK_NAME}_{{ALGO_NAME}}_cfg" + + +def main(launcher: AppLauncher, args: argparse.Namespace): + import json + import time + import traceback + from collections import OrderedDict + + import gymnasium + import numpy as np + import psutil + import robomimic.utils.envutils as EnvUtils + import robomimic.utils.fileutils as FileUtils + import robomimic.utils.obsutils as ObsUtils + import robomimic.utils.torchutils as TorchUtils + import robomimic.utils.trainutils as TrainUtils + import torch + from omni.isaac.kit import SimulationApp + from robomimic.algo import RolloutPolicy, algo_factory + from robomimic.config import config_factory + from robomimic.utils.logutils import DataLogger, PrintLogger + from torch.utils.data import DataLoader + + import space_robotics_bench # noqa: F401 + + ## Extract simulation app + _sim_app: SimulationApp = launcher.app + + # load config + if args.task is not None: + # obtain the configuration entry point + cfg_entry_point_key = ALGO_CFG_ENTRYPOINT_KEY.format(ALGO_NAME=args.algo) + + print(f"Loading configuration for task: {args.task}") + cfg_entry_point_file = gymnasium.spec(args.task).kwargs.pop(cfg_entry_point_key) + # check if entry point exists + if cfg_entry_point_file is None: + raise ValueError( + f"Could not find configuration for the environment: '{args.task}'." + f" Please check that the gym registry has the entry point: '{cfg_entry_point_key}'." + ) + # load config from json file + with open(cfg_entry_point_file) as f: + ext_cfg = json.load(f) + config = config_factory(ext_cfg["algo_name"]) + # update config with external json - this will throw errors if + # the external config has keys not present in the base algo config + with config.values_unlocked(): + config.update(ext_cfg) + else: + raise ValueError("Please provide a task name through CLI arguments.") + + if args.dataset is not None: + config.train.data = args.dataset + else: + config.train.data = os.path.join( + f"./logs/{FRAMEWORK_NAME}", args.task, "hdf_dataset.hdf5" + ) + + if args.name is not None: + config.experiment.name = args.name + + # change location of experiment directory + config.train.output_dir = os.path.abspath( + os.path.join(f"./logs/{FRAMEWORK_NAME}", args.task) + ) + # get torch device + device = TorchUtils.get_torch_device(try_to_use_cuda=config.train.cuda) + + config.lock() + + # catch error during training and print it + res_str = "finished run successfully!" + try: + # first set seeds + np.random.seed(config.train.seed) + torch.manual_seed(config.train.seed) + + print("\n============= New Training Run with Config =============") + print(config) + print("") + log_dir, ckpt_dir, video_dir = TrainUtils.get_exp_dir(config) + print(f">>> Saving logs into directory: {log_dir}") + print(f">>> Saving checkpoints into directory: {ckpt_dir}") + print(f">>> Saving videos into directory: {video_dir}") + + if config.experiment.logging.terminal_output_to_txt: + # log stdout and stderr to a text file + logger = PrintLogger(os.path.join(log_dir, "log.txt")) + sys.stdout = logger + sys.stderr = logger + + # read config to set up metadata for observation modalities (e.g. detecting rgb observations) + ObsUtils.initialize_obsutils_with_config(config) + + # make sure the dataset exists + dataset_path = os.path.expanduser(config.train.data) + if not os.path.exists(dataset_path): + raise FileNotFoundError( + f"Dataset at provided path {dataset_path} not found!" + ) + + # load basic metadata from training file + print("\n============= Loaded Environment Metadata =============") + env_meta = FileUtils.get_env_metadata_from_dataset( + dataset_path=config.train.data + ) + shape_meta = FileUtils.get_shape_metadata_from_dataset( + dataset_path=config.train.data, + all_obs_keys=config.all_obs_keys, + verbose=True, + ) + + if config.experiment.env is not None: + env_meta["env_name"] = config.experiment.env + print( + "=" * 30 + + "\n" + + "Replacing Env to {}\n".format(env_meta["env_name"]) + + "=" * 30 + ) + + # create environment + envs = OrderedDict() + if config.experiment.rollout.enabled: + # create environments for validation runs + env_names = [env_meta["env_name"]] + + if config.experiment.additional_envs is not None: + for name in config.experiment.additional_envs: + env_names.append(name) + + for env_name in env_names: + env = EnvUtils.create_env_from_metadata( + env_meta=env_meta, + env_name=env_name, + render=False, + render_offscreen=config.experiment.render_video, + use_image_obs=shape_meta["use_images"], + ) + envs[env.name] = env + print(envs[env.name]) + + print("") + + # setup for a new training run + data_logger = DataLogger( + log_dir, config=config, log_tb=config.experiment.logging.log_tb + ) + model = algo_factory( + algo_name=config.algo_name, + config=config, + obs_key_shapes=shape_meta["all_shapes"], + ac_dim=shape_meta["ac_dim"], + device=device, + ) + + # save the config as a json file + with open(os.path.join(log_dir, "..", "config.json"), "w") as outfile: + json.dump(config, outfile, indent=4) + + print("\n============= Model Summary =============") + print(model) # print model summary + print("") + + # load training data + trainset, validset = TrainUtils.load_data_for_training( + config, obs_keys=shape_meta["all_obs_keys"] + ) + train_sampler = trainset.get_dataset_sampler() + print("\n============= Training Dataset =============") + print(trainset) + print("") + + # maybe retrieve statistics for normalizing observations + obs_normalization_stats = None + if config.train.hdf5_normalize_obs: + obs_normalization_stats = trainset.get_obs_normalization_stats() + + # initialize data loaders + train_loader = DataLoader( + dataset=trainset, + sampler=train_sampler, + batch_size=config.train.batch_size, + shuffle=(train_sampler is None), + num_workers=config.train.num_data_workers, + drop_last=True, + ) + + if config.experiment.validate: + # cap num workers for validation dataset at 1 + num_workers = min(config.train.num_data_workers, 1) + valid_sampler = validset.get_dataset_sampler() + valid_loader = DataLoader( + dataset=validset, + sampler=valid_sampler, + batch_size=config.train.batch_size, + shuffle=(valid_sampler is None), + num_workers=num_workers, + drop_last=True, + ) + else: + valid_loader = None + + # main training loop + best_valid_loss = None + best_return = ( + {k: -np.inf for k in envs} if config.experiment.rollout.enabled else None + ) + best_success_rate = ( + {k: -1.0 for k in envs} if config.experiment.rollout.enabled else None + ) + last_ckpt_time = time.time() + + # number of learning steps per epoch (defaults to a full dataset pass) + train_num_steps = config.experiment.epoch_every_n_steps + valid_num_steps = config.experiment.validation_epoch_every_n_steps + + for epoch in range(1, config.train.num_epochs + 1): # epoch numbers start at 1 + step_log = TrainUtils.run_epoch( + model=model, + data_loader=train_loader, + epoch=epoch, + num_steps=train_num_steps, + ) + model.on_epoch_end(epoch) + + # setup checkpoint path + epoch_ckpt_name = f"model_epoch_{epoch}" + + # check for recurring checkpoint saving conditions + should_save_ckpt = False + if config.experiment.save.enabled: + time_check = (config.experiment.save.every_n_seconds is not None) and ( + time.time() - last_ckpt_time + > config.experiment.save.every_n_seconds + ) + epoch_check = ( + (config.experiment.save.every_n_epochs is not None) + and (epoch > 0) + and (epoch % config.experiment.save.every_n_epochs == 0) + ) + epoch_list_check = epoch in config.experiment.save.epochs + should_save_ckpt = time_check or epoch_check or epoch_list_check + ckpt_reason = None + if should_save_ckpt: + last_ckpt_time = time.time() + ckpt_reason = "time" + + print(f"Train Epoch {epoch}") + print(json.dumps(step_log, sort_keys=True, indent=4)) + for k, v in step_log.items(): + if k.startswith("Time_"): + data_logger.record(f"Timing_Stats/Train_{k[5:]}", v, epoch) + else: + data_logger.record(f"Train/{k}", v, epoch) + + # Evaluate the model on validation set + if config.experiment.validate: + with torch.no_grad(): + step_log = TrainUtils.run_epoch( + model=model, + data_loader=valid_loader, + epoch=epoch, + validate=True, + num_steps=valid_num_steps, + ) + for k, v in step_log.items(): + if k.startswith("Time_"): + data_logger.record(f"Timing_Stats/Valid_{k[5:]}", v, epoch) + else: + data_logger.record(f"Valid/{k}", v, epoch) + + print(f"Validation Epoch {epoch}") + print(json.dumps(step_log, sort_keys=True, indent=4)) + + # save checkpoint if achieve new best validation loss + valid_check = "Loss" in step_log + if valid_check and ( + best_valid_loss is None or (step_log["Loss"] <= best_valid_loss) + ): + best_valid_loss = step_log["Loss"] + if ( + config.experiment.save.enabled + and config.experiment.save.on_best_validation + ): + epoch_ckpt_name += f"_best_validation_{best_valid_loss}" + should_save_ckpt = True + ckpt_reason = "valid" if ckpt_reason is None else ckpt_reason + + # Evaluate the model by by running rollouts + + # do rollouts at fixed rate or if it's time to save a new ckpt + video_paths = None + rollout_check = (epoch % config.experiment.rollout.rate == 0) or ( + should_save_ckpt and ckpt_reason == "time" + ) + if ( + config.experiment.rollout.enabled + and (epoch > config.experiment.rollout.warmstart) + and rollout_check + ): + # wrap model as a RolloutPolicy to prepare for rollouts + rollout_model = RolloutPolicy( + model, obs_normalization_stats=obs_normalization_stats + ) + + num_episodes = config.experiment.rollout.n + all_rollout_logs, video_paths = TrainUtils.rollout_with_stats( + policy=rollout_model, + envs=envs, + horizon=config.experiment.rollout.horizon, + use_goals=config.use_goals, + num_episodes=num_episodes, + render=False, + video_dir=video_dir if config.experiment.render_video else None, + epoch=epoch, + video_skip=config.experiment.get("video_skip", 5), + terminate_on_success=config.experiment.rollout.terminate_on_success, + ) + + # summarize results from rollouts to tensorboard and terminal + for env_name in all_rollout_logs: + rollout_logs = all_rollout_logs[env_name] + for k, v in rollout_logs.items(): + if k.startswith("Time_"): + data_logger.record( + f"Timing_Stats/Rollout_{env_name}_{k[5:]}", v, epoch + ) + else: + data_logger.record( + f"Rollout/{k}/{env_name}", v, epoch, log_stats=True + ) + + print( + "\nEpoch {} Rollouts took {}s (avg) with results:".format( + epoch, rollout_logs["time"] + ) + ) + print(f"Env: {env_name}") + print(json.dumps(rollout_logs, sort_keys=True, indent=4)) + + # checkpoint and video saving logic + updated_stats = TrainUtils.should_save_from_rollout_logs( + all_rollout_logs=all_rollout_logs, + best_return=best_return, + best_success_rate=best_success_rate, + epoch_ckpt_name=epoch_ckpt_name, + save_on_best_rollout_return=config.experiment.save.on_best_rollout_return, + save_on_best_rollout_success_rate=config.experiment.save.on_best_rollout_success_rate, + ) + best_return = updated_stats["best_return"] + best_success_rate = updated_stats["best_success_rate"] + epoch_ckpt_name = updated_stats["epoch_ckpt_name"] + should_save_ckpt = ( + config.experiment.save.enabled and updated_stats["should_save_ckpt"] + ) or should_save_ckpt + if updated_stats["ckpt_reason"] is not None: + ckpt_reason = updated_stats["ckpt_reason"] + + # Only keep saved videos if the ckpt should be saved (but not because of validation score) + should_save_video = ( + should_save_ckpt and (ckpt_reason != "valid") + ) or config.experiment.keep_all_videos + if video_paths is not None and not should_save_video: + for env_name in video_paths: + os.remove(video_paths[env_name]) + + # Save model checkpoints based on conditions (success rate, validation loss, etc) + if should_save_ckpt: + TrainUtils.save_model( + model=model, + config=config, + env_meta=env_meta, + shape_meta=shape_meta, + ckpt_path=os.path.join(ckpt_dir, epoch_ckpt_name + ".pth"), + obs_normalization_stats=obs_normalization_stats, + ) + + # Finally, log memory usage in MB + process = psutil.Process(os.getpid()) + mem_usage = int(process.memory_info().rss / 1000000) + data_logger.record("System/RAM Usage (MB)", mem_usage, epoch) + print(f"\nEpoch {epoch} Memory Usage: {mem_usage} MB\n") + + # terminate logging + data_logger.close() + + except Exception as e: + res_str = f"run failed with error:\n{e}\n\n{traceback.format_exc()}" + print(res_str) + + +### Helper functions ### +def parse_cli_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + algorithm_group = parser.add_argument_group( + "algorithm arguments", + description="Arguments for algorithm.", + ) + algorithm_group.add_argument( + "--algo", + type=str, + default="bc", + help="Name of the algorithm\n(bc, bcq)", + ) + algorithm_group.add_argument( + "--dataset", + type=str, + default=None, + help="(optional) if provided, override the dataset path defined in the config", + ) + algorithm_group.add_argument( + "--name", + type=str, + default=None, + help="(optional) if provided, override the experiment name defined in the config", + ) + add_default_cli_args(parser) + return parser.parse_args() + + +if __name__ == "__main__": + # Parse arguments + args = parse_cli_args() + + # Launch the app + launcher = launch_app(args) + + # Run the main function + main(launcher=launcher, args=args) + + # Shutdown the app + shutdown_app(launcher) diff --git a/scripts/algo/sb3/play.py b/scripts/algo/sb3/play.py new file mode 100755 index 0000000..23a1088 --- /dev/null +++ b/scripts/algo/sb3/play.py @@ -0,0 +1,127 @@ +#!/root/isaac-sim/python.sh + +import os +import sys + +from omni.isaac.lab.app import AppLauncher + +sys.path.append( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))) +) +from _cli_utils import add_default_cli_args, argparse, launch_app, shutdown_app + +FRAMEWORK_NAME = "sb3" +ALGO_CFG_ENTRYPOINT_KEY = f"{FRAMEWORK_NAME}_{{ALGO_NAME}}_cfg" + + +def main(launcher: AppLauncher, args: argparse.Namespace): + import gymnasium + import numpy as np + import torch + from omni.isaac.kit import SimulationApp + from stable_baselines3 import PPO + from stable_baselines3.common.vec_env import VecNormalize + + import space_robotics_bench # noqa: F401 + from space_robotics_bench.core.wrappers.sb3 import Sb3VecEnvWrapper, process_sb3_cfg + from space_robotics_bench.utils.parsing import ( + get_checkpoint_path, + load_cfg_from_registry, + parse_task_cfg, + ) + + ## Extract simulation app + sim_app: SimulationApp = launcher.app + + # parse configuration + task_cfg = parse_task_cfg( + args.task, + device=args.device, + num_envs=args.num_envs, + use_fabric=not args.disable_fabric, + ) + agent_cfg = load_cfg_from_registry( + args.task, ALGO_CFG_ENTRYPOINT_KEY.format(ALGO_NAME=args.algo) + ) + # post-process agent configuration + agent_cfg = process_sb3_cfg(agent_cfg) + + # create isaac environment + env = gymnasium.make(args.task, cfg=task_cfg) + # wrap around environment for stable baselines + env = Sb3VecEnvWrapper(env) + + # normalize environment (if needed) + if "normalize_input" in agent_cfg: + env = VecNormalize( + env, + training=True, + norm_obs="normalize_input" in agent_cfg + and agent_cfg.pop("normalize_input"), + norm_reward="normalize_value" in agent_cfg + and agent_cfg.pop("normalize_value"), + clip_obs="clip_obs" in agent_cfg and agent_cfg.pop("clip_obs"), + gamma=agent_cfg["gamma"], + clip_reward=np.inf, + ) + + # directory for logging into + log_root_path = os.path.join("logs", FRAMEWORK_NAME, args.algo, args.task) + log_root_path = os.path.abspath(log_root_path) + # check checkpoint is valid + if args.checkpoint is None: + if args.use_last_checkpoint: + checkpoint = "model_.*.zip" + else: + checkpoint = "model.zip" + checkpoint_path = get_checkpoint_path(log_root_path, ".*", checkpoint) + else: + checkpoint_path = args.checkpoint + # create agent from stable baselines + print(f"Loading checkpoint from: {checkpoint_path}") + agent = PPO.load(checkpoint_path, env, print_system_info=True) + + # reset environment + obs = env.reset() + # simulate environment + while sim_app.is_running(): + # run everything in inference mode + with torch.inference_mode(): + # agent stepping + actions, _ = agent.predict(obs, deterministic=True) + # env stepping + obs, _, _, _ = env.step(actions) + + # close the simulator + env.close() + + +### Helper functions ### +def parse_cli_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + algorithm_group = parser.add_argument_group( + "algorithm arguments", + description="Arguments for algorithm.", + ) + algorithm_group.add_argument( + "--algo", + type=str, + default="ppo", + help="Name of the algorithm\n(ppo, sac)", + ) + add_default_cli_args(parser) + return parser.parse_args() + + +if __name__ == "__main__": + # Parse arguments + args = parse_cli_args() + + # Launch the app + launcher = launch_app(args) + + # Run the main function + main(launcher=launcher, args=args) + + # Shutdown the app + shutdown_app(launcher) diff --git a/scripts/algo/sb3/train.py b/scripts/algo/sb3/train.py new file mode 100755 index 0000000..388fe78 --- /dev/null +++ b/scripts/algo/sb3/train.py @@ -0,0 +1,163 @@ +#!/root/isaac-sim/python.sh + +import os +import sys + +from omni.isaac.lab.app import AppLauncher + +sys.path.append( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))) +) +from _cli_utils import add_default_cli_args, argparse, launch_app, shutdown_app + +FRAMEWORK_NAME = "sb3" +ALGO_CFG_ENTRYPOINT_KEY = f"{FRAMEWORK_NAME}_{{ALGO_NAME}}_cfg" + + +def main(launcher: AppLauncher, args: argparse.Namespace): + from datetime import datetime + + import gymnasium + import numpy as np + from omni.isaac.kit import SimulationApp + from omni.isaac.lab.utils.dict import print_dict + from omni.isaac.lab.utils.io import dump_pickle, dump_yaml + from sb3_contrib import TQC, RecurrentPPO + from stable_baselines3 import PPO, SAC + from stable_baselines3.common.callbacks import CheckpointCallback + from stable_baselines3.common.logger import configure + from stable_baselines3.common.vec_env import VecNormalize + + import space_robotics_bench # noqa: F401 + from space_robotics_bench.core.wrappers.sb3 import Sb3VecEnvWrapper, process_sb3_cfg + from space_robotics_bench.utils.parsing import ( + load_cfg_from_registry, + parse_task_cfg, + ) + + ## Extract simulation app + _sim_app: SimulationApp = launcher.app + + # parse configuration + task_cfg = parse_task_cfg( + args.task, + device=args.device, + num_envs=args.num_envs, + use_fabric=not args.disable_fabric, + ) + agent_cfg = load_cfg_from_registry( + args.task, ALGO_CFG_ENTRYPOINT_KEY.format(ALGO_NAME=args.algo) + ) + + # override configuration with command line arguments + if args.seed is not None: + agent_cfg["seed"] = args.seed + + # directory for logging into + log_dir = os.path.join( + "logs", + FRAMEWORK_NAME, + args.algo, + args.task, + datetime.now().strftime("%Y-%m-%d_%H-%M-%S"), + ) + # dump the configuration into log-directory + dump_yaml(os.path.join(log_dir, "params", "task.yaml"), task_cfg) + dump_yaml(os.path.join(log_dir, "params", "agent.yaml"), agent_cfg) + dump_pickle(os.path.join(log_dir, "params", "task.pkl"), task_cfg) + dump_pickle(os.path.join(log_dir, "params", "agent.pkl"), agent_cfg) + + # post-process agent configuration + agent_cfg = process_sb3_cfg(agent_cfg) + # read configurations about the agent-training + policy_arch = agent_cfg.pop("policy") + n_timesteps = agent_cfg.pop("n_timesteps") + + # create isaac environment + env = gymnasium.make( + args.task, cfg=task_cfg, render_mode="rgb_array" if args.video else None + ) + # wrap for video recording + if args.video: + video_kwargs = { + "video_folder": os.path.join(log_dir, "videos"), + "step_trigger": lambda step: step % args.video_interval == 0, + "video_length": args.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gymnasium.wrappers.RecordVideo(env, **video_kwargs) + # wrap around environment for stable baselines + env = Sb3VecEnvWrapper(env) + # set the seed + env.seed(seed=agent_cfg["seed"]) + + if "normalize_input" in agent_cfg: + env = VecNormalize( + env, + training=True, + norm_obs="normalize_input" in agent_cfg + and agent_cfg.pop("normalize_input"), + norm_reward="normalize_value" in agent_cfg + and agent_cfg.pop("normalize_value"), + clip_obs="clip_obs" in agent_cfg and agent_cfg.pop("clip_obs"), + gamma=agent_cfg["gamma"], + clip_reward=np.inf, + ) + + # create agent from stable baselines + if args.algo == "ppo": + agent = PPO(policy_arch, env, verbose=1, **agent_cfg) + elif args.algo == "ppo_lstm": + agent = RecurrentPPO(policy_arch, env, verbose=1, **agent_cfg) + elif args.algo == "sac": + agent = SAC(policy_arch, env, verbose=1, **agent_cfg) + elif args.algo == "tqc": + agent = TQC(policy_arch, env, verbose=1, **agent_cfg) + # configure the logger + new_logger = configure(log_dir, ["stdout", "tensorboard"]) + agent.set_logger(new_logger) + + # callbacks for agent + checkpoint_callback = CheckpointCallback( + save_freq=1000, save_path=log_dir, name_prefix="model", verbose=2 + ) + # train the agent + agent.learn(total_timesteps=n_timesteps, callback=checkpoint_callback) + # save the final model + agent.save(os.path.join(log_dir, "model")) + + # close the simulator + env.close() + + +### Helper functions ### +def parse_cli_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + algorithm_group = parser.add_argument_group( + "algorithm arguments", + description="Arguments for algorithm.", + ) + algorithm_group.add_argument( + "--algo", + type=str, + default="ppo", + help="Name of the algorithm\n(ppo, sac, ppo_lstm)", + ) + add_default_cli_args(parser) + return parser.parse_args() + + +if __name__ == "__main__": + # Parse arguments + args = parse_cli_args() + + # Launch the app + launcher = launch_app(args) + + # Run the main function + main(launcher=launcher, args=args) + + # Shutdown the app + shutdown_app(launcher) diff --git a/scripts/blender/__init__.py b/scripts/blender/__init__.py new file mode 100755 index 0000000..636498f --- /dev/null +++ b/scripts/blender/__init__.py @@ -0,0 +1,23 @@ +#!/usr/bin/env -S blender --python-exit-code 1 --python +""" +Script for setting up Blender preferences. + +- Use Cycles with CUDA GPU +- Enable Node Wrangler addon +""" + +import bpy + +# Use Cycles with GPU +bpy.context.preferences.addons["cycles"].preferences.compute_device_type = "CUDA" +bpy.context.preferences.addons["cycles"].preferences.get_devices() +for device in bpy.context.preferences.addons["cycles"].preferences.devices: + device.use = device.type == "CUDA" +bpy.context.scene.render.engine = "CYCLES" +bpy.context.scene.cycles.device = "GPU" + +# Enable Node Wrangler +bpy.ops.preferences.addon_enable(module="node_wrangler") + +# Save preferences +bpy.ops.wm.save_userpref() diff --git a/scripts/blender/procgen_assets.py b/scripts/blender/procgen_assets.py new file mode 100755 index 0000000..b2900be --- /dev/null +++ b/scripts/blender/procgen_assets.py @@ -0,0 +1,913 @@ +#!/usr/bin/env -S blender --factory-startup --background --offline-mode --enable-autoexec --python-exit-code 1 --python +""" +Script for automated procedural asset generation using Blender that revolves around its rich +node-based system for geometry (Geometry Nodes) and materials (Shader Nodes). + +Overview: + The requested node trees are constructed via scripts defined by `--autorun_scripts`. A sequence of + Geometry Nodes moodifiers is then applied to a prototype object that is duplicated for each generated + variant. Once the geometry is finalized, a procedural material is applied and baked into PBR textures + before exporting the final model. + +Example (manual invocation is not recommended): + blender --python procgen_assets.py -- \ + --autorun_scripts path/to/nodes_0.py path/to/nodes_1.py ... \ + --geometry_nodes '{"NodeModifierName": {"input_name": input_value, ...}, ...}' \ + --material MaterialName \ + --texture_resolution 4096 \ + --num_assets 10 \ + --outdir path/to/output/directory +""" + +from __future__ import annotations + +import argparse +import contextlib +import enum +import io +import json +import os +import re +import sys +import time +from copy import deepcopy +from os import path +from pathlib import Path +from typing import Any, Dict, List, Literal, Optional, TextIO, Tuple, Union + +import bpy + + +def main(**kwargs): + print_bpy(f"[INFO]: Starting procedural asset generation with kwargs: {kwargs}") + verify_requirements() + ProceduralGenerator.generate(**kwargs) + + +class ProceduralGenerator: + """ + Generator of procedural models using Blender. + """ + + @classmethod + def generate( + cls, + *, + ## Input + autorun_scripts: List[str], + ## Output + outdir: str, + name: str, + ext: str, + overwrite_min_age: int, + ## Generator + seed: int, + num_assets: int, + ## Export + export_kwargs: Dict[str, Any], + ## Geometry + geometry_nodes: Dict[str, Dict[str, Any]], + decimate_angle_limit: Optional[float], + decimate_face_count: Optional[int], + ## Material + material: Optional[str], + texture_resolution: int, + ): + """ + Entrypoint method for generating a set of procedural models. + """ + + # Preprocess the input arguments + outdir = Path(outdir).absolute().as_posix() + if not ext.startswith("."): + ext = f".{ext}" + + # Create the output directory + os.makedirs(name=outdir, exist_ok=True) + + # Reset to factory settings with an empty scene + bpy.ops.wm.read_factory_settings(use_empty=True) + + # Autorun all input scripts + for script in autorun_scripts: + print_bpy(f"[INFO]: Running script: {script}") + bpy.ops.script.python_file_run(filepath=script) + + # Extract a map of the input socket mapping for all loaded node groups + node_group_input_socket_maps = cls.extract_aliased_input_socket_id_mappings() + + # Create an empty mesh object and treat it as a prototype that will be duplicated for each processed seed + bpy.ops.object.add(type="MESH") + proto_obj: bpy.types.Object = bpy.context.active_object + proto_obj.name = name + proto_obj.data.name = name + + # Apply the requested list of geometry nodes to the prototype object + cls.apply_geometry_nodes_modifiers( + obj=proto_obj, + geometry_nodes=geometry_nodes, + node_group_input_socket_maps=node_group_input_socket_maps, + ) + + # Preheat the oven for baking the material into PBR textures + if material: + ProceduralGenerator.Baker.preheat_oven() + + # Generate models over the specified range + for current_seed in range(seed, seed + num_assets): + # Form the output filepath + filepath = path.join(outdir, f"{name}{current_seed}{ext}") + + # Skip generation if the file already exists and is too recent + if path.exists(filepath) and ( + overwrite_min_age < 0 + or (overwrite_min_age > time.time() - path.getmtime(filepath)) + ): + print_bpy( + f"[INFO]: Skipping generation of '{filepath}' because it was generated in the last {overwrite_min_age} seconds" + ) + continue + + # Duplicate the prototype object, rename it, select it and mark it as the active object + obj = cls.duplicate_object(obj=proto_obj, seed=current_seed) + + # Update the random seed of all geometry nodes modifiers that have a seed input socket + geometry_nodes_modifiers = [ + modifier + for modifier in obj.modifiers.values() + if modifier.type == "NODES" + ] + for modifier in geometry_nodes_modifiers: + if seed_id := node_group_input_socket_maps[ + modifier.node_group.name + ].get("seed"): + modifier[seed_id] = current_seed + + # Apply changes via mesh update + obj.data.update() + + # Apply all modifiers + for modifier in obj.modifiers: + bpy.ops.object.modifier_apply(modifier=modifier.name) + + # If specified, bake the material into PBR textures + if material: + obj.data.materials.clear() + obj.data.materials.append(bpy.data.materials.get(material)) + ProceduralGenerator.Baker.bake_into_pbr_material( + obj=obj, texture_resolution=texture_resolution + ) + + # Decimate the mesh if necessary + if decimate_angle_limit: + bpy.ops.object.modifier_add(type="DECIMATE") + obj.modifiers["Decimate"].decimate_type = "DISSOLVE" + obj.modifiers["Decimate"].angle_limit = decimate_angle_limit + + if decimate_face_count: + # Decimate the mesh + bpy.ops.object.modifier_add(type="DECIMATE") + obj.modifiers["Decimate"].ratio = decimate_face_count / len( + obj.data.polygons + ) + bpy.ops.object.modifier_apply(modifier="Decimate") + + # Export the model + cls.Exporter.usd_export( + filepath=filepath, ext=ext, makedirs=False, **export_kwargs + ) + print_bpy( + f"[LOG]: Generated asset #{current_seed} ({current_seed-seed+1}/{num_assets}): {filepath}" + ) + + # Update the viewport to keep track of progress + if not bpy.app.background: + bpy.ops.wm.redraw_timer(type="DRAW_WIN_SWAP", iterations=1) + + # Remove the generated object + bpy.data.objects.remove(obj) + + print_bpy("[INFO]: Generation completed") + + @staticmethod + def extract_aliased_input_socket_id_mappings() -> Dict[str, Dict[str, str]]: + def _extract_input_socket_id_mapping( + node_group: bpy.types.NodeTree, + ) -> Dict[str, str]: + return { + canonicalize_str(item.name): item.identifier + for item in node_group.interface.items_tree.values() + if item.item_type == "SOCKET" and item.in_out == "INPUT" + } + + # Extract a map of the input socket mapping for all node groups + node_group_input_socket_maps = { + node_group_name: _extract_input_socket_id_mapping(node_group) + for node_group_name, node_group in bpy.data.node_groups.items() + } + + # Rename common aliases for convenience + COMMON_ALIASES: Dict[str, List[str]] = { + "seed": [ + "pseodorandomseed", + "randomseed", + "rng", + ], + "detail": [ + "detaillevel", + "detailobject", + "levelofdetail", + "subdivisionlevel", + "subdivisions", + "subdivlevel", + ], + } + for node_group_input_socket_map in node_group_input_socket_maps.values(): + for target, possible_alias in COMMON_ALIASES.items(): + original_alias: Optional[str] = ( + target if target in node_group_input_socket_map.keys() else None + ) + for key in node_group_input_socket_map.keys(): + if key in possible_alias: + if original_alias is not None: + raise ValueError( + "Ambiguous name of the input socket '{target}' (canonicalized): '{original_alias}', '{key}'" + ) + original_alias = key + if original_alias is not None and original_alias != target: + node_group_input_socket_map[target] = node_group_input_socket_map[ + original_alias + ] + + return node_group_input_socket_maps + + @staticmethod + def apply_geometry_nodes_modifiers( + obj: bpy.types.Object, + *, + geometry_nodes: Dict[str, Dict[str, Any]], + node_group_input_socket_maps: Dict[str, Dict[str, str]], + ): + for i, (node_group_name, node_group_inputs) in enumerate( + geometry_nodes.items() + ): + # Create a new nodes modifier + modifier: bpy.types.NodesModifier = obj.modifiers.new( + name=f"node{i}", type="NODES" + ) + + # Assign the requested node group + if node_group := bpy.data.node_groups.get(node_group_name): + modifier.node_group = node_group + else: + raise ValueError( + f"Node group '{node_group_name}' not found in the list of available groups: {bpy.data.node_groups.keys()}" + ) + + # Set inputs accordingly + for key, value in node_group_inputs.items(): + socket_id = node_group_input_socket_maps[node_group_name][ + canonicalize_str(key) + ] + modifier[socket_id] = value + + # Apply changes via mesh update + obj.data.update() + + @staticmethod + def duplicate_object(obj: bpy.types.Object, seed: int) -> bpy.types.Object: + # Select the object to duplicate + bpy.ops.object.select_all(action="DESELECT") + obj.select_set(True) + bpy.context.view_layer.objects.active = obj + + # Duplicate the object + bpy.ops.object.duplicate() + + # Get the new duplicated object (it will be the active object) + new_obj = bpy.context.active_object + + # Rename the new object based on the seed + new_obj.name = f"{obj.name}_{seed}" + new_obj.data.name = f"{obj.data.name}_{seed}" + + # Select the new object and make it the active object + bpy.ops.object.select_all(action="DESELECT") + new_obj.select_set(True) + bpy.context.view_layer.objects.active = new_obj + + return new_obj + + class Recipe(enum.Enum): + ALBEDO = enum.auto() + METALLIC = enum.auto() + SPECULAR = enum.auto() + ROUGHNESS = enum.auto() + NORMAL = enum.auto() + + @staticmethod + def enabled_recipes(): + return ( + ProceduralGenerator.Recipe.ALBEDO, + ProceduralGenerator.Recipe.METALLIC, + # ProceduralGenerator.Recipe.SPECULAR, + ProceduralGenerator.Recipe.ROUGHNESS, + ProceduralGenerator.Recipe.NORMAL, + ) + + @property + def bake_type(self): + match self: + case self.ALBEDO: + return "DIFFUSE" + case self.METALLIC: + return "EMIT" + case self.SPECULAR: + return "GLOSSY" + case _: + return self.name + + @property + def color_space(self): + return "sRGB" if self.ALBEDO == self else "Non-Color" + + @property + def shader_socket_name(self): + match self: + case self.ALBEDO: + return "Base Color" + case self.METALLIC: + return "Metallic" + case self.SPECULAR: + return "Specular Tint" + case self.ROUGHNESS: + return "Roughness" + case self.NORMAL: + return "Normal" + + def prep(self, material: bpy.types.Material) -> Dict[str, Any]: + match self: + # If the shader node has a "Metallic" input, it needs to be disabled for baking + case self.ALBEDO: + # Get the output material node + output_material_node = [ + node + for node in material.node_tree.nodes + if node.type == "OUTPUT_MATERIAL" + ][0] + + # Get surface shader socket + shader_node_socket = output_material_node.inputs["Surface"] + if not shader_node_socket.is_linked: + return {} + + # Get the shader node connected to the socket + shader_node = shader_node_socket.links[0].from_node + + # No need to do anything if the shader node does not have a "Metallic" input + if "Metallic" not in shader_node.inputs: + return {} + + metallic_socket = shader_node.inputs["Metallic"] + ingredients = {} + + # If the metallic input is linked, store the original link source and temporarily disconnect it + if metallic_socket.is_linked: + metallic_socket_links = metallic_socket.links[0] + ingredients["orig_metallic_from_socket"] = ( + metallic_socket_links.from_socket + ) + material.node_tree.links.remove(metallic_socket_links) + + # Always store the original default value and set it to 0.0 + ingredients["orig_metallic_default_value"] = ( + metallic_socket.default_value + ) + metallic_socket.default_value = 0.0 + + return ingredients + + # Render the metallic input as emission originating from the surface shader's metallic input + case self.METALLIC: + # Get the output material node + output_material_node = [ + node + for node in material.node_tree.nodes + if node.type == "OUTPUT_MATERIAL" + ][0] + + # Get surface shader socket + shader_node_socket = output_material_node.inputs["Surface"] + if not shader_node_socket.is_linked: + return {} + + # Get the shader link + shader_link = shader_node_socket.links[0] + + # Get the shader node connected to the socket + shader_node = shader_link.from_node + + # Store the original link source and temporarily disconnect it + ingredients = { + "orig_surface_shader_source": shader_link.from_socket + } + material.node_tree.links.remove(shader_link) + + # No need to do anything if the shader node does not have a "Metallic" input + orig_metallic_default_value = 0.0 + with_emissive_rgb_node = False + if "Metallic" in shader_node.inputs: + metallic_socket = shader_node.inputs["Metallic"] + + # If the metallic input is linked, store the original link source and temporarily disconnect it + if metallic_socket.is_linked: + metallic_socket_links = metallic_socket.links[0] + material.node_tree.links.new( + metallic_socket_links.from_socket, + shader_node_socket, + ) + else: + with_emissive_rgb_node = True + orig_metallic_default_value = metallic_socket.default_value + else: + with_emissive_rgb_node = True + + if with_emissive_rgb_node: + rgb_node = material.node_tree.nodes.new(type="ShaderNodeRGB") + rgb_node.outputs[0].default_value = ( + *((orig_metallic_default_value,) * 3), + 1.0, + ) + material.node_tree.links.new( + rgb_node.outputs["Color"], shader_node_socket + ) + ingredients["emissive_rgb_node"] = rgb_node + + return ingredients + + case _: + return {} + + def cleanup( + self, + material: bpy.types.Material, + *, + orig_metallic_from_socket: Optional[bpy.types.NodeSocket] = None, + orig_metallic_default_value: Optional[float] = None, + orig_surface_shader_source: Optional[bpy.types.NodeSocket] = None, + emissive_rgb_node: Optional[bpy.types.Node] = None, + ): + match self: + case self.ALBEDO: + if orig_metallic_default_value or orig_metallic_from_socket: + metallic_socket = ( + [ + node + for node in material.node_tree.nodes + if node.type == "OUTPUT_MATERIAL" + ][0] + .inputs["Surface"] + .links[0] + .inputs["Metallic"] + ) + metallic_socket.default_value = orig_metallic_default_value + if orig_metallic_from_socket: + material.node_tree.links.new( + orig_metallic_from_socket, + metallic_socket, + ) + + case self.METALLIC: + if orig_surface_shader_source: + shader_node_socket = [ + node + for node in material.node_tree.nodes + if node.type == "OUTPUT_MATERIAL" + ][0].inputs["Surface"] + shader_link = shader_node_socket.links[0] + material.node_tree.links.remove(shader_link) + material.node_tree.links.new( + orig_surface_shader_source, + shader_node_socket, + ) + if emissive_rgb_node: + material.node_tree.nodes.remove(emissive_rgb_node) + + class Baker: + """ + Simple wrapper around Blender baking capabilities. + """ + + @staticmethod + def preheat_oven(): + # Only Cycles supports texture baking + bpy.context.scene.render.engine = "CYCLES" + bpy.data.scenes[0].render.engine = "CYCLES" + + # Bake using GPU + bpy.context.preferences.addons["cycles"].preferences.compute_device_type = ( + "CUDA" + ) + bpy.context.preferences.addons["cycles"].preferences.get_devices() + for device in bpy.context.preferences.addons["cycles"].preferences.devices: + device.use = device.type == "CUDA" + bpy.context.scene.render.engine = "CYCLES" + bpy.context.scene.cycles.device = "GPU" + + # Improve performance + bpy.context.scene.cycles.samples = 1 + bpy.context.scene.cycles.use_auto_tile = True + + # Consider only the color pass (no environment lighting) + bpy.context.scene.render.bake.use_pass_direct = False + bpy.context.scene.render.bake.use_pass_indirect = False + bpy.context.scene.render.bake.use_pass_color = True + + @classmethod + def bake_into_pbr_material( + cls, + obj: bpy.types.Object, + *, + texture_resolution: int, + ): + # Adjust the recipe according to the object + bpy.context.scene.render.bake.margin = texture_resolution // 64 + + # Unwrap the object if necessary + if not obj.data.uv_layers: + cls._unwrap_uv_on_active_obj(obj, texture_resolution) + + # Get the material + material = obj.data.materials[0] + + # Bake all textures from the recipe + baked_textures = {} + for recipe in ProceduralGenerator.Recipe.enabled_recipes(): + ingredients = recipe.prep(material=material) + + # Create the image node into which the texture will be baked + image_node = cls._create_image_node( + material=material, + recipe=recipe, + texture_resolution=texture_resolution, + ) + + # Bake the texture + bpy.ops.object.bake(type=recipe.bake_type) + baked_textures[recipe] = image_node.image + + # Remove the image node and cleanup the material + material.node_tree.nodes.remove(image_node) + recipe.cleanup(material=material, **ingredients) + + # Create a new PBR material with the baked textures + pbr_material = cls._bake_into_pbr_material( + name=f"PBR_{obj.name}", baked_textures=baked_textures + ) + + # Replace the original material with the new PBR material + obj.data.materials.clear() + obj.data.materials.append(pbr_material) + + @classmethod + def _bake_into_pbr_material( + cls, + name: str, + baked_textures: Dict[ProceduralGenerator.Recipe, bpy.types.Image], + ) -> bpy.types.Material: + # Create a new material + pbr_material = bpy.data.materials.new(name=name) + pbr_material.use_nodes = True + + # Get handles to the nodes and links + nodes = pbr_material.node_tree.nodes + links = pbr_material.node_tree.links + + # Clear the existing nodes + nodes.clear() + + # Create Material Output and Principled BSDF nodes + shader_node = nodes.new(type="ShaderNodeBsdfPrincipled") + shader_node.location = (-300, 0) + output_node = nodes.new(type="ShaderNodeOutputMaterial") + output_node.location = (0, 0) + links.new(shader_node.outputs["BSDF"], output_node.inputs["Surface"]) + + # Create Texture Coordinate and Mapping nodes + texcoord_node = nodes.new(type="ShaderNodeTexCoord") + texcoord_node.location = (-1200, 0) + mapping_node = nodes.new(type="ShaderNodeMapping") + mapping_node.location = (-1000, 0) + links.new(mapping_node.inputs["Vector"], texcoord_node.outputs["UV"]) + + # Create baked textures + for i, (recipe, texture) in enumerate(baked_textures.items()): + # Create Image Texture node + img_texture = nodes.new(type="ShaderNodeTexImage") + img_texture.image = texture + img_texture.location = ( + -800, + 375 * (0.5 * len(baked_textures) + 0.5 - i), + ) + links.new(img_texture.inputs["Vector"], mapping_node.outputs["Vector"]) + + match recipe: + # Normal map requires the Normal Map node + case ProceduralGenerator.Recipe.NORMAL: + normal_map_node = nodes.new(type="ShaderNodeNormalMap") + normal_map_node.location = (-500, -127) + links.new( + img_texture.outputs["Color"], + normal_map_node.inputs["Color"], + ) + links.new( + normal_map_node.outputs["Normal"], + shader_node.inputs[recipe.shader_socket_name], + ) + + # Other textures are directly linked to the shader node + case _: + links.new( + img_texture.outputs["Color"], + shader_node.inputs[recipe.shader_socket_name], + ) + + return pbr_material + + @staticmethod + def _unwrap_uv_on_active_obj(obj: bpy.types.Object, texture_resolution: int): + bpy.ops.object.mode_set(mode="EDIT") + bpy.ops.mesh.select_all(action="SELECT") + + # Try with the default unwrap first, but capture the output in case it fails (printed as warning to stdout) + stdout_str = io.StringIO() + with contextlib.redirect_stdout(stdout_str): + bpy.ops.uv.unwrap() + + # If the default unwrap failed, try with the Smart UV Project + if "Unwrap failed" in stdout_str.getvalue(): + bpy.ops.uv.smart_project( + rotate_method="AXIS_ALIGNED", island_margin=0.5 / texture_resolution + ) + + bpy.ops.object.mode_set(mode="OBJECT") + + @classmethod + def _create_image_node( + cls, + material: bpy.types.Material, + recipe: ProceduralGenerator.Recipe, + texture_resolution: int, + ): + node = material.node_tree.nodes.new("ShaderNodeTexImage") + node.image = cls._create_image_texture( + name=recipe.name.lower(), + texture_resolution=texture_resolution, + color_space=recipe.color_space, + ) + + material.node_tree.nodes.active = node + + return node + + @staticmethod + def _create_image_texture( + name: str, + *, + texture_resolution: int, + color_space: Literal["Non-Color", "sRGB"], + ): + image = bpy.data.images.new( + name, + width=texture_resolution, + height=texture_resolution, + alpha=False, + float_buffer=False, + tiled=False, + ) + image.alpha_mode = "NONE" + image.colorspace_settings.name = color_space + + return image + + class Exporter: + """ + Simple wrapper around Blender export capabilities. + """ + + DEFAULT_USD_EXPORT_OVERRIDES = { + "check_existing": False, + "selected_objects_only": True, + "author_blender_name": False, + "use_instancing": True, + } + + @classmethod + def usd_export( + cls, + filepath: Union[str, Path], + *, + ext: str = ".usdz", + makedirs: bool = True, + **kwargs, + ): + """ + Export via `bpy.ops.wm.usd_export()` + """ + + # Prepare the output path + if not isinstance(filepath, Path): + filepath = Path(filepath) + filepath = filepath.with_suffix(ext).absolute() + + # Create parent directories if necessary + if makedirs: + os.makedirs(name=filepath.parent, exist_ok=True) + + # Forward export kwargs + export_kwargs = deepcopy(cls.DEFAULT_USD_EXPORT_OVERRIDES) + export_kwargs.update(kwargs) + + # Export the USD file + bpy.ops.wm.usd_export( + filepath=filepath.as_posix(), + **export_kwargs, + ) + + +### String utils ### +REGEX_CANONICALIZE_STR_PATTERN: re.Pattern = re.compile("[\W_]+") + + +def canonicalize_str(input: str) -> str: + """ + Canonicalizes a string by converting it to lowercase and removing unwanted characters. + + This function processes the input string to ensure it is in a standardized format, making it suitable for consistent usage in applications. It utilizes a predefined regular expression pattern to eliminate any characters that do not meet the specified criteria. + + Args: + input (str): The string to be canonicalized. + + Returns: + str: The canonicalized version of the input string. + """ + return REGEX_CANONICALIZE_STR_PATTERN.sub("", input.lower()) + + +### Misc utils ### +def print_bpy(msg: Any, file: Optional[TextIO] = sys.stdout, *args, **kwargs): + """ + Helper print function that also provides output inside the Blender console in addition to the system console. + """ + + print(msg, file=file, *args, **kwargs) + for window in bpy.context.window_manager.windows: + for area in window.screen.areas: + if area.type == "CONSOLE": + with bpy.context.temp_override( + window=window, screen=window.screen, area=area + ): + bpy.ops.console.scrollback_append( + text=str(msg), + type="ERROR" if file == sys.stderr else "OUTPUT", + ) + + +def verify_requirements(): + VERSION_BPY_MIN: Tuple[int, int] = (4, 2) + if ( + bpy.app.version[0] != VERSION_BPY_MIN[0] + or bpy.app.version[1] < VERSION_BPY_MIN[1] + ): + print_bpy( + f"[WARNING]: Blender {bpy.app.version_string} is likely incompatible with this script (written for Blender {VERSION_BPY_MIN[0]}.{VERSION_BPY_MIN[1]})", + file=sys.stderr, + ) + + +### CLI ### +def parse_cli_args() -> argparse.Namespace: + """ + Parse command-line arguments for this script. + """ + + parser = argparse.ArgumentParser( + description="Generate procedural dataset of USD assets using Blender", + usage=f"{sys.argv[0] if sys.argv[0].endswith('blender') else 'blender'} --python {path.realpath(__file__)} -- [options]", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + argument_default=argparse.SUPPRESS, + ) + + group = parser.add_argument_group("Input") + group.add_argument( + "-i", + "--autorun_scripts", + type=str, + nargs="*", + help="List of Blender scripts to execute", + required=True, + ) + + group = parser.add_argument_group("Output") + group.add_argument( + "-o", + "--outdir", + type=str, + help="The output directory", + required=True, + ) + group.add_argument( + "--name", + type=str, + help="The base name of the exported models", + default="model", + ) + group.add_argument( + "--ext", + type=str, + help="The file extension of the exported models", + default=".usdz", + ) + group.add_argument( + "--overwrite_min_age", + type=int, + help="Number of seconds after which to overwrite the generated assets if they already exist (disabled if negative)", + default=0, + ) + + group = parser.add_argument_group("Generator") + group.add_argument( + "-s", + "--seed", + type=int, + help="The initial seed of the random number generator", + default=0, + ) + group.add_argument( + "-n", + "--num_assets", + type=int, + help="Number of assets to generate", + default=1, + ) + + group = parser.add_argument_group("Export") + group.add_argument( + "--export_kwargs", + type=json.loads, + help="Keyword arguments for the USD export", + default={}, + ) + + group = parser.add_argument_group("Geometry") + group.add_argument( + "--geometry_nodes", + type=json.loads, + help="List of Geometry Nodes modifiers from `--autorun_scripts` for generating the geometry, with an optional dictionary for configuring their inputs", + required=True, + ) + group.add_argument( + "--decimate_angle_limit", + type=float, + help="If specified, decimate the generated geometry to the specified target angle limit", + default=None, + ) + group.add_argument( + "--decimate_face_count", + type=int, + help="If specified, decimate the generated geometry to the specified target face count", + default=None, + ) + + group = parser.add_argument_group("Material") + group.add_argument( + "--material", + type=str, + help="Material of the generated models from `--autorun_scripts`, which will be baked as PBR textures into the USD file", + default=None, + ) + group.add_argument( + "--texture_resolution", + type=int, + help="Resolution of the baked PBR textures", + default=1024, + ) + + if "--" in sys.argv: + args = parser.parse_args(sys.argv[sys.argv.index("--") + 1 :]) + else: + args, unknown_args = parser.parse_known_args() + if unknown_args: + print_bpy( + f"[WARNING]: Unknown args: {unknown_args}", + file=sys.stderr, + ) + print_bpy( + '[HINT]: Consider delimiting your args for Python script with "--"' + ) + + return args + + +if __name__ == "__main__": + main(**vars(parse_cli_args())) diff --git a/scripts/gui.bash b/scripts/gui.bash new file mode 100755 index 0000000..8f5a61f --- /dev/null +++ b/scripts/gui.bash @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +### Run the GUI application +### Usage: gui.bash [--release] [CARGO_RUN_ARGS] +set -e + +SCRIPT_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" &>/dev/null && pwd)" +REPOSITORY_DIR="$(dirname "${SCRIPT_DIR}")" +CARGO_MANIFEST_PATH="${REPOSITORY_DIR}/Cargo.toml" + +## Config +# Options for running the application +RUN_OPTS="${RUN_OPTS:-}" +# Default to release build +WITH_RELEASE_BUILD="${WITH_RELEASE_BUILD:-true}" +if [[ "${WITH_RELEASE_BUILD,,}" = true ]]; then + RUN_OPTS+=" --release" +fi + +## Run with Cargo +CARGO_RUN_CMD=( + cargo run + --manifest-path "${CARGO_MANIFEST_PATH}" + --package space_robotics_bench_gui + --bin gui + "${RUN_OPTS}" + "${*:1}" +) +echo -e "\033[1;90m[TRACE] ${CARGO_RUN_CMD[*]}\033[0m" | xargs +# shellcheck disable=SC2048 +exec ${CARGO_RUN_CMD[*]} diff --git a/scripts/list_envs.py b/scripts/list_envs.py new file mode 100755 index 0000000..fcf7532 --- /dev/null +++ b/scripts/list_envs.py @@ -0,0 +1,65 @@ +#!/root/isaac-sim/python.sh +""" +Utility script for listing all registered the Space Robotics Bench + +Usage: + ros2 run space_robotics_bench list_envs.py +""" + +from _cli_utils import add_default_cli_args, argparse, launch_app, shutdown_app + + +def main(): + import gymnasium + from prettytable import PrettyTable + + import space_robotics_bench # noqa: F401 + from space_robotics_bench.utils.registry import get_srb_tasks + + # Table config + table = PrettyTable(["id (task/demo)", "entrypoint", "config"]) + table.title = "Space Robotics Bench" + table.align = "l" + + ## Fill the table with tasks + for task_id in get_srb_tasks(): + spec = gymnasium.registry[task_id] + table.add_row( + [ + task_id, + str(spec.entry_point), + ( + str(spec.kwargs["task_cfg"]) + .removeprefix("") + .replace(spec.entry_point.split(":")[0], "$MOD") + ), + ] + ) + + ## Print the results + print(table) + + +### Helper functions ### +def parse_cli_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + add_default_cli_args(parser) + return parser.parse_args() + + +if __name__ == "__main__": + # Parse arguments + args = parse_cli_args() + + # Set headless mode + args.headless = True + + # Launch the app + launcher = launch_app(args) + + # Run the main function + main() + + # Shutdown the app + shutdown_app(launcher) diff --git a/scripts/random_agent.py b/scripts/random_agent.py new file mode 100755 index 0000000..835baf9 --- /dev/null +++ b/scripts/random_agent.py @@ -0,0 +1,101 @@ +#!/root/isaac-sim/python.sh +""" +Testing script for running tasks with random actions + +Examples: + ros2 run space_robotics_bench random_agent.py + ros2 run space_robotics_bench random_agent.py --task sample_collection --num_envs 4 +""" + +from _cli_utils import add_default_cli_args, argparse, launch_app, shutdown_app +from omni.isaac.lab.app import AppLauncher + + +def main(launcher: AppLauncher, args: argparse.Namespace): + ## Note: Importing modules here due to delayed Omniverse Kit extension loading + from os import path + + import gymnasium + import torch + from omni.isaac.kit import SimulationApp + from omni.isaac.lab.utils.dict import print_dict + + import space_robotics_bench # noqa: F401 + from space_robotics_bench.utils.parsing import create_logdir_path, parse_task_cfg + + ## Extract simulation app + sim_app: SimulationApp = launcher.app + + ## Configuration + task_cfg = parse_task_cfg( + task_name=args.task, + device=args.device, + num_envs=args.num_envs, + use_fabric=not args.disable_fabric, + ) + + ## Create the environment + env = gymnasium.make( + id=args.task, cfg=task_cfg, render_mode="rgb_array" if args.video else None + ) + + ## Initialize the environment + observation, info = env.reset() + + ## Add wrapper for video recording (if enabled) + if args.video: + logdir = create_logdir_path("random", args.task) + video_kwargs = { + "video_folder": path.join(logdir, "videos"), + "step_trigger": lambda step: step % args.video_interval == 0, + "video_length": args.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gymnasium.wrappers.RecordVideo(env, **video_kwargs) + + ## Run the environment with random actions + with torch.inference_mode(): + while sim_app.is_running(): + # Sample random actions + actions = torch.from_numpy(env.action_space.sample()).to( + device=env.unwrapped.device + ) + + # Step the environment + observation, reward, terminated, truncated, info = env.step(actions) + # print( + # f"> actions: {actions}\n" + # f"> observation: {observation}\n" + # f"> reward: {reward}\n" + # f"> terminated: {terminated}\n" + # f"> truncated: {truncated}\n" + # f"> info: {info}\n" + # ) + + # Note: Each environment is automatically reset (independently) when terminated or truncated + + ## Close the environment + env.close() + + +### Helper functions ### +def parse_cli_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + add_default_cli_args(parser) + return parser.parse_args() + + +if __name__ == "__main__": + # Parse arguments + args = parse_cli_args() + + # Launch the app + launcher = launch_app(args) + + # Run the main function + main(launcher=launcher, args=args) + + # Shutdown the app + shutdown_app(launcher) diff --git a/scripts/ros2.py b/scripts/ros2.py new file mode 100755 index 0000000..0e2f005 --- /dev/null +++ b/scripts/ros2.py @@ -0,0 +1,104 @@ +#!/root/isaac-sim/python.sh +""" +Entrypoint script for iterfacing over ROS 2 + +Examples: + ros2 run space_robotics_bench ros2.py + ros2 run space_robotics_bench ros2.py --task sample_collection --num_envs 4 +""" + +from _cli_utils import add_default_cli_args, argparse, launch_app, shutdown_app +from omni.isaac.lab.app import AppLauncher + + +def main(launcher: AppLauncher, args: argparse.Namespace): + ## Note: Importing modules here due to delayed Omniverse Kit extension loading + from os import path + + import gymnasium + import torch + from omni.isaac.kit import SimulationApp + from omni.isaac.lab.utils.dict import print_dict + + import space_robotics_bench # Noqa: F401 + from space_robotics_bench.core.interfaces import ROS2 + from space_robotics_bench.utils.parsing import create_logdir_path, parse_task_cfg + + ## Extract simulation app + sim_app: SimulationApp = launcher.app + + # Parse configuration + task_cfg = parse_task_cfg( + task_name=args.task, + device=args.device, + num_envs=args.num_envs, + use_fabric=not args.disable_fabric, + ) + # Disable truncation + if hasattr(task_cfg, "enable_truncation"): + task_cfg.enable_truncation = False + + ## Create the environment + env = gymnasium.make( + id=args.task, cfg=task_cfg, render_mode="rgb_array" if args.video else None + ) + + ## Initialize the environment + observation, info = env.reset() + + ## Create ROS 2 interface + ros2_interface = ROS2(env) + + ## Add wrapper for video recording (if enabled) + if args.video: + logdir = create_logdir_path("ros2", args.task) + video_kwargs = { + "video_folder": path.join(logdir, "videos"), + "step_trigger": lambda step: step % args.video_interval == 0, + "video_length": args.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gymnasium.wrappers.RecordVideo(env, **video_kwargs) + + ## Run the environment with ROS 2 interface + with torch.inference_mode(): + while sim_app.is_running(): + # Get actions from ROS 2 + actions = ros2_interface.actions + + # Step the environment + observation, reward, terminated, truncated, info = env.step(actions) + + # Publish to ROS 2 + ros2_interface.publish(observation, reward, terminated, truncated, info) + + # Process requests from ROS 2 + ros2_interface.update() + + # Note: Each environment is automatically reset (independently) when terminated or truncated + + ## Close the environment + env.close() + + +### Helper functions ### +def parse_cli_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + add_default_cli_args(parser) + return parser.parse_args() + + +if __name__ == "__main__": + # Parse arguments + args = parse_cli_args() + + # Launch the app + launcher = launch_app(args) + + # Run the main function + main(launcher=launcher, args=args) + + # Shutdown the app + shutdown_app(launcher) diff --git a/scripts/teleop.py b/scripts/teleop.py new file mode 100755 index 0000000..8a85faf --- /dev/null +++ b/scripts/teleop.py @@ -0,0 +1,185 @@ +#!/root/isaac-sim/python.sh +""" +Teleoperation script for running tasks with human-in-the-loop control + +Examples: + ros2 run space_robotics_bench teleop.py + ros2 run space_robotics_bench teleop.py --task sample_collection --num_envs 4 --teleop_device spacemouse +""" + +from _cli_utils import add_default_cli_args, argparse, launch_app, shutdown_app +from omni.isaac.lab.app import AppLauncher + + +def main(launcher: AppLauncher, args: argparse.Namespace): + ## Note: Importing modules here due to delayed Omniverse Kit extension loading + from os import path + + import gymnasium + import numpy as np + import torch + from omni.isaac.kit import SimulationApp + from omni.isaac.lab.utils.dict import print_dict + + import space_robotics_bench # Noqa: F401 + from space_robotics_bench.core import mdp + from space_robotics_bench.core.actions import ( + ManipulatorTaskSpaceActionCfg, + MultiCopterActionGroupCfg, + WheeledRoverActionGroupCfg, + ) + from space_robotics_bench.core.managers import SceneEntityCfg + from space_robotics_bench.core.teleop_devices import CombinedInterface + from space_robotics_bench.utils.parsing import create_logdir_path, parse_task_cfg + + if args.headless and "keyboard" in args.teleop_device: + raise ValueError("Native teleoperation is only supported in GUI mode.") + + ## Extract simulation app + sim_app: SimulationApp = launcher.app + + # Parse configuration + task_cfg = parse_task_cfg( + task_name=args.task, + device=args.device, + num_envs=args.num_envs, + use_fabric=not args.disable_fabric, + ) + # Disable truncation + if hasattr(task_cfg, "enable_truncation"): + task_cfg.enable_truncation = False + + ## Create the environment + env = gymnasium.make( + id=args.task, cfg=task_cfg, render_mode="rgb_array" if args.video else None + ) + + ## Create controller + teleop_interface = CombinedInterface( + devices=args.teleop_device, + pos_sensitivity=args.pos_sensitivity, + rot_sensitivity=args.rot_sensitivity, + action_cfg=env.unwrapped.cfg.actions, + ) + teleop_interface.reset() + teleop_interface.add_callback("L", env.reset) + print(teleop_interface) + + ## Initialize the environment + observation, info = env.reset() + + ## Add wrapper for video recording (if enabled) + if args.video: + logdir = create_logdir_path("srb", args.task) + video_kwargs = { + "video_folder": path.join(logdir, "videos"), + "step_trigger": lambda step: step % args.video_interval == 0, + "video_length": args.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gymnasium.wrappers.RecordVideo(env, **video_kwargs) + + def process_actions(twist: np.ndarray, gripper_cmd: bool) -> torch.Tensor: + twist = torch.tensor( + twist, dtype=torch.float32, device=env.unwrapped.device + ).repeat(env.unwrapped.num_envs, 1) + if isinstance(env.unwrapped.cfg.actions, ManipulatorTaskSpaceActionCfg): + if not args.disable_control_scheme_inversion: + twist[:, :2] *= -1.0 + gripper_action = torch.zeros(twist.shape[0], 1, device=twist.device) + gripper_action[:] = -1.0 if gripper_cmd else 1.0 + return torch.concat([twist, gripper_action], dim=1) + elif isinstance(env.unwrapped.cfg.actions, MultiCopterActionGroupCfg): + return torch.concat( + [ + twist[:, :3], + twist[:, 5].unsqueeze(1), + ], + dim=1, + ) + + elif isinstance(env.unwrapped.cfg.actions, WheeledRoverActionGroupCfg): + return twist[:, :2] + + ## Run the environment + with torch.inference_mode(): + while sim_app.is_running(): + # Get actions from the teleoperation interface + actions = process_actions(*teleop_interface.advance()) + + # Step the environment + observation, reward, terminated, truncated, info = env.step(actions) + + # Note: Each environment is automatically reset (independently) when terminated or truncated + + # Provide force feedback for teleop devices + if isinstance(env.unwrapped.cfg.actions, ManipulatorTaskSpaceActionCfg): + FT_FEEDBACK_SCALE = torch.tensor([0.16, 0.16, 0.16, 0.0, 0.0, 0.0]) + ft_feedback_asset_cfg = SceneEntityCfg( + "robot", + body_names=task_cfg.robot_cfg.regex_links_hand, + ) + ft_feedback_asset_cfg.resolve(env.unwrapped.scene) + ft_feedback = ( + FT_FEEDBACK_SCALE + * mdp.body_incoming_wrench_mean( + env=env.unwrapped, + asset_cfg=ft_feedback_asset_cfg, + )[0, ...].cpu() + ) + teleop_interface.set_ft_feedback(ft_feedback) + + ## Close the environment + env.close() + + +### Helper functions ### +def parse_cli_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + algorithm_group = parser.add_argument_group( + "teleop", + description="Arguments for teleoperation", + ) + algorithm_group.add_argument( + "--teleop_device", + type=str, + nargs="+", + default=["keyboard"], + help="Device for interacting with environment", + ) + algorithm_group.add_argument( + "--pos_sensitivity", + type=float, + default=10.0, + help="Sensitivity factor for translation.", + ) + algorithm_group.add_argument( + "--rot_sensitivity", + type=float, + default=40.0, + help="Sensitivity factor for rotation.", + ) + algorithm_group.add_argument( + "--disable_control_scheme_inversion", + action="store_true", + default=False, + help="Flag to disable inverting the control scheme due to view for manipulation-based tasks.", + ) + add_default_cli_args(parser) + return parser.parse_args() + + +if __name__ == "__main__": + # Parse arguments + args = parse_cli_args() + + # Launch the app + launcher = launch_app(args) + + # Run the main function + main(launcher=launcher, args=args) + + # Shutdown the app + shutdown_app(launcher) diff --git a/scripts/utils/clean_procgen_cache.py b/scripts/utils/clean_procgen_cache.py new file mode 100755 index 0000000..0c91168 --- /dev/null +++ b/scripts/utils/clean_procgen_cache.py @@ -0,0 +1,29 @@ +#!/root/isaac-sim/python.sh +""" +Utility script for removing the cache of procedural assets generated by the Space Robotics Bench + +Usage: + ros2 run space_robotics_bench clean_procgen_cache.py +""" + +import shutil + +import platformdirs + + +def main(): + cache_dir = platformdirs.user_cache_dir("srb") + + response = input( + f"Do you want to remove the cache of space_robotics_bench located at '{cache_dir}'? (y/N): " + ) + + if response.lower() in ["y", "yes"]: + shutil.rmtree(cache_dir, ignore_errors=True) + print("Cache removed") + else: + print("Exiting without removing the cache") + + +if __name__ == "__main__": + main() diff --git a/scripts/utils/tensorboard.bash b/scripts/utils/tensorboard.bash new file mode 100755 index 0000000..7984b6a --- /dev/null +++ b/scripts/utils/tensorboard.bash @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +### Start a TensorBoard server with the default logdir relative to the repository +### Usage: tensorboard.bash +set -e + +SCRIPT_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" &>/dev/null && pwd)" +REPOSITORY_DIR="$(dirname "$(dirname "${SCRIPT_DIR}")")" + +## If desired, use the last edited log directory +USE_LAST="${USE_LAST:-"false"}" +if [ "${USE_LAST}" == "true" ]; then + LOGDIR="$(find "${LOGDIR}" -type d -exec ls -dt {} + | head -n 1)" +fi + +## Start TensorBoard +TENSORBOARD_CMD=( + "${ISAAC_SIM_PYTHON:-"python3"}" "${ISAAC_SIM_PATH:-"/root/isaac-sim"}/tensorboard" + --logdir "${LOGDIR:-"${REPOSITORY_DIR}/logs"}" + --bind_all +) +echo -e "\033[1;90m[TRACE] ${TENSORBOARD_CMD[*]}\033[0m" | xargs +# shellcheck disable=SC2048 +exec ${TENSORBOARD_CMD[*]} diff --git a/scripts/utils/update_assets.bash b/scripts/utils/update_assets.bash new file mode 100755 index 0000000..efc63ab --- /dev/null +++ b/scripts/utils/update_assets.bash @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +### Update assets by updating git submodules +### Usage: update_assets.bash +set -e + +SCRIPT_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" &>/dev/null && pwd)" +REPOSITORY_DIR="$(dirname "$(dirname "${SCRIPT_DIR}")")" +ASSETS_DIR="${REPOSITORY_DIR}/assets" + + +## Update git submodules +GIT_SUBMODULE_UPDATE_CMD=( + git + -C "${REPOSITORY_DIR}" + submodule + update + --remote + --recursive + "${ASSETS_DIR}" +) +echo -e "\033[1;90m[TRACE] ${GIT_SUBMODULE_UPDATE_CMD[*]}\033[0m" | xargs +# shellcheck disable=SC2048 +exec ${GIT_SUBMODULE_UPDATE_CMD[*]} diff --git a/scripts/zero_agent.py b/scripts/zero_agent.py new file mode 100755 index 0000000..a18d7bf --- /dev/null +++ b/scripts/zero_agent.py @@ -0,0 +1,99 @@ +#!/root/isaac-sim/python.sh +""" +Testing script for running tasks with zero-valued actions + +Examples: + ros2 run space_robotics_bench zero_agent.py + ros2 run space_robotics_bench zero_agent.py --task sample_collection --num_envs 4 +""" + +from _cli_utils import add_default_cli_args, argparse, launch_app, shutdown_app +from omni.isaac.lab.app import AppLauncher + + +def main(launcher: AppLauncher, args: argparse.Namespace): + ## Note: Importing modules here due to delayed Omniverse Kit extension loading + from os import path + + import gymnasium + import torch + from omni.isaac.kit import SimulationApp + from omni.isaac.lab.utils.dict import print_dict + + import space_robotics_bench # noqa: F401 + from space_robotics_bench.utils.parsing import create_logdir_path, parse_task_cfg + + ## Extract simulation app + sim_app: SimulationApp = launcher.app + + ## Configuration + task_cfg = parse_task_cfg( + task_name=args.task, + device=args.device, + num_envs=args.num_envs, + use_fabric=not args.disable_fabric, + ) + + ## Create the environment + env = gymnasium.make( + id=args.task, cfg=task_cfg, render_mode="rgb_array" if args.video else None + ) + + ## Initialize the environment + observation, info = env.reset() + + ## Add wrapper for video recording (if enabled) + if args.video: + logdir = create_logdir_path("zero", args.task) + video_kwargs = { + "video_folder": path.join(logdir, "videos"), + "step_trigger": lambda step: step % args.video_interval == 0, + "video_length": args.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gymnasium.wrappers.RecordVideo(env, **video_kwargs) + + ## Run the environment with random actions + with torch.inference_mode(): + while sim_app.is_running(): + # Construct zero-valued actions + actions = torch.zeros(env.action_space.shape, device=env.unwrapped.device) + + # Step the environment + observation, reward, terminated, truncated, info = env.step(actions) + print( + f"> actions: {actions}\n" + f"> observation: {observation}\n" + f"> reward: {reward}\n" + f"> terminated: {terminated}\n" + f"> truncated: {truncated}\n" + f"> info: {info}\n" + ) + + # Note: Each environment is automatically reset (independently) when terminated or truncated + + ## Close the environment + env.close() + + +### Helper functions ### +def parse_cli_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + add_default_cli_args(parser) + return parser.parse_args() + + +if __name__ == "__main__": + # Parse arguments + args = parse_cli_args() + + # Launch the app + launcher = launch_app(args) + + # Run the main function + main(launcher=launcher, args=args) + + # Shutdown the app + shutdown_app(launcher) diff --git a/space_robotics_bench/__init__.py b/space_robotics_bench/__init__.py new file mode 100644 index 0000000..6992800 --- /dev/null +++ b/space_robotics_bench/__init__.py @@ -0,0 +1,39 @@ +from os import environ + +from .utils.importer import import_modules_recursively +from .utils.sim_app import is_sim_app_started +from .utils.traceback import try_enable_rich_traceback + +## Try enabling rich traceback +try_enable_rich_traceback() + +## Verify the availability of the Rust extension module +try: + from . import _rs # noqa: F401 +except Exception: + raise ModuleNotFoundError( + "Failed to import Python submodule 'space_robotics_bench._rs' that contains the Rust " + "extension module. Please ensure that the package has been installed correctly." + ) + +## If the simulation app is started, register all tasks by recursively importing +## the "{__name__}.tasks" submodule +if environ.get("SRB_SKIP_REGISTRATION", "false").lower() in ["true", "1"]: + print( + f"INFO: [SRB_SKIP_REGISTRATION={environ.get('SRB_SKIP_REGISTRATION')}] Skipping " + "the registration of the Space Robotics Bench tasks." + ) +elif is_sim_app_started(): + import_modules_recursively( + module_name=f"{__name__}.tasks", + ignorelist=[ + "common", + ], + ) +else: + raise RuntimeError( + "Tasks of the Space Robotics Bench cannot be registered because the simulation " + "is not running. Please import the 'space_robotics_bench' module after starting the " + "Omniverse simulation app. Alternatively, set the 'SRB_SKIP_REGISTRATION' environment " + "variable to 'true' to skip the registration of tasks." + ) diff --git a/space_robotics_bench/_rs/__init__.pyi b/space_robotics_bench/_rs/__init__.pyi new file mode 100644 index 0000000..0c9a35e --- /dev/null +++ b/space_robotics_bench/_rs/__init__.pyi @@ -0,0 +1 @@ +from . import envs, utils # noqa: F401 diff --git a/space_robotics_bench/_rs/envs/__init__.pyi b/space_robotics_bench/_rs/envs/__init__.pyi new file mode 100644 index 0000000..54fba44 --- /dev/null +++ b/space_robotics_bench/_rs/envs/__init__.pyi @@ -0,0 +1,88 @@ +from typing import Optional, Tuple + +class EnvironmentConfig: + def __init__( + self, + scenario: Scenario, + assets: Assets, + seed: int, + detail: float, + ): ... + @staticmethod + def extract( + cls, + cfg_path: Optional[str] = None, + env_prefix: Optional[str] = "SRB_", + other: Optional[EnvironmentConfig] = None, + ) -> EnvironmentConfig: ... + def write(self, path: str): ... + @property + def scenario(self) -> Scenario: ... + @property + def assets(self) -> Assets: ... + @property + def seed(self) -> int: ... + @property + def detail(self) -> float: ... + +class Assets: + def __init__( + self, + robot: Asset, + object: Asset, + terrain: Asset, + vehicle: Asset, + ): ... + @property + def robot(self) -> Asset: ... + @property + def object(self) -> Asset: ... + @property + def terrain(self) -> Asset: ... + @property + def vehicle(self) -> Asset: ... + +class Asset: + def __init__( + self, + variant: AssetVariant, + ): ... + @property + def variant(self) -> AssetVariant: ... + +class AssetVariant: + NONE: int = ... + PRIMITIVE: int = ... + DATASET: int = ... + PROCEDURAL: int = ... + +class Scenario: + ASTEROID: int = ... + EARTH: int = ... + MARS: int = ... + MOON: int = ... + ORBIT: int = ... + @property + def gravity_magnitude(self) -> float: ... + @property + def gravity_variation(self) -> float: ... + @property + def gravity_range(self) -> Tuple[float, float]: ... + @property + def light_intensity(self) -> float: ... + @property + def light_intensity_variation(self) -> float: ... + @property + def light_intensity_range(self) -> Tuple[float, float]: ... + @property + def light_angular_diameter(self) -> float: ... + @property + def light_angular_diameter_variation(self) -> float: ... + @property + def light_angular_diameter_range(self) -> Tuple[float, float]: ... + @property + def light_color_temperature(self) -> float: ... + @property + def light_color_temperature_variation(self) -> float: ... + @property + def light_color_temperature_range(self) -> Tuple[float, float]: ... diff --git a/space_robotics_bench/_rs/utils/__init__.pyi b/space_robotics_bench/_rs/utils/__init__.pyi new file mode 100644 index 0000000..c644377 --- /dev/null +++ b/space_robotics_bench/_rs/utils/__init__.pyi @@ -0,0 +1 @@ +from . import sampling # noqa: F401 diff --git a/space_robotics_bench/_rs/utils/sampling/__init__.pyi b/space_robotics_bench/_rs/utils/sampling/__init__.pyi new file mode 100644 index 0000000..8de3c9b --- /dev/null +++ b/space_robotics_bench/_rs/utils/sampling/__init__.pyi @@ -0,0 +1,22 @@ +from typing import List, Tuple + +def sample_poisson_disk_2d( + num_samples: int, + bounds: Tuple[Tuple[float, float], Tuple[float, float]], + radius: float, +) -> List[Tuple[float, float]]: ... +def sample_poisson_disk_2d_looped( + num_samples: Tuple[int, int], + bounds: Tuple[Tuple[float, float], Tuple[float, float]], + radius: float, +) -> List[List[Tuple[float, float]]]: ... +def sample_poisson_disk_3d( + num_samples: int, + bounds: Tuple[Tuple[float, float, float], Tuple[float, float, float]], + radius: float, +) -> List[Tuple[float, float, float]]: ... +def sample_poisson_disk_3d_looped( + num_samples: Tuple[int, int], + bounds: Tuple[Tuple[float, float, float], Tuple[float, float, float]], + radius: float, +) -> List[List[Tuple[float, float, float]]]: ... diff --git a/space_robotics_bench/assets/__init__.py b/space_robotics_bench/assets/__init__.py new file mode 100644 index 0000000..da9feb7 --- /dev/null +++ b/space_robotics_bench/assets/__init__.py @@ -0,0 +1,5 @@ +from .light import * # noqa: F403 +from .object import * # noqa: F403 +from .robot import * # noqa: F403 +from .terrain import * # noqa: F403 +from .vehicle import * # noqa: F403 diff --git a/space_robotics_bench/assets/light/__init__.py b/space_robotics_bench/assets/light/__init__.py new file mode 100644 index 0000000..4c0b205 --- /dev/null +++ b/space_robotics_bench/assets/light/__init__.py @@ -0,0 +1,2 @@ +from .sky import * # noqa: F403 +from .sunlight import * # noqa: F403 diff --git a/space_robotics_bench/assets/light/sky.py b/space_robotics_bench/assets/light/sky.py new file mode 100644 index 0000000..472d332 --- /dev/null +++ b/space_robotics_bench/assets/light/sky.py @@ -0,0 +1,41 @@ +from os import path +from typing import Any, Dict, Optional + +from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR + +import space_robotics_bench.core.envs as env_utils +import space_robotics_bench.core.sim as sim_utils +from space_robotics_bench.core.assets import AssetBaseCfg +from space_robotics_bench.paths import SRB_ASSETS_DIR_SRB_HDRI + + +def sky_from_env_cfg( + env_cfg: env_utils.EnvironmentConfig, + *, + prim_path: str = "/World/sky", + spawn_kwargs: Dict[str, Any] = {}, + **kwargs, +) -> Optional[AssetBaseCfg]: + texture_file: Optional[str] = None + + match env_cfg.scenario: + case env_utils.Scenario.EARTH: + texture_file = ( + f"{ISAAC_NUCLEUS_DIR}/Materials/Textures/Skies/PolyHaven/kloofendal_43d_clear_puresky_4k.hdr", + ) + case env_utils.Scenario.MARS: + texture_file = path.join(SRB_ASSETS_DIR_SRB_HDRI, "martian_sky_day.hdr") + case env_utils.Scenario.ORBIT: + texture_file = path.join(SRB_ASSETS_DIR_SRB_HDRI, "low_lunar_orbit.jpg") + + if texture_file is None: + return None + return AssetBaseCfg( + prim_path=prim_path, + spawn=sim_utils.DomeLightCfg( + intensity=0.25 * env_cfg.scenario.light_intensity, + texture_file=texture_file, + **spawn_kwargs, + ), + **kwargs, + ) diff --git a/space_robotics_bench/assets/light/sunlight.py b/space_robotics_bench/assets/light/sunlight.py new file mode 100644 index 0000000..4a2eb74 --- /dev/null +++ b/space_robotics_bench/assets/light/sunlight.py @@ -0,0 +1,25 @@ +from typing import Any, Dict + +import space_robotics_bench.core.envs as env_utils +import space_robotics_bench.core.sim as sim_utils +from space_robotics_bench.core.assets import AssetBaseCfg + + +def sunlight_from_env_cfg( + env_cfg: env_utils.EnvironmentConfig, + *, + prim_path: str = "/World/light", + spawn_kwargs: Dict[str, Any] = {}, + **kwargs, +) -> AssetBaseCfg: + return AssetBaseCfg( + prim_path=prim_path, + spawn=sim_utils.DistantLightCfg( + intensity=env_cfg.scenario.light_intensity, + angle=env_cfg.scenario.light_angular_diameter, + color_temperature=env_cfg.scenario.light_color_temperature, + enable_color_temperature=True, + **spawn_kwargs, + ), + **kwargs, + ) diff --git a/space_robotics_bench/assets/object/__init__.py b/space_robotics_bench/assets/object/__init__.py new file mode 100644 index 0000000..8903d87 --- /dev/null +++ b/space_robotics_bench/assets/object/__init__.py @@ -0,0 +1,259 @@ +from typing import Any, Dict, Optional, Tuple + +import space_robotics_bench.core.envs as env_utils +import space_robotics_bench.core.sim as sim_utils +from space_robotics_bench.core.assets import AssetBaseCfg, RigidObjectCfg +from space_robotics_bench.utils import color as color_utils + +from .lunar_rock_procgen import LunarRockProcgenCfg +from .martian_rock_procgen import MartianRockProcgenCfg +from .peg_in_hole_procgen import HoleProcgenCfg, PegProcgenCfg +from .peg_in_hole_profile import HoleProfileCfg, PegProfileCfg, PegProfileShortCfg +from .sample_tube import SampleTubeCfg +from .solar_panel import SolarPanelCfg + + +@staticmethod +def object_of_interest_from_env_cfg( + env_cfg: env_utils.EnvironmentConfig, + *, + prim_path: str = "{ENV_REGEX_NS}/object", + num_assets: int = 1, + size: Tuple[float, float] = (0.06, 0.06, 0.04), + spawn_kwargs: Dict[str, Any] = {}, + procgen_seed_offset: int = 0, + procgen_kwargs: Dict[str, Any] = {}, + **kwargs, +) -> Optional[RigidObjectCfg]: + spawn: Optional[sim_utils.SpawnerCfg] = None + + if spawn_kwargs.get("collision_props", None) is None: + spawn_kwargs["collision_props"] = sim_utils.CollisionPropertiesCfg() + if spawn_kwargs.get("rigid_props", None) is None: + spawn_kwargs["rigid_props"] = sim_utils.RigidBodyPropertiesCfg() + if spawn_kwargs.get("mass_props", None) is None: + spawn_kwargs["mass_props"] = sim_utils.MassPropertiesCfg(density=2000.0) + + match env_cfg.assets.object.variant: + case env_utils.AssetVariant.PRIMITIVE: + if spawn_kwargs.get("visual_material", None) is None: + spawn_kwargs["visual_material"] = ( + color_utils.preview_surface_from_env_cfg(env_cfg) + ) + + spawn = sim_utils.MultiShapeCfg( + size=size, + shape_cfg=sim_utils.ShapeCfg(**spawn_kwargs), + ) + + case env_utils.AssetVariant.DATASET: + match env_cfg.scenario: + case env_utils.Scenario.MARS: + if spawn_kwargs.get("mesh_collision_props", None) is None: + spawn_kwargs["mesh_collision_props"] = ( + sim_utils.MeshCollisionPropertiesCfg( + mesh_approximation="sdf", + ) + ) + spawn = SampleTubeCfg(**spawn_kwargs) + case _: + if spawn_kwargs.get("mesh_collision_props", None) is None: + spawn_kwargs["mesh_collision_props"] = ( + sim_utils.MeshCollisionPropertiesCfg( + mesh_approximation="boundingCube" + ) + ) + if spawn_kwargs.get("visual_material", None) is None: + spawn_kwargs["visual_material"] = ( + color_utils.preview_surface_from_env_cfg(env_cfg) + ) + + spawn = sim_utils.MultiAssetCfg( + assets_cfg=[ + PegProfileCfg(**spawn_kwargs), + PegProfileShortCfg(**spawn_kwargs), + ] + ) + + case env_utils.AssetVariant.PROCEDURAL: + if spawn_kwargs.get("mesh_collision_props", None) is None: + spawn_kwargs["mesh_collision_props"] = ( + sim_utils.MeshCollisionPropertiesCfg( + mesh_approximation="sdf", + ) + ) + usd_file_cfg = sim_utils.UsdFileCfg( + usd_path="IGNORED", + **spawn_kwargs, + ) + + match env_cfg.scenario: + case env_utils.Scenario.MOON | env_utils.Scenario.ORBIT: + spawn = LunarRockProcgenCfg( + num_assets=num_assets, + usd_file_cfg=usd_file_cfg, + seed=env_cfg.seed + procgen_seed_offset, + detail=env_cfg.detail, + ) + + case env_utils.Scenario.MARS: + spawn = MartianRockProcgenCfg( + num_assets=num_assets, + usd_file_cfg=usd_file_cfg, + seed=env_cfg.seed + procgen_seed_offset, + detail=env_cfg.detail, + ) + case _: + return None + + for node_cfg in spawn.geometry_nodes.values(): + if "scale" in node_cfg: + node_cfg["scale"] = size + + for key, value in procgen_kwargs.items(): + for node_cfg in spawn.geometry_nodes.values(): + if key in node_cfg: + node_cfg[key] = value + elif hasattr(spawn, key): + setattr(spawn, key, value) + + if spawn is None: + raise NotImplementedError + return RigidObjectCfg( + prim_path=prim_path, + spawn=spawn, + **kwargs, + ) + + +@staticmethod +def peg_in_hole_from_env_cfg( + env_cfg: env_utils.EnvironmentConfig, + *, + prim_path_peg: str = "{ENV_REGEX_NS}/peg", + prim_path_hole: str = "{ENV_REGEX_NS}/hole", + num_assets: int = 1, + size: Tuple[float, float, float] = (0.05, 0.05, 0.05), + spawn_kwargs_peg: Dict[str, Any] = {}, + spawn_kwargs_hole: Dict[str, Any] = {}, + procgen_seed_offset: int = 0, + procgen_kwargs_peg: Dict[str, Any] = {}, + procgen_kwargs_hole: Dict[str, Any] = {}, + short_peg: bool = False, + **kwargs, +) -> Optional[Tuple[RigidObjectCfg, AssetBaseCfg]]: + spawn_peg: Optional[sim_utils.SpawnerCfg] = None + spawn_hole: Optional[sim_utils.SpawnerCfg] = None + + if spawn_kwargs_peg.get("collision_props", None) is None: + spawn_kwargs_peg["collision_props"] = sim_utils.CollisionPropertiesCfg() + if spawn_kwargs_hole.get("collision_props", None) is None: + spawn_kwargs_hole["collision_props"] = sim_utils.CollisionPropertiesCfg() + if spawn_kwargs_peg.get("rigid_props", None) is None: + spawn_kwargs_peg["rigid_props"] = sim_utils.RigidBodyPropertiesCfg() + if spawn_kwargs_peg.get("mass_props", None) is None: + spawn_kwargs_peg["mass_props"] = sim_utils.MassPropertiesCfg(density=2000.0) + + match env_cfg.assets.object.variant: + case env_utils.AssetVariant.DATASET: + if spawn_kwargs_peg.get("visual_material", None) is None: + spawn_kwargs_peg["visual_material"] = ( + color_utils.preview_surface_from_env_cfg(env_cfg) + ) + if spawn_kwargs_hole.get("visual_material", None) is None: + spawn_kwargs_hole["visual_material"] = ( + color_utils.preview_surface_from_env_cfg(env_cfg) + ) + if spawn_kwargs_peg.get("mesh_collision_props", None) is None: + spawn_kwargs_peg["mesh_collision_props"] = ( + sim_utils.MeshCollisionPropertiesCfg( + mesh_approximation="boundingCube", + ) + ) + spawn_peg = ( + PegProfileShortCfg(**spawn_kwargs_peg) + if short_peg + else PegProfileCfg(**spawn_kwargs_peg) + ) + spawn_hole = HoleProfileCfg(**spawn_kwargs_hole) + case env_utils.AssetVariant.PROCEDURAL: + if spawn_kwargs_peg.get("mesh_collision_props", None) is None: + spawn_kwargs_peg["mesh_collision_props"] = ( + sim_utils.MeshCollisionPropertiesCfg( + mesh_approximation="sdf", + ) + ) + spawn_peg = PegProcgenCfg( + num_assets=num_assets, + usd_file_cfg=sim_utils.UsdFileCfg( + usd_path="IGNORED", + **spawn_kwargs_peg, + ), + seed=env_cfg.seed + procgen_seed_offset, + detail=env_cfg.detail, + ) + spawn_hole = HoleProcgenCfg( + num_assets=num_assets, + usd_file_cfg=sim_utils.UsdFileCfg( + usd_path="IGNORED", + **spawn_kwargs_hole, + ), + seed=env_cfg.seed + procgen_seed_offset, + detail=env_cfg.detail, + ) + + for spawn, procgen_kwargs in ( + (spawn_peg, procgen_kwargs_peg), + (spawn_hole, procgen_kwargs_hole), + ): + for node_cfg in spawn.geometry_nodes.values(): + if "scale" in node_cfg: + node_cfg["scale"] = size + + for key, value in procgen_kwargs.items(): + for node_cfg in spawn.geometry_nodes.values(): + if key in node_cfg: + node_cfg[key] = value + elif hasattr(spawn, key): + setattr(spawn, key, value) + + return RigidObjectCfg( + prim_path=prim_path_peg, + spawn=spawn_peg, + **kwargs, + ), AssetBaseCfg( + prim_path=prim_path_hole, + spawn=spawn_hole, + **kwargs, + ) + + +@staticmethod +def solar_panel_from_env_cfg( + env_cfg: env_utils.EnvironmentConfig, + *, + prim_path: str = "{ENV_REGEX_NS}/panel", + spawn_kwargs: Dict[str, Any] = {}, + **kwargs, +) -> Optional[RigidObjectCfg]: + spawn: Optional[sim_utils.SpawnerCfg] = None + + if spawn_kwargs.get("collision_props", None) is None: + spawn_kwargs["collision_props"] = sim_utils.CollisionPropertiesCfg() + if spawn_kwargs.get("rigid_props", None) is None: + spawn_kwargs["rigid_props"] = sim_utils.RigidBodyPropertiesCfg() + if spawn_kwargs.get("mass_props", None) is None: + spawn_kwargs["mass_props"] = sim_utils.MassPropertiesCfg(density=1000.0) + + if spawn_kwargs.get("mesh_collision_props", None) is None: + spawn_kwargs["mesh_collision_props"] = sim_utils.MeshCollisionPropertiesCfg( + mesh_approximation="sdf", + ) + + spawn = SolarPanelCfg(**spawn_kwargs) + + return RigidObjectCfg( + prim_path=prim_path, + spawn=spawn, + **kwargs, + ) diff --git a/space_robotics_bench/assets/object/lunar_rock_procgen.py b/space_robotics_bench/assets/object/lunar_rock_procgen.py new file mode 100644 index 0000000..7dfc4b5 --- /dev/null +++ b/space_robotics_bench/assets/object/lunar_rock_procgen.py @@ -0,0 +1,33 @@ +from os import path +from typing import Any, Dict, List, Optional + +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.sim as sim_utils +from space_robotics_bench.paths import SRB_ASSETS_DIR_SRB_OBJECT +from space_robotics_bench.utils.path import abs_listdir + + +@configclass +class LunarRockProcgenCfg(sim_utils.BlenderNodesAssetCfg): + name: str = "lunar_rock" + autorun_scripts: List[str] = abs_listdir( + path.join(SRB_ASSETS_DIR_SRB_OBJECT, "lunar_rock_procgen") + ) + + # Geometry + geometry_nodes: Dict[str, Dict[str, Any]] = { + "LunarRock": { + "detail": 5, # Level of the subdivision (resolution of the mesh) + "scale": [0.1, 0.1, 0.075], # Metric scale of the mesh + "scale_std": [0.01, 0.01, 0.005], # Standard deviation of the scale + "horizontal_cut": False, # Flag to enable horizontal cut + "horizontal_cut_offset": 0.0, # Offset of the horizontal cut with respect to the mesh center + } + } + decimate_face_count: Optional[int] = None + decimate_angle_limit: Optional[float] = None + + # Material + material: Optional[str] = "LunarRock" + texture_resolution: int = 1024 diff --git a/space_robotics_bench/assets/object/martian_rock_procgen.py b/space_robotics_bench/assets/object/martian_rock_procgen.py new file mode 100644 index 0000000..40a134d --- /dev/null +++ b/space_robotics_bench/assets/object/martian_rock_procgen.py @@ -0,0 +1,33 @@ +from os import path +from typing import Any, Dict, List, Optional + +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.sim as sim_utils +from space_robotics_bench.paths import SRB_ASSETS_DIR_SRB_OBJECT +from space_robotics_bench.utils.path import abs_listdir + + +@configclass +class MartianRockProcgenCfg(sim_utils.BlenderNodesAssetCfg): + name: str = "martian_rock" + autorun_scripts: List[str] = abs_listdir( + path.join(SRB_ASSETS_DIR_SRB_OBJECT, "martian_rock_procgen") + ) + + # Geometry + geometry_nodes: Dict[str, Dict[str, Any]] = { + "MartianRock": { + "detail": 4, # Level of the subdivision (resolution of the mesh) + "scale": [0.1, 0.1, 0.075], # Metric scale of the mesh + "scale_std": [0.01, 0.01, 0.005], # Standard deviation of the scale + "horizontal_cut": False, # Flag to enable horizontal cut + "horizontal_cut_offset": 0.0, # Offset of the horizontal cut with respect to the mesh center + } + } + decimate_face_count: Optional[int] = None + decimate_angle_limit: Optional[float] = None + + # Material + material: Optional[str] = "MartianRock" + texture_resolution: int = 1024 diff --git a/space_robotics_bench/assets/object/peg_in_hole_procgen.py b/space_robotics_bench/assets/object/peg_in_hole_procgen.py new file mode 100644 index 0000000..d8856e6 --- /dev/null +++ b/space_robotics_bench/assets/object/peg_in_hole_procgen.py @@ -0,0 +1,44 @@ +from os import path +from typing import Any, Dict, List, Optional + +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.sim as sim_utils +from space_robotics_bench.paths import SRB_ASSETS_DIR_SRB_OBJECT +from space_robotics_bench.utils.path import abs_listdir + + +@configclass +class PegProcgenCfg(sim_utils.BlenderNodesAssetCfg): + name: str = "peg" + autorun_scripts: List[str] = abs_listdir( + path.join(SRB_ASSETS_DIR_SRB_OBJECT, "peg_in_hole_procgen") + ) + + # Geometry + geometry_nodes: Dict[str, Dict[str, Any]] = {"Peg": {}} + decimate_face_count: Optional[int] = None + decimate_angle_limit: Optional[float] = None + + # Material + material: Optional[str] = "Metal" + texture_resolution: int = 512 + + +@configclass +class HoleProcgenCfg(sim_utils.BlenderNodesAssetCfg): + name: str = "hole" + autorun_scripts: List[str] = abs_listdir( + path.join(SRB_ASSETS_DIR_SRB_OBJECT, "peg_in_hole_procgen") + ) + + export_kwargs: Dict[str, Any] = {"triangulate_meshes": True} + + # Geometry + geometry_nodes: Dict[str, Dict[str, Any]] = {"BlankModule": {}, "Hole": {}} + decimate_face_count: Optional[int] = None + decimate_angle_limit: Optional[float] = None + + # Material + material: Optional[str] = "Metal" + texture_resolution: int = 1024 diff --git a/space_robotics_bench/assets/object/peg_in_hole_profile.py b/space_robotics_bench/assets/object/peg_in_hole_profile.py new file mode 100644 index 0000000..95f584f --- /dev/null +++ b/space_robotics_bench/assets/object/peg_in_hole_profile.py @@ -0,0 +1,25 @@ +from os import path + +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.sim as sim_utils +from space_robotics_bench.paths import SRB_ASSETS_DIR_SRB_OBJECT + + +@configclass +class PegProfileCfg(sim_utils.UsdFileCfg): + usd_path = path.join( + SRB_ASSETS_DIR_SRB_OBJECT, "peg_in_hole_profile", "profile.usdc" + ) + + +@configclass +class PegProfileShortCfg(sim_utils.UsdFileCfg): + usd_path = path.join( + SRB_ASSETS_DIR_SRB_OBJECT, "peg_in_hole_profile", "profile_short.usdc" + ) + + +@configclass +class HoleProfileCfg(sim_utils.UsdFileCfg): + usd_path = path.join(SRB_ASSETS_DIR_SRB_OBJECT, "peg_in_hole_profile", "hole.usdc") diff --git a/space_robotics_bench/assets/object/sample_tube.py b/space_robotics_bench/assets/object/sample_tube.py new file mode 100644 index 0000000..6fe35a0 --- /dev/null +++ b/space_robotics_bench/assets/object/sample_tube.py @@ -0,0 +1,11 @@ +from os import path + +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.sim as sim_utils +from space_robotics_bench.paths import SRB_ASSETS_DIR_SRB_OBJECT + + +@configclass +class SampleTubeCfg(sim_utils.UsdFileCfg): + usd_path = path.join(SRB_ASSETS_DIR_SRB_OBJECT, "sample_tube", "sample_tube.usdc") diff --git a/space_robotics_bench/assets/object/solar_panel.py b/space_robotics_bench/assets/object/solar_panel.py new file mode 100644 index 0000000..c3999f3 --- /dev/null +++ b/space_robotics_bench/assets/object/solar_panel.py @@ -0,0 +1,11 @@ +from os import path + +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.sim as sim_utils +from space_robotics_bench.paths import SRB_ASSETS_DIR_SRB_OBJECT + + +@configclass +class SolarPanelCfg(sim_utils.UsdFileCfg): + usd_path = path.join(SRB_ASSETS_DIR_SRB_OBJECT, "solar_panel", "solar_panel.usdc") diff --git a/space_robotics_bench/assets/robot/__init__.py b/space_robotics_bench/assets/robot/__init__.py new file mode 100644 index 0000000..1bf15a2 --- /dev/null +++ b/space_robotics_bench/assets/robot/__init__.py @@ -0,0 +1,34 @@ +import space_robotics_bench.core.assets as asset_utils +import space_robotics_bench.core.envs as env_utils + +from .canadarm3_large import canadarm3_large_cfg # noqa: F401 +from .franka import franka_cfg +from .ingenuity import ingenuity_cfg +from .perseverance import perseverance_cfg + + +def rover_from_env_cfg( + env_cfg: env_utils.EnvironmentConfig, + *, + prim_path: str = "{ENV_REGEX_NS}/robot", + **kwargs, +) -> asset_utils.MobileRobotCfg: + return perseverance_cfg(prim_path=prim_path, **kwargs) + + +def manipulator_from_env_cfg( + env_cfg: env_utils.EnvironmentConfig, + *, + prim_path: str = "{ENV_REGEX_NS}/robot", + **kwargs, +) -> asset_utils.ManipulatorCfg: + return franka_cfg(prim_path=prim_path, **kwargs) + + +def aerial_robot_from_env_cfg( + env_cfg: env_utils.EnvironmentConfig, + *, + prim_path: str = "{ENV_REGEX_NS}/robot", + **kwargs, +) -> asset_utils.AerialRobotCfg: + return ingenuity_cfg(prim_path=prim_path, **kwargs) diff --git a/space_robotics_bench/assets/robot/canadarm3_large.py b/space_robotics_bench/assets/robot/canadarm3_large.py new file mode 100644 index 0000000..3f32923 --- /dev/null +++ b/space_robotics_bench/assets/robot/canadarm3_large.py @@ -0,0 +1,125 @@ +from os import path + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.actuators import ImplicitActuatorCfg +from omni.isaac.lab.assets.articulation import ArticulationCfg +from omni.isaac.lab.controllers import DifferentialIKControllerCfg +from omni.isaac.lab.envs.mdp import DifferentialInverseKinematicsActionCfg +from torch import pi + +import space_robotics_bench.core.assets as asset_utils +import space_robotics_bench.utils.math as math_utils +from space_robotics_bench.core.actions import ManipulatorTaskSpaceActionCfg +from space_robotics_bench.core.envs import mdp +from space_robotics_bench.paths import SRB_ASSETS_DIR_SRB_ROBOT + + +def canadarm3_large_cfg( + *, + prim_path: str = "{ENV_REGEX_NS}/robot", + asset_name: str = "robot", + use_relative_mode: bool = True, + action_scale: float = 0.1, + **kwargs, +) -> asset_utils.ManipulatorCfg: + frame_base = "canadarm3_large_0" + frame_ee = "canadarm3_large_7" + regex_joints_arm = "canadarm3_large_joint_[1-7]" + regex_joints_hand = "canadarm3_large_joint_7" + + return asset_utils.ManipulatorCfg( + ## Model + asset_cfg=ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=path.join( + SRB_ASSETS_DIR_SRB_ROBOT, + "canadarm3_large", + "canadarm3_large.usdc", + ), + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=True, + max_depenetration_velocity=5.0, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=False, + solver_position_iteration_count=12, + solver_velocity_iteration_count=1, + ), + activate_contact_sensors=True, + ), + init_state=ArticulationCfg.InitialStateCfg( + joint_pos={ + "canadarm3_large_joint_1": 50.0 * pi / 180.0, + "canadarm3_large_joint_2": 0.0 * pi / 180.0, + "canadarm3_large_joint_3": 55.0 * pi / 180.0, + "canadarm3_large_joint_4": 75.0 * pi / 180.0, + "canadarm3_large_joint_5": -30.0 * pi / 180.0, + "canadarm3_large_joint_6": 0.0 * pi / 180.0, + "canadarm3_large_joint_7": 0.0 * pi / 180.0, + }, + ), + actuators={ + "joints": ImplicitActuatorCfg( + joint_names_expr=["canadarm3_large_joint_[1-7]"], + effort_limit=2500.0, + velocity_limit=5.0, + stiffness=40000.0, + damping=25000.0, + ), + }, + soft_joint_pos_limit_factor=1.0, + ).replace(prim_path=prim_path, **kwargs), + ## Actions + action_cfg=ManipulatorTaskSpaceActionCfg( + arm=mdp.DifferentialInverseKinematicsActionCfg( + asset_name=asset_name, + joint_names=[regex_joints_arm], + body_name=frame_ee, + controller=DifferentialIKControllerCfg( + command_type="pose", + use_relative_mode=use_relative_mode, + ik_method="dls", + ), + scale=action_scale, + body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg( + pos=(0.0, 0.0, 0.0) + ), + ), + hand=mdp.BinaryJointPositionActionCfg( + asset_name=asset_name, + joint_names=[regex_joints_hand], + close_command_expr={regex_joints_hand: 0.0}, + open_command_expr={regex_joints_hand: 0.0}, + ), + ), + ## Frames + frame_base=asset_utils.FrameCfg( + prim_relpath=frame_base, + ), + frame_ee=asset_utils.FrameCfg( + prim_relpath=frame_ee, + offset=asset_utils.TransformCfg( + translation=(0.0, 0.0, 0.1034), + ), + ), + frame_camera_base=asset_utils.FrameCfg( + prim_relpath=f"{frame_base}/camera_base", + offset=asset_utils.TransformCfg( + translation=(0.06, 0.0, 0.15), + rotation=math_utils.quat_from_rpy(0.0, -10.0, 0.0), + ), + ), + frame_camera_wrist=asset_utils.FrameCfg( + prim_relpath=f"{frame_ee}/camera_wrist", + offset=asset_utils.TransformCfg( + translation=(0.0, 0.0, -0.45), + rotation=math_utils.quat_from_rpy(0.0, 90.0, 180.0), + ), + ), + ## Links + regex_links_arm="canadarm3_large_[0-6]", + regex_links_hand="canadarm3_large_7", + ## Joints + regex_joints_arm=regex_joints_arm, + regex_joints_hand=regex_joints_hand, + ) diff --git a/space_robotics_bench/assets/robot/franka.py b/space_robotics_bench/assets/robot/franka.py new file mode 100644 index 0000000..07c03c4 --- /dev/null +++ b/space_robotics_bench/assets/robot/franka.py @@ -0,0 +1,141 @@ +from os import path + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.actuators import ImplicitActuatorCfg +from omni.isaac.lab.assets.articulation import ArticulationCfg +from omni.isaac.lab.controllers import DifferentialIKControllerCfg +from omni.isaac.lab.envs.mdp import DifferentialInverseKinematicsActionCfg +from torch import pi + +import space_robotics_bench.core.assets as asset_utils +import space_robotics_bench.utils.math as math_utils +from space_robotics_bench.core.actions import ManipulatorTaskSpaceActionCfg +from space_robotics_bench.core.envs import mdp +from space_robotics_bench.paths import SRB_ASSETS_DIR_SRB_ROBOT + + +def franka_cfg( + *, + prim_path: str = "{ENV_REGEX_NS}/robot", + asset_name: str = "robot", + use_relative_mode: bool = True, + action_scale: float = 0.1, + **kwargs, +) -> asset_utils.ManipulatorCfg: + frame_base = "panda_link0" + frame_ee = "panda_hand" + regex_joints_arm = "panda_joint.*" + regex_joints_hand = "panda_finger_joint.*" + + return asset_utils.ManipulatorCfg( + ## Model + asset_cfg=ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=path.join( + SRB_ASSETS_DIR_SRB_ROBOT, + "franka", + "panda.usdc", + ), + activate_contact_sensors=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=True, + max_depenetration_velocity=5.0, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=False, + solver_position_iteration_count=12, + solver_velocity_iteration_count=1, + ), + # collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + ), + init_state=ArticulationCfg.InitialStateCfg( + joint_pos={ + "panda_joint1": 0.0, + "panda_joint2": -(pi / 8.0), + "panda_joint3": 0.0, + "panda_joint4": -(pi - (pi / 8.0)), + "panda_joint5": 0.0, + "panda_joint6": pi - (pi / 4.0), + "panda_joint7": (pi / 4.0), + "panda_finger_joint.*": 0.04, + }, + ), + actuators={ + "panda_shoulder": ImplicitActuatorCfg( + joint_names_expr=["panda_joint[1-4]"], + effort_limit=87.0, + velocity_limit=2.175, + stiffness=4000.0, + damping=800.0, + ), + "panda_forearm": ImplicitActuatorCfg( + joint_names_expr=["panda_joint[5-7]"], + effort_limit=12.0, + velocity_limit=2.61, + stiffness=4000.0, + damping=800.0, + ), + "panda_hand": ImplicitActuatorCfg( + joint_names_expr=["panda_finger_joint.*"], + effort_limit=200.0, + velocity_limit=0.2, + stiffness=2e3, + damping=1e2, + ), + }, + soft_joint_pos_limit_factor=1.0, + ).replace(prim_path=prim_path, **kwargs), + ## Actions + action_cfg=ManipulatorTaskSpaceActionCfg( + arm=mdp.DifferentialInverseKinematicsActionCfg( + asset_name=asset_name, + joint_names=[regex_joints_arm], + body_name=frame_ee, + controller=DifferentialIKControllerCfg( + command_type="pose", + use_relative_mode=use_relative_mode, + ik_method="dls", + ), + scale=action_scale, + body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg( + pos=(0.0, 0.0, 0.107) + ), + ), + hand=mdp.BinaryJointPositionActionCfg( + asset_name=asset_name, + joint_names=[regex_joints_hand], + close_command_expr={regex_joints_hand: 0.0}, + open_command_expr={regex_joints_hand: 0.04}, + ), + ), + ## Frames + frame_base=asset_utils.FrameCfg( + prim_relpath=frame_base, + ), + frame_ee=asset_utils.FrameCfg( + prim_relpath=frame_ee, + offset=asset_utils.TransformCfg( + translation=(0.0, 0.0, 0.1034), + ), + ), + frame_camera_base=asset_utils.FrameCfg( + prim_relpath=f"{frame_base}/camera_base", + offset=asset_utils.TransformCfg( + translation=(0.06, 0.0, 0.15), + rotation=math_utils.quat_from_rpy(0.0, -10.0, 0.0), + ), + ), + frame_camera_wrist=asset_utils.FrameCfg( + prim_relpath=f"{frame_ee}/camera_wrist", + offset=asset_utils.TransformCfg( + translation=(0.07, 0.0, 0.05), + rotation=math_utils.quat_from_rpy(0.0, -60.0, 180.0), + ), + ), + ## Links + regex_links_arm="panda_link[1-7]", + regex_links_hand="panda_(hand|leftfinger|rightfinger)", + ## Joints + regex_joints_arm=regex_joints_arm, + regex_joints_hand=regex_joints_hand, + ) diff --git a/space_robotics_bench/assets/robot/ingenuity.py b/space_robotics_bench/assets/robot/ingenuity.py new file mode 100644 index 0000000..93b0431 --- /dev/null +++ b/space_robotics_bench/assets/robot/ingenuity.py @@ -0,0 +1,87 @@ +from os import path + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.actuators import ImplicitActuatorCfg +from omni.isaac.lab.assets.articulation import ArticulationCfg +from torch import pi + +import space_robotics_bench.core.assets as asset_utils +import space_robotics_bench.utils.math as math_utils +from space_robotics_bench.core.actions import ( + MultiCopterActionCfg, + MultiCopterActionGroupCfg, +) +from space_robotics_bench.paths import SRB_ASSETS_DIR_SRB_ROBOT + + +def ingenuity_cfg( + *, + prim_path: str = "{ENV_REGEX_NS}/robot", + asset_name: str = "robot", + action_scale: float = 4.0, + **kwargs, +) -> asset_utils.MultiCopterCfg: + frame_base = "body" + regex_links_rotors = "rotor_[1-2]" + regex_joints_rotors = "rotor_joint_[1-2]" + + return asset_utils.MultiCopterCfg( + ## Model + asset_cfg=ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=path.join( + SRB_ASSETS_DIR_SRB_ROBOT, "ingenuity", "ingenuity.usdc" + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=False, + solver_position_iteration_count=8, + solver_velocity_iteration_count=1, + ), + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=True, + max_depenetration_velocity=5.0, + ), + activate_contact_sensors=True, + ), + init_state=ArticulationCfg.InitialStateCfg(), + actuators={ + "rotors": ImplicitActuatorCfg( + joint_names_expr=[regex_joints_rotors], + velocity_limit=2500 / 60 * 2 * pi, # 2500 RPM + effort_limit=7.5, + stiffness=0.0, + damping=1000.0, + ), + }, + soft_joint_pos_limit_factor=0.0, + ).replace(prim_path=prim_path, **kwargs), + ## Actions + action_cfg=MultiCopterActionGroupCfg( + flight=MultiCopterActionCfg( + asset_name=asset_name, + frame_base=frame_base, + regex_joints_rotors=regex_joints_rotors, + nominal_rpm={ + "rotor_joint_1": 2500.0, + "rotor_joint_2": -2500.0, + }, + tilt_magnitude=0.125, + scale=action_scale, + ) + ), + ## Frames + frame_base=asset_utils.FrameCfg( + prim_relpath=frame_base, + ), + frame_camera_bottom=asset_utils.FrameCfg( + prim_relpath=f"{frame_base}/camera_bottom", + offset=asset_utils.TransformCfg( + translation=(0.045, 0.0, 0.1275), + rotation=math_utils.quat_from_rpy(0.0, 90.0, 0.0), + ), + ), + ## Links + regex_links_rotors=regex_links_rotors, + ## Joints + regex_joints_rotors=regex_joints_rotors, + ) diff --git a/space_robotics_bench/assets/robot/perseverance.py b/space_robotics_bench/assets/robot/perseverance.py new file mode 100644 index 0000000..20ce617 --- /dev/null +++ b/space_robotics_bench/assets/robot/perseverance.py @@ -0,0 +1,123 @@ +from os import path + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.actuators import ImplicitActuatorCfg +from omni.isaac.lab.assets.articulation import ArticulationCfg + +import space_robotics_bench.core.assets as asset_utils +import space_robotics_bench.utils.math as math_utils +from space_robotics_bench.core.actions import ( + WheeledRoverActionCfg, + WheeledRoverActionGroupCfg, +) +from space_robotics_bench.paths import SRB_ASSETS_DIR_SRB_ROBOT + + +def perseverance_cfg( + *, + prim_path: str = "{ENV_REGEX_NS}/robot", + asset_name: str = "robot", + action_scale: float = 1.0, + **kwargs, +) -> asset_utils.WheeledRoverCfg: + frame_base = "body" + regex_drive_joints = "drive_joint.*" + regex_steer_joints = "steer_joint.*" + + return asset_utils.WheeledRoverCfg( + ## Model + asset_cfg=ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=path.join( + SRB_ASSETS_DIR_SRB_ROBOT, + "perseverance", + "perseverance.usdc", + ), + activate_contact_sensors=True, + collision_props=sim_utils.CollisionPropertiesCfg( + contact_offset=0.02, rest_offset=0.005 + ), + rigid_props=sim_utils.RigidBodyPropertiesCfg( + max_linear_velocity=1.5, + max_angular_velocity=1000.0, + max_depenetration_velocity=1.0, + disable_gravity=False, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=False, + solver_position_iteration_count=16, + solver_velocity_iteration_count=4, + ), + ), + init_state=ArticulationCfg.InitialStateCfg(), + actuators={ + "drive_joints": ImplicitActuatorCfg( + joint_names_expr=[regex_drive_joints], + velocity_limit=40.0, + effort_limit=150.0, + damping=25000.0, + stiffness=0.0, + ), + "steer_joints": ImplicitActuatorCfg( + joint_names_expr=[regex_steer_joints], + velocity_limit=2.0, + effort_limit=400.0, + damping=200.0, + stiffness=500.0, + ), + "rocker_joints": ImplicitActuatorCfg( + joint_names_expr=["suspension_joint_rocker.*"], + velocity_limit=5.0, + effort_limit=2500.0, + damping=400.0, + stiffness=4000.0, + ), + "bogie_joints": ImplicitActuatorCfg( + joint_names_expr=["suspension_joint_bogie.*"], + velocity_limit=4.0, + effort_limit=500.0, + damping=25.0, + stiffness=200.0, + ), + }, + ).replace(prim_path=prim_path, **kwargs), + ## Actions + action_cfg=WheeledRoverActionGroupCfg( + WheeledRoverActionCfg( + asset_name=asset_name, + wheelbase=(2.26, 2.14764), + wheelbase_mid=2.39164, + wheel_radius=0.26268, + steering_joint_names=[ + "steer_joint_front_left", + "steer_joint_front_right", + "steer_joint_rear_left", + "steer_joint_rear_right", + ], + drive_joint_names=[ + "drive_joint_front_left", + "drive_joint_front_right", + "drive_joint_rear_left", + "drive_joint_rear_right", + "drive_joint_mid_left", + "drive_joint_mid_right", + ], + scale=action_scale, + ) + ), + ## Frames + frame_base=asset_utils.FrameCfg( + prim_relpath=frame_base, + ), + frame_camera_front=asset_utils.FrameCfg( + prim_relpath=f"{frame_base}/camera_front", + offset=asset_utils.TransformCfg( + # translation=(-0.3437, -0.8537, 1.9793), # Left Navcam + translation=(-0.7675, -0.8537, 1.9793), # Right Navcam + rotation=math_utils.quat_from_rpy(0.0, 15.0, -90.0), + ), + ), + ## Joints + regex_drive_joints=regex_drive_joints, + regex_steer_joints=regex_steer_joints, + ) diff --git a/space_robotics_bench/assets/terrain/__init__.py b/space_robotics_bench/assets/terrain/__init__.py new file mode 100644 index 0000000..143a205 --- /dev/null +++ b/space_robotics_bench/assets/terrain/__init__.py @@ -0,0 +1,81 @@ +import math +from typing import Any, Dict, Optional, Tuple + +import space_robotics_bench.core.envs as env_utils +import space_robotics_bench.core.sim as sim_utils +from space_robotics_bench.core.assets import AssetBaseCfg + +from .lunar_surface_procgen import LunarSurfaceProcgenCfg +from .martian_surface_procgen import MartianSurfaceProcgenCfg + + +def terrain_from_env_cfg( + env_cfg: env_utils.EnvironmentConfig, + *, + size: Tuple[float, float] = (10.0, 10.0), + num_assets: int = 1, + prim_path: str = "{ENV_REGEX_NS}/terrain", + spawn_kwargs: Dict[str, Any] = {}, + procgen_kwargs: Dict[str, Any] = {}, + **kwargs, +) -> Optional[AssetBaseCfg]: + spawn: Optional[sim_utils.SpawnerCfg] = None + + match env_cfg.assets.terrain.variant: + case env_utils.AssetVariant.PRIMITIVE: + prim_path = "/World/terrain" + size = ( + 10 * math.sqrt(num_assets) * size[0], + 10 * math.sqrt(num_assets) * size[1], + ) + spawn = sim_utils.GroundPlaneCfg( + size=size, color=(0.0, 158.0 / 255.0, 218.0 / 255.0), **spawn_kwargs + ) + + case env_utils.AssetVariant.PROCEDURAL: + if spawn_kwargs.get("collision_props") is None: + spawn_kwargs["collision_props"] = sim_utils.CollisionPropertiesCfg() + usd_file_cfg = sim_utils.UsdFileCfg( + usd_path="IGNORED", + **spawn_kwargs, + ) + + match env_cfg.scenario: + case env_utils.Scenario.MOON: + spawn = LunarSurfaceProcgenCfg( + num_assets=num_assets, + usd_file_cfg=usd_file_cfg, + seed=env_cfg.seed, + detail=env_cfg.detail, + ) + + case env_utils.Scenario.MARS: + spawn = MartianSurfaceProcgenCfg( + num_assets=num_assets, + usd_file_cfg=usd_file_cfg, + seed=env_cfg.seed, + detail=env_cfg.detail, + ) + case _: + return None + + # Set height to 10% of the average planar size + scale = (*size, (size[0] + size[1]) / 20.0) + for node_cfg in spawn.geometry_nodes.values(): + if node_cfg.get("scale") is not None: + node_cfg["scale"] = scale + + for key, value in procgen_kwargs.items(): + for node_cfg in spawn.geometry_nodes.values(): + if node_cfg.get(key) is not None: + node_cfg[key] = value + elif hasattr(spawn, key): + setattr(spawn, key, value) + + if spawn is None: + return None + return AssetBaseCfg( + prim_path=prim_path, + spawn=spawn, + **kwargs, + ) diff --git a/space_robotics_bench/assets/terrain/lunar_surface_procgen.py b/space_robotics_bench/assets/terrain/lunar_surface_procgen.py new file mode 100644 index 0000000..c4c2fac --- /dev/null +++ b/space_robotics_bench/assets/terrain/lunar_surface_procgen.py @@ -0,0 +1,32 @@ +from os import path +from typing import Any, Dict, List, Optional + +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.sim as sim_utils +from space_robotics_bench.paths import SRB_ASSETS_DIR_SRB_TERRAIN +from space_robotics_bench.utils.path import abs_listdir + + +@configclass +class LunarSurfaceProcgenCfg(sim_utils.BlenderNodesAssetCfg): + name: str = "lunar_surface" + autorun_scripts: List[str] = abs_listdir( + path.join(SRB_ASSETS_DIR_SRB_TERRAIN, "lunar_surface_procgen") + ) + + # Geometry + geometry_nodes: Dict[str, Dict[str, Any]] = { + "LunarTerrain": { + "density": 0.1, # Density of the terrain in meters per vertex + "scale": (10.0, 10.0, 1.0), # Metric scale of the mesh + "flat_area_size": 0.0, # Size of a flat round area at the centre of the mesh in meters + "rock_mesh_boolean": False, # Flag to enable mesh boolean between the terrain and rocks + } + } + decimate_face_count: Optional[int] = None + decimate_angle_limit: Optional[float] = None + + # Material + material: Optional[str] = "LunarSurface" + texture_resolution: int = 4096 diff --git a/space_robotics_bench/assets/terrain/martian_surface_procgen.py b/space_robotics_bench/assets/terrain/martian_surface_procgen.py new file mode 100644 index 0000000..665321f --- /dev/null +++ b/space_robotics_bench/assets/terrain/martian_surface_procgen.py @@ -0,0 +1,32 @@ +from os import path +from typing import Any, Dict, List, Optional + +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.sim as sim_utils +from space_robotics_bench.paths import SRB_ASSETS_DIR_SRB_TERRAIN +from space_robotics_bench.utils.path import abs_listdir + + +@configclass +class MartianSurfaceProcgenCfg(sim_utils.BlenderNodesAssetCfg): + name: str = "martian_surface" + autorun_scripts: List[str] = abs_listdir( + path.join(SRB_ASSETS_DIR_SRB_TERRAIN, "martian_surface_procgen") + ) + + # Geometry + geometry_nodes: Dict[str, Dict[str, Any]] = { + "MartianTerrain": { + "density": 0.1, # Density of the terrain in meters per vertex + "scale": (10.0, 10.0, 1.0), # Metric scale of the mesh + "flat_area_size": 0.0, # Size of a flat round area at the centre of the mesh in meters + "rock_mesh_boolean": False, # Flag to enable mesh boolean between the terrain and rocks + } + } + decimate_face_count: Optional[int] = None + decimate_angle_limit: Optional[float] = None + + # Material + material: Optional[str] = "MartianSurface" + texture_resolution: int = 4096 diff --git a/space_robotics_bench/assets/vehicle/__init__.py b/space_robotics_bench/assets/vehicle/__init__.py new file mode 100644 index 0000000..b22e11b --- /dev/null +++ b/space_robotics_bench/assets/vehicle/__init__.py @@ -0,0 +1,33 @@ +from typing import Any, Dict, Optional + +import space_robotics_bench.core.assets as asset_utils +import space_robotics_bench.core.envs as env_utils + +from .construction_rover import construction_rover_cfg +from .gateway import gateway_cfg + + +def vehicle_from_env_cfg( + env_cfg: env_utils.EnvironmentConfig, + *, + prim_path: str = "{ENV_REGEX_NS}/vehicle", + spawn_kwargs: Dict[str, Any] = {}, + **kwargs, +) -> Optional[asset_utils.VehicleCfg]: + vehicle_cfg: Optional[asset_utils.VehicleCfg] = None + match env_cfg.assets.vehicle.variant: + case env_utils.AssetVariant.NONE: + return None + + case _: + match env_cfg.scenario: + case env_utils.Scenario.MOON | env_utils.Scenario.MARS: + vehicle_cfg = construction_rover_cfg( + prim_path=prim_path, spawn_kwargs=spawn_kwargs, **kwargs + ) + case env_utils.Scenario.ORBIT: + vehicle_cfg = gateway_cfg( + prim_path=prim_path, spawn_kwargs=spawn_kwargs, **kwargs + ) + + return vehicle_cfg diff --git a/space_robotics_bench/assets/vehicle/construction_rover.py b/space_robotics_bench/assets/vehicle/construction_rover.py new file mode 100644 index 0000000..fbda556 --- /dev/null +++ b/space_robotics_bench/assets/vehicle/construction_rover.py @@ -0,0 +1,55 @@ +from os import path +from typing import Any, Dict + +import space_robotics_bench.core.assets as asset_utils +import space_robotics_bench.core.sim as sim_utils +import space_robotics_bench.utils.math as math_utils +from space_robotics_bench.paths import SRB_ASSETS_DIR_SRB_VEHICLE + + +def construction_rover_cfg( + *, + prim_path: str = "{ENV_REGEX_NS}/vehicle", + spawn_kwargs: Dict[str, Any] = {}, + **kwargs, +) -> asset_utils.VehicleCfg: + if spawn_kwargs.get("collision_props") is None: + spawn_kwargs["collision_props"] = sim_utils.CollisionPropertiesCfg() + + return asset_utils.VehicleCfg( + ## Model + asset_cfg=asset_utils.AssetBaseCfg( + prim_path=prim_path, + spawn=sim_utils.UsdFileCfg( + usd_path=path.join( + SRB_ASSETS_DIR_SRB_VEHICLE, + "construction_rover", + "construction_rover.usdc", + ), + **spawn_kwargs, + ), + **kwargs, + ), + ## Frames + frame_manipulator_base=asset_utils.FrameCfg( + prim_relpath="manipulator_base", + offset=asset_utils.TransformCfg( + translation=(0.0, 0.0, 0.25), + ), + ), + frame_camera_base=asset_utils.FrameCfg( + prim_relpath="camera_base", + offset=asset_utils.TransformCfg( + translation=(0.21, 0.0, 0.0), + rotation=math_utils.quat_from_rpy(0.0, 45.0, 0.0), + ), + ), + frame_cargo_bay=asset_utils.FrameCfg( + prim_relpath="cargo_bay", + offset=asset_utils.TransformCfg( + translation=(-0.6, 0.0, 0.3), + ), + ), + ## Properties + height=0.25, + ) diff --git a/space_robotics_bench/assets/vehicle/gateway.py b/space_robotics_bench/assets/vehicle/gateway.py new file mode 100644 index 0000000..c3e3830 --- /dev/null +++ b/space_robotics_bench/assets/vehicle/gateway.py @@ -0,0 +1,53 @@ +from os import path +from typing import Any, Dict + +import space_robotics_bench.core.assets as asset_utils +import space_robotics_bench.core.sim as sim_utils +import space_robotics_bench.utils.math as math_utils +from space_robotics_bench.paths import SRB_ASSETS_DIR_SRB_VEHICLE + + +def gateway_cfg( + *, + prim_path: str = "{ENV_REGEX_NS}/vehicle", + spawn_kwargs: Dict[str, Any] = {}, + **kwargs, +) -> asset_utils.VehicleCfg: + if spawn_kwargs.get("collision_props") is None: + spawn_kwargs["collision_props"] = sim_utils.CollisionPropertiesCfg() + + return asset_utils.VehicleCfg( + ## Model + asset_cfg=asset_utils.AssetBaseCfg( + prim_path=prim_path, + spawn=sim_utils.UsdFileCfg( + usd_path=path.join( + SRB_ASSETS_DIR_SRB_VEHICLE, "gateway", "gateway.usdc" + ), + **spawn_kwargs, + ), + **kwargs, + ), + ## Frames + frame_manipulator_base=asset_utils.FrameCfg( + prim_relpath="base", + offset=asset_utils.TransformCfg( + translation=(0.0, 0.0, 0.0), + ), + ), + frame_camera_base=asset_utils.FrameCfg( + prim_relpath="camera_base", + offset=asset_utils.TransformCfg( + translation=(0.21, 0.0, 0.0), + rotation=math_utils.quat_from_rpy(0.0, 15.0, 0.0), + ), + ), + frame_cargo_bay=asset_utils.FrameCfg( + prim_relpath="cargo_bay", + offset=asset_utils.TransformCfg( + translation=(-0.6, 0.0, 0.3), + ), + ), + ## Properties + height=0.0, + ) diff --git a/space_robotics_bench/core/__init__.py b/space_robotics_bench/core/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/space_robotics_bench/core/actions/__init__.py b/space_robotics_bench/core/actions/__init__.py new file mode 100644 index 0000000..295b068 --- /dev/null +++ b/space_robotics_bench/core/actions/__init__.py @@ -0,0 +1,5 @@ +from omni.isaac.lab.envs.mdp.actions import * # noqa: F403 + +from .aerial import * # noqa: F403 +from .manipulator import * # noqa: F403 +from .mobile import * # noqa: F403 diff --git a/space_robotics_bench/core/actions/aerial.py b/space_robotics_bench/core/actions/aerial.py new file mode 100644 index 0000000..30c98c6 --- /dev/null +++ b/space_robotics_bench/core/actions/aerial.py @@ -0,0 +1,122 @@ +from collections.abc import Sequence +from dataclasses import MISSING +from typing import Dict, Union + +import torch +from omni.isaac.lab.managers import ActionTerm, ActionTermCfg +from omni.isaac.lab.utils import configclass + +from space_robotics_bench.core.assets import Articulation +from space_robotics_bench.core.envs import BaseEnv +from space_robotics_bench.utils.math import euler_xyz_from_quat, quat_from_euler_xyz + + +class MultiCopterAction(ActionTerm): + cfg: "MultiCopterActionCfg" + _asset: Articulation + + def __init__(self, cfg: "MultiCopterActionCfg", env: BaseEnv): + super().__init__(cfg, env) + + self._body_index = self._asset.find_bodies(self.cfg.frame_base)[0] + self._rotor_indices, rotor_names = self._asset.find_joints( + self.cfg.regex_joints_rotors + ) + + if isinstance(self.cfg.nominal_rpm, Dict): + self._nominal_rpm = torch.tensor( + [self.cfg.nominal_rpm[name] for name in rotor_names], + device=self.device, + dtype=torch.float32, + ).repeat(self._asset.num_instances, 1) + else: + self._nominal_rpm = torch.tensor( + [self.cfg.nominal_rpm] * len(rotor_names), + device=self.device, + dtype=torch.float32, + ).repeat(self._asset.num_instances, 1) + + @property + def action_dim(self) -> int: + return 4 + + @property + def raw_actions(self) -> torch.Tensor: + return self._raw_actions + + @property + def processed_actions(self) -> torch.Tensor: + return self._processed_actions + + def process_actions(self, actions): + self._raw_actions = actions + self._processed_actions = self.raw_actions * self.cfg.scale + + def apply_actions(self): + self._asset.set_joint_velocity_target(self._nominal_rpm) + + current_velocity = self._asset._data.body_vel_w[:, self._body_index].squeeze(1) + current_yaw = euler_xyz_from_quat(self._asset._data.root_quat_w)[2] + + applied_velocity_lin = ( + self.cfg.controller_damping * current_velocity[:, :3] + + (1 - self.cfg.controller_damping) * self.processed_actions[:, :3] + ) + applied_velocity_rot_yaw = ( + self.cfg.controller_damping * current_velocity[:, 5] + + (1 - self.cfg.controller_damping) * self.processed_actions[:, 3] + ) + applied_velocities = torch.cat( + ( + applied_velocity_lin, + torch.zeros_like(applied_velocity_rot_yaw)[:, None], + torch.zeros_like(applied_velocity_rot_yaw)[:, None], + applied_velocity_rot_yaw[:, None], + ), + dim=1, + ) + if self.cfg.noise > 0: + applied_velocities += self.cfg.noise * torch.randn_like(applied_velocities) + self._asset.write_root_velocity_to_sim(applied_velocities) + + target_tilt = ( + torch.atan2(applied_velocity_lin[:, 1], applied_velocity_lin[:, 0]) + - current_yaw + ) + tilt_factor = self.cfg.tilt_magnitude * torch.norm( + current_velocity[:, :2], dim=1 + ) + target_roll = -torch.sin(target_tilt) * tilt_factor + target_pitch = torch.cos(target_tilt) * tilt_factor + self._asset.write_root_pose_to_sim( + torch.cat( + ( + self._asset._data.root_pos_w, + quat_from_euler_xyz(target_roll, target_pitch, current_yaw), + ), + dim=1, + ) + ) + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + pass + + +@configclass +class MultiCopterActionCfg(ActionTermCfg): + class_type: ActionTerm = MultiCopterAction + + frame_base: str = MISSING + regex_joints_rotors: str = MISSING + + nominal_rpm: Union[float, Dict["str", float]] = 1000.0 + tilt_magnitude: float = 1.0 + controller_damping: float = 0.98 + + scale: float = 1.0 + noise: float = 0.0025 + + +@configclass +class MultiCopterActionGroupCfg: + flight: MultiCopterActionCfg = MISSING diff --git a/space_robotics_bench/core/actions/manipulator.py b/space_robotics_bench/core/actions/manipulator.py new file mode 100644 index 0000000..efed163 --- /dev/null +++ b/space_robotics_bench/core/actions/manipulator.py @@ -0,0 +1,14 @@ +from dataclasses import MISSING + +from omni.isaac.lab.utils import configclass + +from space_robotics_bench.core.actions import ( + BinaryJointPositionActionCfg, + DifferentialInverseKinematicsActionCfg, +) + + +@configclass +class ManipulatorTaskSpaceActionCfg: + arm: DifferentialInverseKinematicsActionCfg = MISSING + hand: BinaryJointPositionActionCfg = MISSING diff --git a/space_robotics_bench/core/actions/mobile.py b/space_robotics_bench/core/actions/mobile.py new file mode 100644 index 0000000..7052552 --- /dev/null +++ b/space_robotics_bench/core/actions/mobile.py @@ -0,0 +1,247 @@ +from collections.abc import Sequence +from dataclasses import MISSING +from typing import List, Optional, Tuple + +import torch +from omni.isaac.lab.managers import ActionTerm, ActionTermCfg +from omni.isaac.lab.utils import configclass + +from space_robotics_bench.core.assets import Articulation +from space_robotics_bench.core.envs import BaseEnv + + +class WheeledRoverAction(ActionTerm): + cfg: "WheeledRoverActionCfg" + _asset: Articulation + + def __init__(self, cfg: "WheeledRoverActionCfg", env: BaseEnv): + super().__init__(cfg, env) + + self._steering_joint_indices = self._asset.find_joints( + self.cfg.steering_joint_names, preserve_order=True + )[0] + self._drive_joint_indices = self._asset.find_joints( + self.cfg.drive_joint_names, preserve_order=True + )[0] + + @property + def action_dim(self) -> int: + return 2 + + @property + def raw_actions(self) -> torch.Tensor: + return self._raw_actions + + @property + def processed_actions(self) -> torch.Tensor: + return self._processed_actions + + def process_actions(self, actions): + self._raw_actions = actions + self._processed_actions = self.raw_actions * self.cfg.scale + self._processed_actions[:, 1] /= torch.pi + + def apply_actions(self): + steer_joint_positions, drive_joint_velocities = self._process_actions( + velocity_lin=self._processed_actions[:, 0], + velocity_ang=self._processed_actions[:, 1], + wheelbase=self.cfg.wheelbase, + wheelbase_mid=self.cfg.wheelbase_mid, + wheel_radius=self.cfg.wheel_radius, + ) + self._asset.set_joint_position_target( + steer_joint_positions, joint_ids=self._steering_joint_indices + ) + self._asset.set_joint_velocity_target( + drive_joint_velocities, joint_ids=self._drive_joint_indices + ) + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + pass + + @staticmethod + @torch.jit.script + def _process_actions( + *, + velocity_lin: torch.Tensor, + velocity_ang: torch.Tensor, + wheelbase: Tuple[float, float], + wheelbase_mid: float, + wheel_radius: float, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Computes the steering joint positions and drive joint velocities from linear and angular velocities of a wheeled rover. + + Adapted from: + """ + num_envs = velocity_lin.size(0) + device = velocity_lin.device + + abs_velocity_lin = torch.abs(velocity_lin) + abs_velocity_ang = torch.abs(velocity_ang) + turn_direction = torch.sign(velocity_ang) + turning_radius = torch.where( + torch.logical_not(abs_velocity_ang == 0) + | torch.logical_not(abs_velocity_lin == 0), + abs_velocity_lin / abs_velocity_ang, + torch.tensor(10e10, device=device), + ) + is_point_turn = turning_radius < wheelbase_mid / 2.0 + + r_front_left = turning_radius - (wheelbase[0] / 2) * turn_direction + steer_angle_front_left = torch.where( + is_point_turn, + torch.tensor(-torch.pi / 4, device=device), + turn_direction + * torch.atan2( + (torch.tensor(wheelbase[1], device=device).repeat(num_envs) / 2), + r_front_left, + ), + ) + r_front_right = turning_radius + (wheelbase[0] / 2) * turn_direction + steer_angle_front_right = torch.where( + is_point_turn, + torch.tensor(torch.pi / 4, device=device), + turn_direction + * torch.atan2( + (torch.tensor(wheelbase[1], device=device).repeat(num_envs) / 2), + r_front_right, + ), + ) + r_rear_left = turning_radius - (wheelbase[0] / 2) * turn_direction + steer_angle_rear_left = torch.where( + is_point_turn, + torch.tensor(torch.pi / 4, device=device), + -turn_direction + * torch.atan2( + (torch.tensor(wheelbase[1], device=device).repeat(num_envs) / 2), + r_rear_left, + ), + ) + r_rear_right = turning_radius + (wheelbase[0] / 2) * turn_direction + steer_angle_rear_right = torch.where( + is_point_turn, + torch.tensor(-torch.pi / 4, device=device), + -turn_direction + * torch.atan2( + (torch.tensor(wheelbase[1], device=device).repeat(num_envs) / 2), + r_rear_right, + ), + ) + steer_joint_positions = torch.stack( + [ + steer_angle_front_left, + steer_angle_front_right, + steer_angle_rear_left, + steer_angle_rear_right, + ], + dim=1, + ) + + drive_direction = torch.sign(velocity_lin) + drive_direction = torch.where( + drive_direction == 0, drive_direction + 1, drive_direction + ) + velocity_lin_front_left = torch.where( + is_point_turn, + -turn_direction * abs_velocity_ang, + drive_direction + * torch.where( + abs_velocity_ang == 0, + abs_velocity_lin, + (r_front_left * abs_velocity_ang), + ), + ) + velocity_lin_front_right = torch.where( + is_point_turn, + turn_direction * abs_velocity_ang, + drive_direction + * torch.where( + abs_velocity_ang == 0, + abs_velocity_lin, + (r_front_right * abs_velocity_ang), + ), + ) + velocity_lin_rear_left = torch.where( + is_point_turn, + -turn_direction * abs_velocity_ang, + drive_direction + * torch.where( + abs_velocity_ang == 0, + abs_velocity_lin, + (r_rear_left * abs_velocity_ang), + ), + ) + velocity_lin_rear_right = torch.where( + is_point_turn, + turn_direction * abs_velocity_ang, + drive_direction + * torch.where( + abs_velocity_ang == 0, + abs_velocity_lin, + (r_rear_right * abs_velocity_ang), + ), + ) + velocity_lin_mid_left = torch.where( + is_point_turn, + -turn_direction * abs_velocity_ang, + drive_direction + * torch.where( + abs_velocity_ang == 0, + abs_velocity_lin, + ( + turning_radius + - (wheelbase_mid / 2) * turn_direction * abs_velocity_ang + ), + ), + ) + velocity_lin_mid_right = torch.where( + is_point_turn, + turn_direction * abs_velocity_ang, + drive_direction + * torch.where( + abs_velocity_ang == 0, + abs_velocity_lin, + ( + turning_radius + + (wheelbase_mid / 2) * turn_direction * abs_velocity_ang + ), + ), + ) + drive_joint_velocities = torch.stack( + [ + velocity_lin_front_left, + velocity_lin_front_right, + velocity_lin_rear_left, + velocity_lin_rear_right, + velocity_lin_mid_left, + velocity_lin_mid_right, + ], + dim=1, + ) / (2.0 * wheel_radius) + + return steer_joint_positions, drive_joint_velocities + + +@configclass +class WheeledRoverActionCfg(ActionTermCfg): + class_type: type = WheeledRoverAction + + steering_joint_names: List[str] = MISSING + drive_joint_names: List[str] = MISSING + + wheelbase: Tuple[float, float] = MISSING + wheelbase_mid: Optional[float] = None + + wheel_radius: float = MISSING + + scale: float = 1.0 + + def __post_init__(self): + if self.wheelbase_mid is None: + self.wheelbase_mid = self.wheelbase[1] + + +@configclass +class WheeledRoverActionGroupCfg: + drive: WheeledRoverActionCfg = MISSING diff --git a/space_robotics_bench/core/assets/__init__.py b/space_robotics_bench/core/assets/__init__.py new file mode 100644 index 0000000..2126639 --- /dev/null +++ b/space_robotics_bench/core/assets/__init__.py @@ -0,0 +1,6 @@ +from omni.isaac.lab.assets import * # noqa: F403 + +from .common import * # noqa: F403 +from .robots import * # noqa: F403 +from .terrains import * # noqa: F403 +from .vehicles import * # noqa: F403 diff --git a/space_robotics_bench/core/assets/common/__init__.py b/space_robotics_bench/core/assets/common/__init__.py new file mode 100644 index 0000000..244878f --- /dev/null +++ b/space_robotics_bench/core/assets/common/__init__.py @@ -0,0 +1,2 @@ +from .asset import * # noqa: F403 +from .transform import * # noqa: F403 diff --git a/space_robotics_bench/core/assets/common/asset.py b/space_robotics_bench/core/assets/common/asset.py new file mode 100644 index 0000000..a380f6c --- /dev/null +++ b/space_robotics_bench/core/assets/common/asset.py @@ -0,0 +1,8 @@ +from pydantic import BaseModel + +from space_robotics_bench.core.assets import AssetBaseCfg + + +class AssetCfg(BaseModel): + ## Model + asset_cfg: AssetBaseCfg diff --git a/space_robotics_bench/core/assets/common/transform.py b/space_robotics_bench/core/assets/common/transform.py new file mode 100644 index 0000000..a03c659 --- /dev/null +++ b/space_robotics_bench/core/assets/common/transform.py @@ -0,0 +1,13 @@ +from typing import Tuple + +from pydantic import BaseModel + + +class TransformCfg(BaseModel): + translation: Tuple[float, float, float] = (0.0, 0.0, 0.0) + rotation: Tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0) # w, x, y, z + + +class FrameCfg(BaseModel): + prim_relpath: str + offset: TransformCfg = TransformCfg() diff --git a/space_robotics_bench/core/assets/objects/__init__.py b/space_robotics_bench/core/assets/objects/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/space_robotics_bench/core/assets/robots/__init__.py b/space_robotics_bench/core/assets/robots/__init__.py new file mode 100644 index 0000000..868a2b3 --- /dev/null +++ b/space_robotics_bench/core/assets/robots/__init__.py @@ -0,0 +1,4 @@ +from .base import * # noqa: F403 isort:skip +from .aerial import * # noqa: F403 +from .manipulator import * # noqa: F403 +from .mobile import * # noqa: F403 diff --git a/space_robotics_bench/core/assets/robots/aerial.py b/space_robotics_bench/core/assets/robots/aerial.py new file mode 100644 index 0000000..c8cc050 --- /dev/null +++ b/space_robotics_bench/core/assets/robots/aerial.py @@ -0,0 +1,28 @@ +from typing import Any + +from space_robotics_bench.core.actions import MultiCopterActionGroupCfg +from space_robotics_bench.core.assets import FrameCfg + +from . import RobotCfg + + +class AerialRobotCfg(RobotCfg): + ## Actions + action_cfg: Any + + +class MultiCopterCfg(AerialRobotCfg): + class Config: + arbitrary_types_allowed = True # Due to MultiCopterActionGroupCfg + + ## Actions + action_cfg: MultiCopterActionGroupCfg + + ## Frames + frame_camera_bottom: FrameCfg + + ## Links + regex_links_rotors: str + + ## Joints + regex_joints_rotors: str diff --git a/space_robotics_bench/core/assets/robots/base.py b/space_robotics_bench/core/assets/robots/base.py new file mode 100644 index 0000000..e9136b6 --- /dev/null +++ b/space_robotics_bench/core/assets/robots/base.py @@ -0,0 +1,18 @@ +from typing import Any, Union + +from space_robotics_bench.core.assets import ( + ArticulationCfg, + AssetCfg, + FrameCfg, + RigidObjectCfg, +) + + +class RobotCfg(AssetCfg): + asset_cfg: Union[ArticulationCfg, RigidObjectCfg] + + ## Actions + action_cfg: Any + + ## Frames + frame_base: FrameCfg diff --git a/space_robotics_bench/core/assets/robots/manipulator.py b/space_robotics_bench/core/assets/robots/manipulator.py new file mode 100644 index 0000000..739d968 --- /dev/null +++ b/space_robotics_bench/core/assets/robots/manipulator.py @@ -0,0 +1,22 @@ +from space_robotics_bench.core.actions import ManipulatorTaskSpaceActionCfg +from space_robotics_bench.core.assets import FrameCfg + +from . import RobotCfg + + +class ManipulatorCfg(RobotCfg): + ## Actions + action_cfg: ManipulatorTaskSpaceActionCfg + + ## Frames + frame_ee: FrameCfg + frame_camera_base: FrameCfg + frame_camera_wrist: FrameCfg + + ## Links + regex_links_arm: str + regex_links_hand: str + + ## Joints + regex_joints_arm: str + regex_joints_hand: str diff --git a/space_robotics_bench/core/assets/robots/mobile.py b/space_robotics_bench/core/assets/robots/mobile.py new file mode 100644 index 0000000..0ffd7c2 --- /dev/null +++ b/space_robotics_bench/core/assets/robots/mobile.py @@ -0,0 +1,23 @@ +from typing import Any + +from space_robotics_bench.core.actions import WheeledRoverActionGroupCfg +from space_robotics_bench.core.assets import FrameCfg + +from . import RobotCfg + + +class MobileRobotCfg(RobotCfg): + ## Actions + action_cfg: Any + + +class WheeledRoverCfg(MobileRobotCfg): + ## Actions + action_cfg: WheeledRoverActionGroupCfg + + ## Frames + frame_camera_front: FrameCfg + + ## Joints + regex_drive_joints: str + regex_steer_joints: str diff --git a/space_robotics_bench/core/assets/robots/propelled.py b/space_robotics_bench/core/assets/robots/propelled.py new file mode 100644 index 0000000..9d7f0e6 --- /dev/null +++ b/space_robotics_bench/core/assets/robots/propelled.py @@ -0,0 +1,20 @@ +from space_robotics_bench.core.actions import PropelledRobotTaskSpaceActionCfg +from space_robotics_bench.core.assets import FrameCfg + +from . import RobotCfg + + +class PropelledRobotCfg(RobotCfg): + ## Actions + action_cfg: PropelledRobotTaskSpaceActionCfg + + +class MultiCopterCfg(PropelledRobotCfg): + ## Frames + frame_camera_bottom: FrameCfg + + # ## Links + # regex_links_wheels: str + + # ## Joints + # regex_joints_wheels: str diff --git a/space_robotics_bench/core/assets/terrains/__init__.py b/space_robotics_bench/core/assets/terrains/__init__.py new file mode 100644 index 0000000..69a898d --- /dev/null +++ b/space_robotics_bench/core/assets/terrains/__init__.py @@ -0,0 +1 @@ +from omni.isaac.lab.terrains import * # noqa: F403 diff --git a/space_robotics_bench/core/assets/vehicles/__init__.py b/space_robotics_bench/core/assets/vehicles/__init__.py new file mode 100644 index 0000000..12e69f4 --- /dev/null +++ b/space_robotics_bench/core/assets/vehicles/__init__.py @@ -0,0 +1 @@ +from .base import * # noqa: F403 diff --git a/space_robotics_bench/core/assets/vehicles/base.py b/space_robotics_bench/core/assets/vehicles/base.py new file mode 100644 index 0000000..b3b5d63 --- /dev/null +++ b/space_robotics_bench/core/assets/vehicles/base.py @@ -0,0 +1,19 @@ +from typing import Optional, Union + +from space_robotics_bench.core.assets import ( + ArticulationCfg, + AssetBaseCfg, + AssetCfg, + FrameCfg, + RigidObjectCfg, +) + + +class VehicleCfg(AssetCfg): + ## Model + asset_cfg: Union[AssetBaseCfg, ArticulationCfg, RigidObjectCfg] + + ## Frames + frame_manipulator_base: FrameCfg + frame_camera_base: Optional[FrameCfg] + frame_cargo_bay: FrameCfg diff --git a/space_robotics_bench/core/envs/__init__.py b/space_robotics_bench/core/envs/__init__.py new file mode 100644 index 0000000..c0bf712 --- /dev/null +++ b/space_robotics_bench/core/envs/__init__.py @@ -0,0 +1,5 @@ +from omni.isaac.lab.envs import * # noqa: F403 + +from space_robotics_bench._rs.envs import * # noqa: F403 + +from .base import * # noqa: F403 diff --git a/space_robotics_bench/core/envs/base/__init__.py b/space_robotics_bench/core/envs/base/__init__.py new file mode 100644 index 0000000..cb13bc4 --- /dev/null +++ b/space_robotics_bench/core/envs/base/__init__.py @@ -0,0 +1,2 @@ +from .direct import * # noqa: F403 +from .managed import * # noqa: F403 diff --git a/space_robotics_bench/core/envs/base/direct/__init__.py b/space_robotics_bench/core/envs/base/direct/__init__.py new file mode 100644 index 0000000..623a285 --- /dev/null +++ b/space_robotics_bench/core/envs/base/direct/__init__.py @@ -0,0 +1,2 @@ +from .cfg import * # noqa: F403 +from .impl import * # noqa: F403 diff --git a/space_robotics_bench/core/envs/base/direct/cfg.py b/space_robotics_bench/core/envs/base/direct/cfg.py new file mode 100644 index 0000000..4f5c737 --- /dev/null +++ b/space_robotics_bench/core/envs/base/direct/cfg.py @@ -0,0 +1,31 @@ +from os import path + +from omni.isaac.lab.envs import DirectRLEnvCfg +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.envs as env_utils +from space_robotics_bench.paths import SRB_CONFIG_DIR + + +@configclass +class BaseEnvCfg(DirectRLEnvCfg): + """ + Extended version of :class:`omni.isaac.lab.envs.DirectRLEnvCfg`. + """ + + ## Updated defaults + # Disable UI window by default + ui_window_class_type: type | None = None + # Redundant: spaces are automatically extracted + num_actions: int = 0 + # Redundant: spaces are automatically extracted + num_observations: int = 0 + + ## Environment + env_cfg: env_utils.EnvironmentConfig = env_utils.EnvironmentConfig.extract( + cfg_path=path.join(SRB_CONFIG_DIR, "env.yaml") + ) + + ## Misc + # Flag that disables the timeout for the environment + enable_truncation: bool = True diff --git a/space_robotics_bench/core/envs/base/direct/impl.py b/space_robotics_bench/core/envs/base/direct/impl.py new file mode 100644 index 0000000..ce4cad9 --- /dev/null +++ b/space_robotics_bench/core/envs/base/direct/impl.py @@ -0,0 +1,83 @@ +from typing import Sequence + +import gymnasium +import numpy as np +import torch +from omni.isaac.lab.envs import DirectRLEnv +from omni.isaac.lab.managers import ActionManager + +from . import BaseEnvCfg + + +class __PostInitCaller(type): + def __call__(cls, *args, **kwargs): + obj = type.__call__(cls, *args, **kwargs) + obj.__post_init__() + return obj + + +class BaseEnv(DirectRLEnv, metaclass=__PostInitCaller): + """ + Extended version of :class:`omni.isaac.lab.envs.DirectRLEnv`. + """ + + def __init__(self, cfg: BaseEnvCfg, render_mode: str | None = None, **kwargs): + super().__init__(cfg, render_mode, **kwargs) + + if self.cfg.actions: + self.action_manager = ActionManager(self.cfg.actions, self) + self.cfg.num_actions = self.action_manager.total_action_dim + print("[INFO] Action Manager: ", self.action_manager) + + def __post_init__(self): + self._update_gym_env_spaces() + + def close(self): + if not self._is_closed: + if self.cfg.actions: + del self.action_manager + + super().close() + + def _reset_idx(self, env_ids: Sequence[int]): + if self.cfg.actions: + self.action_manager.reset(env_ids) + + if self.cfg.events: + self.event_manager.reset(env_ids) + + super()._reset_idx(env_ids) + + def _pre_physics_step(self, actions: torch.Tensor): + if self.cfg.actions: + self.action_manager.process_action(actions) + else: + super()._pre_physics_step(actions) + + def _apply_action(self): + if self.cfg.actions: + self.action_manager.apply_action() + else: + super()._apply_action() + + def _update_gym_env_spaces(self): + # Action space + self.single_action_space = gymnasium.spaces.Box( + low=-np.inf, high=np.inf, shape=(self.cfg.num_actions,) + ) + self.action_space = gymnasium.vector.utils.batch_space( + self.single_action_space, self.num_envs + ) + + # Observation space + self.single_observation_space = gymnasium.spaces.Dict({}) + for ( + obs_key, + obs_buf, + ) in self._get_observations().items(): + self.single_observation_space[obs_key] = gymnasium.spaces.Box( + low=-np.inf, high=np.inf, shape=obs_buf.shape[1:] + ) + self.observation_space = gymnasium.vector.utils.batch_space( + self.single_observation_space, self.num_envs + ) diff --git a/space_robotics_bench/core/envs/base/managed/__init__.py b/space_robotics_bench/core/envs/base/managed/__init__.py new file mode 100644 index 0000000..623a285 --- /dev/null +++ b/space_robotics_bench/core/envs/base/managed/__init__.py @@ -0,0 +1,2 @@ +from .cfg import * # noqa: F403 +from .impl import * # noqa: F403 diff --git a/space_robotics_bench/core/envs/base/managed/cfg.py b/space_robotics_bench/core/envs/base/managed/cfg.py new file mode 100644 index 0000000..d57495b --- /dev/null +++ b/space_robotics_bench/core/envs/base/managed/cfg.py @@ -0,0 +1,12 @@ +from omni.isaac.lab.envs import ManagerBasedRLEnvCfg +from omni.isaac.lab.utils import configclass + + +@configclass +class BaseEnvManagedCfg(ManagerBasedRLEnvCfg): + """ + Extended version of :class:`omni.isaac.lab.envs.ManagerBasedRLEnvCfg`. + """ + + # Disable UI window by default + ui_window_class_type: type | None = None diff --git a/space_robotics_bench/core/envs/base/managed/impl.py b/space_robotics_bench/core/envs/base/managed/impl.py new file mode 100644 index 0000000..391e699 --- /dev/null +++ b/space_robotics_bench/core/envs/base/managed/impl.py @@ -0,0 +1,7 @@ +from omni.isaac.lab.envs import ManagerBasedRLEnv + + +class BaseEnvManaged(ManagerBasedRLEnv): + """ + Extended version of :class:`omni.isaac.lab.envs.ManagerBasedRLEnv`. + """ diff --git a/space_robotics_bench/core/interfaces/__init__.py b/space_robotics_bench/core/interfaces/__init__.py new file mode 100644 index 0000000..d23cfd1 --- /dev/null +++ b/space_robotics_bench/core/interfaces/__init__.py @@ -0,0 +1 @@ +from .ros2 import * # noqa: F403 diff --git a/space_robotics_bench/core/interfaces/ros2.py b/space_robotics_bench/core/interfaces/ros2.py new file mode 100644 index 0000000..c0443f7 --- /dev/null +++ b/space_robotics_bench/core/interfaces/ros2.py @@ -0,0 +1,532 @@ +import json +import threading +from collections.abc import Callable +from queue import Queue +from typing import Any, Dict, Optional, Tuple, Union + +import numpy as np +import rclpy +import torch +from builtin_interfaces.msg import Time +from geometry_msgs.msg import ( + Quaternion, + Transform, + TransformStamped, + Twist, + Vector3, + Wrench, +) +from rclpy.node import Node +from rclpy.qos import DurabilityPolicy, HistoryPolicy, QoSProfile, ReliabilityPolicy +from sensor_msgs.msg import Image, JointState +from std_msgs.msg import Bool +from std_msgs.msg import Empty as EmptyMsg +from std_msgs.msg import Float32, Header, String +from std_srvs.srv import Empty as EmptySrv +from tf2_ros import TransformBroadcaster + +from space_robotics_bench.core.actions import ( + ManipulatorTaskSpaceActionCfg, + MultiCopterActionGroupCfg, + WheeledRoverActionGroupCfg, +) +from space_robotics_bench.core.envs import BaseEnv +from space_robotics_bench.envs import ( + BaseAerialRoboticsEnv, + BaseManipulationEnv, + BaseMobileRoboticsEnv, +) +from space_robotics_bench.utils.string import canonicalize_str + + +class ROS2: + OBSERVATION_MAPPING = { + "image_scene_rgb": ("camera_scene/image_raw", Image), + "image_scene_depth": ("camera_scene/depth/image_raw", Image), + "image_front_rgb": ("robot/camera_front/image_raw", Image), + "image_front_depth": ("robot/camera_front/depth/image_raw", Image), + "image_bottom_rgb": ("robot/camera_bottom/image_raw", Image), + "image_bottom_depth": ("robot/camera_bottom/depth/image_raw", Image), + "image_base_rgb": ("robot/camera_base/image_raw", Image), + "image_base_depth": ("robot/camera_base/depth/image_raw", Image), + "image_wrist_rgb": ("robot/camera_wrist/image_raw", Image), + "image_wrist_depth": ("robot/camera_wrist/depth/image_raw", Image), + } + + def __init__( + self, + env: Union[ + BaseEnv, + BaseAerialRoboticsEnv, + BaseManipulationEnv, + BaseMobileRoboticsEnv, + ], + node: Optional[object] = None, + force_multienv: bool = True, + ): + self._env = env + self._is_multi_env = force_multienv or self._env.unwrapped.num_envs > 1 + + ## Initialize node + if node is None: + rclpy.init(args=None) + self._node = Node("srb") + else: + self._node = node + + ## Execution queue for actions and services that must be executed in the main thread between environment steps via `update()` + self._exec_queue = Queue() + + ## Clock publisher + self._clock_pub = self._node.create_publisher( + Time, + "clock", + QoSProfile( + reliability=ReliabilityPolicy.RELIABLE, + durability=DurabilityPolicy.TRANSIENT_LOCAL, + history=HistoryPolicy.KEEP_LAST, + depth=1, + ), + ) + self._time = 0.0 + + ## Initialize transform broadcaster + self._tf_broadcaster = TransformBroadcaster(self._node) + + ## Setup interfaces + self._setup_actions() + self._setup_observation() + self._setup_reward() + self._setup_terminated() + self._setup_truncated() + self._setup_misc() + + # Run a thread for listening to device + if node is None: + self._thread = threading.Thread(target=rclpy.spin, args=(self._node,)) + self._thread.daemon = True + self._thread.start() + + def __del__(self): + self._thread.join() + + @property + def actions(self) -> torch.Tensor: + return torch.tensor( + self._actions, + dtype=torch.float32, + device=self._env.unwrapped.unwrapped.device, + ) + + def publish( + self, + observation: Dict[str, torch.Tensor], + reward: torch.Tensor, + terminated: torch.Tensor, + truncated: torch.Tensor, + info: Dict[str, Any], + ): + if not self._pub_observation: + self._setup_observation(observation) + stamp = self._node.get_clock().now().to_msg() + for key, value in observation.items(): + mapping = self.OBSERVATION_MAPPING.get(key) + if mapping is None: + continue + if mapping[0] == "tf": + for i in range(self._env.unwrapped.num_envs): + self._broadcast_tf( + value[i], + mapping, + stamp, + prefix=f"{f'env{i}/' if self._is_multi_env else ''}", + ) + else: + pub = self._pub_observation[key] + for i in range(self._env.unwrapped.num_envs): + pub[i].publish(self._wrap_msg(value[i], mapping, stamp=stamp)) + + for i in range(self._env.unwrapped.num_envs): + self._pub_reward[i].publish(Float32(data=float(reward[i]))) + if terminated[i]: + self._pub_terminated[i].publish(EmptyMsg()) + if truncated[i]: + self._pub_truncated[i].publish(EmptyMsg()) + + if info: + self._pub_info.publish(String(data=json.dumps(info))) + + def reset(self): + self._env.reset() + self._actions = np.zeros( + (self._env.unwrapped.num_envs, self._env.unwrapped.cfg.num_actions) + ) + + def update(self): + self._time += self._env.unwrapped.cfg.sim.dt + self._clock_pub.publish( + Time(sec=int(self._time), nanosec=int((self._time % 1) * 1e9)) + ) + + while not self._exec_queue.empty(): + request, kwargs = self._exec_queue.get() + request(**kwargs) + + def _setup_actions(self): + self._actions = np.zeros( + (self._env.unwrapped.num_envs, self._env.unwrapped.cfg.num_actions) + ) + + robot_name = self._env.unwrapped.cfg.robot_cfg.asset_cfg.prim_path.split("/")[ + -1 + ] + + qos_profile = QoSProfile( + reliability=ReliabilityPolicy.RELIABLE, + durability=DurabilityPolicy.VOLATILE, + history=HistoryPolicy.KEEP_LAST, + depth=1, + ) + + if isinstance(self._env.unwrapped.cfg.actions, ManipulatorTaskSpaceActionCfg): + + def _create_actions_cb_cmd_vel( + cb_name: str, env_id: Optional[int] = None + ) -> Callable: + def cb_single(self, msg: Twist): + self._actions[env_id, :6] = np.array( + [ + msg.linear.x, + msg.linear.y, + msg.linear.z, + msg.angular.x, + msg.angular.y, + msg.angular.z, + ] + ) + + def cb_all(self, msg: Twist): + self._actions[:, :6] = np.array( + [ + msg.linear.x, + msg.linear.y, + msg.linear.z, + msg.angular.x, + msg.angular.y, + msg.angular.z, + ] + ) + + cb = cb_single if env_id else cb_all + cb_name = f"cb_{canonicalize_str(cb_name)}{env_id or ''}" + setattr(self, cb_name, cb.__get__(self, self.__class__)) + return getattr(self, cb_name) + + def _create_actions_cb_gripper( + cb_name: str, env_id: Optional[int] = None + ) -> Callable: + def cb_single(self, msg: Bool): + self._actions[env_id, 6] = -1.0 if msg.data else 1.0 + + def cb_all(self, msg: Bool): + self._actions[:, 6] = -1.0 if msg.data else 1.0 + + cb = cb_single if env_id else cb_all + cb_name = f"cb_{canonicalize_str(cb_name)}{env_id or ''}" + setattr(self, cb_name, cb.__get__(self, self.__class__)) + return getattr(self, cb_name) + + if self._is_multi_env: + self._sub_actions = ( + self._node.create_subscription( + Twist, + f"envs/{robot_name}/cmd_vel", + _create_actions_cb_cmd_vel(cb_name="robot_cmd_vel"), + qos_profile, + ), + self._node.create_subscription( + Bool, + f"envs/{robot_name}/gripper", + _create_actions_cb_gripper(cb_name="robot_gripper"), + qos_profile, + ), + *( + self._node.create_subscription( + Twist, + f"env{i}/{robot_name}/cmd_vel", + _create_actions_cb_cmd_vel("robot_cmd_vel", i), + qos_profile, + ) + for i in range(self._env.unwrapped.num_envs) + ), + *( + self._node.create_subscription( + Bool, + f"env{i}/{robot_name}/gripper", + _create_actions_cb_gripper("robot_gripper", i), + qos_profile, + ) + for i in range(self._env.unwrapped.num_envs) + ), + ) + else: + self._sub_actions = ( + self._node.create_subscription( + Twist, + f"{robot_name}/cmd_vel", + _create_actions_cb_cmd_vel(cb_name="robot_cmd_vel"), + qos_profile, + ), + self._node.create_subscription( + Bool, + f"{robot_name}/gripper", + _create_actions_cb_gripper(cb_name="robot_gripper"), + qos_profile, + ), + ) + + elif isinstance(self._env.unwrapped.cfg.actions, MultiCopterActionGroupCfg): + + def _create_actions_cb_cmd_vel( + cb_name: str, env_id: Optional[int] = None + ) -> Callable: + def cb_single(self, msg: Twist): + self._actions[env_id] = np.array( + [msg.linear.x, msg.linear.y, msg.linear.z, msg.angular.z] + ) + + def cb_all(self, msg: Twist): + self._actions[:] = np.array( + [msg.linear.x, msg.linear.y, msg.linear.z, msg.angular.z] + ) + + cb = cb_single if env_id else cb_all + cb_name = f"cb_{canonicalize_str(cb_name)}{env_id or ''}" + setattr(self, cb_name, cb.__get__(self, self.__class__)) + return getattr(self, cb_name) + + if self._is_multi_env: + self._sub_actions = ( + self._node.create_subscription( + Twist, + f"envs/{robot_name}/cmd_vel", + _create_actions_cb_cmd_vel(cb_name="robot_cmd_vel"), + qos_profile, + ), + *( + self._node.create_subscription( + Twist, + f"env{i}/{robot_name}/cmd_vel", + _create_actions_cb_cmd_vel("robot_cmd_vel", i), + qos_profile, + ) + for i in range(self._env.unwrapped.num_envs) + ), + ) + else: + self._sub_actions = self._node.create_subscription( + Twist, + f"{robot_name}/cmd_vel", + _create_actions_cb_cmd_vel(cb_name="robot_cmd_vel"), + qos_profile, + ) + + elif isinstance(self._env.unwrapped.cfg.actions, WheeledRoverActionGroupCfg): + + def _create_actions_cb_cmd_vel( + cb_name: str, env_id: Optional[int] = None + ) -> Callable: + def cb_single(self, msg: Twist): + self._actions[env_id] = np.array([msg.linear.x, msg.angular.z]) + + def cb_all(self, msg: Twist): + self._actions[:] = np.array([msg.linear.x, msg.angular.z]) + + cb = cb_single if env_id else cb_all + cb_name = f"cb_{canonicalize_str(cb_name)}{env_id or ''}" + setattr(self, cb_name, cb.__get__(self, self.__class__)) + return getattr(self, cb_name) + + if self._is_multi_env: + self._sub_actions = ( + self._node.create_subscription( + Twist, + f"envs/{robot_name}/cmd_vel", + _create_actions_cb_cmd_vel(cb_name="robot_cmd_vel"), + qos_profile, + ), + *( + self._node.create_subscription( + Twist, + f"env{i}/{robot_name}/cmd_vel", + _create_actions_cb_cmd_vel("robot_cmd_vel", i), + qos_profile, + ) + for i in range(self._env.unwrapped.num_envs) + ), + ) + else: + self._sub_actions = self._node.create_subscription( + Twist, + f"{robot_name}/cmd_vel", + _create_actions_cb_cmd_vel(cb_name="robot_cmd_vel"), + qos_profile, + ) + + def _setup_observation(self, observation: Optional[Dict[str, torch.Tensor]] = None): + if observation is None: + self._pub_observation = {} + return + + qos_profile = QoSProfile( + reliability=ReliabilityPolicy.RELIABLE, + durability=DurabilityPolicy.VOLATILE, + history=HistoryPolicy.KEEP_LAST, + depth=1, + ) + for key in observation.keys(): + mapping = self.OBSERVATION_MAPPING.get(key) + if mapping is None: + continue + if mapping[0] == "tf": + continue + self._pub_observation[key] = tuple( + self._node.create_publisher( + mapping[1], + f"{f'env{i}/' if self._is_multi_env else ''}{mapping[0]}", + qos_profile, + ) + for i in range(self._env.unwrapped.num_envs) + ) + + def _broadcast_tf( + self, + tensor: torch.Tensor, + mapping: Tuple[str, type, str, str], + stamp: Time, + prefix: str = "", + ): + _, _, parent_frame_id, child_frame_id = mapping + self._tf_broadcaster.sendTransform( + TransformStamped( + header=Header(frame_id=parent_frame_id, stamp=stamp), + child_frame_id=f"{prefix}{child_frame_id}", + transform=Transform( + translation=Vector3( + x=tensor[0].item(), + y=tensor[1].item(), + z=tensor[2].item(), + ), + rotation=Quaternion( + x=tensor[4].item(), + y=tensor[5].item(), + z=tensor[6].item(), + w=tensor[3].item(), + ), + ), + ) + ) + + def _wrap_msg( + self, + tensor: torch.Tensor, + mapping: Tuple[str, type], + stamp: Time, + ): + topic_name, msg_type = mapping + + if msg_type is JointState: + return JointState( + header=Header(frame_id=topic_name.rsplit("/", 1)[0], stamp=stamp), + name=self._env.unwrapped._robot.joint_names, + position=tensor.cpu().numpy(), + ) + elif msg_type is Wrench: + return Wrench( + force=Vector3( + x=tensor[0].item(), + y=tensor[1].item(), + z=tensor[2].item(), + ), + torque=Vector3( + x=tensor[3].item(), + y=tensor[4].item(), + z=tensor[5].item(), + ), + ) + elif msg_type is Image: + return Image( + header=Header(frame_id=topic_name.rsplit("/", 1)[0], stamp=stamp), + width=tensor.shape[1], + height=tensor.shape[0], + step=tensor.shape[1] * 3, + data=(255.0 * tensor).cpu().numpy().astype(np.uint8).tobytes(), + is_bigendian=0, + encoding="rgb8" if tensor.shape[2] == 3 else "mono8", + ) + else: + return msg_type(*tensor.cpu().numpy()) + + def _setup_reward(self): + qos_profile = QoSProfile( + reliability=ReliabilityPolicy.RELIABLE, + durability=DurabilityPolicy.VOLATILE, + history=HistoryPolicy.KEEP_LAST, + depth=1, + ) + self._pub_reward = tuple( + self._node.create_publisher( + Float32, f"env{i if self._is_multi_env else ''}/reward", qos_profile + ) + for i in range(self._env.unwrapped.num_envs) + ) + + def _setup_terminated(self): + qos_profile = QoSProfile( + reliability=ReliabilityPolicy.RELIABLE, + durability=DurabilityPolicy.TRANSIENT_LOCAL, + history=HistoryPolicy.KEEP_LAST, + depth=1, + ) + self._pub_terminated = tuple( + self._node.create_publisher( + EmptyMsg, + f"env{i if self._is_multi_env else ''}/terminated", + qos_profile, + ) + for i in range(self._env.unwrapped.num_envs) + ) + + def _setup_truncated(self): + qos_profile = QoSProfile( + reliability=ReliabilityPolicy.RELIABLE, + durability=DurabilityPolicy.TRANSIENT_LOCAL, + history=HistoryPolicy.KEEP_LAST, + depth=1, + ) + self._pub_truncated = tuple( + self._node.create_publisher( + EmptyMsg, f"env{i if self._is_multi_env else ''}/truncated", qos_profile + ) + for i in range(self._env.unwrapped.num_envs) + ) + + def _setup_misc(self): + self._srv_reset = self._node.create_service( + EmptySrv, "sim/reset", self._cb_reset + ) + + self._pub_info = self._node.create_publisher( + String, + f"env{'s' if self._is_multi_env else ''}/info", + QoSProfile( + reliability=ReliabilityPolicy.RELIABLE, + durability=DurabilityPolicy.TRANSIENT_LOCAL, + history=HistoryPolicy.KEEP_LAST, + depth=1, + ), + ) + + def _cb_reset(self, request: EmptySrv.Request, response: EmptySrv.Response): + self._exec_queue.put((self.reset, {})) + return response diff --git a/space_robotics_bench/core/managers/__init__.py b/space_robotics_bench/core/managers/__init__.py new file mode 100644 index 0000000..7931436 --- /dev/null +++ b/space_robotics_bench/core/managers/__init__.py @@ -0,0 +1 @@ +from omni.isaac.lab.managers import * # noqa: F403 diff --git a/space_robotics_bench/core/markers/__init__.py b/space_robotics_bench/core/markers/__init__.py new file mode 100644 index 0000000..0920c19 --- /dev/null +++ b/space_robotics_bench/core/markers/__init__.py @@ -0,0 +1,3 @@ +from omni.isaac.lab.markers import * # noqa: F403 + +from .frame_marker import * # noqa: F403 diff --git a/space_robotics_bench/core/markers/frame_marker.py b/space_robotics_bench/core/markers/frame_marker.py new file mode 100644 index 0000000..acbb137 --- /dev/null +++ b/space_robotics_bench/core/markers/frame_marker.py @@ -0,0 +1,4 @@ +from space_robotics_bench.core.markers import FRAME_MARKER_CFG + +FRAME_MARKER_SMALL_CFG = FRAME_MARKER_CFG.copy() +FRAME_MARKER_SMALL_CFG.markers["frame"].scale = (0.025, 0.025, 0.025) diff --git a/space_robotics_bench/core/mdp/__init__.py b/space_robotics_bench/core/mdp/__init__.py new file mode 100644 index 0000000..2e14eb1 --- /dev/null +++ b/space_robotics_bench/core/mdp/__init__.py @@ -0,0 +1,4 @@ +from omni.isaac.lab.envs.mdp import * # noqa: F403 + +from .events import * # noqa: F403 +from .observations import * # noqa: F403 diff --git a/space_robotics_bench/core/mdp/events.py b/space_robotics_bench/core/mdp/events.py new file mode 100644 index 0000000..cca1aa7 --- /dev/null +++ b/space_robotics_bench/core/mdp/events.py @@ -0,0 +1,185 @@ +from typing import TYPE_CHECKING, Dict, List, Tuple + +import torch +from omni.isaac.core.prims.xform_prim_view import XFormPrimView +from omni.isaac.lab.managers import SceneEntityCfg +from pxr import Usd + +import space_robotics_bench.utils.math as math_utils +import space_robotics_bench.utils.sampling as sampling_utils +from space_robotics_bench.core.assets import Articulation, RigidObject + +if TYPE_CHECKING: + from space_robotics_bench.core.envs import BaseEnv + + +def reset_xform_orientation_uniform( + env: "BaseEnv", + env_ids: torch.Tensor, + orientation_distribution_params: Dict[str, Tuple[float, float]], + asset_cfg: SceneEntityCfg = SceneEntityCfg("object"), +) -> Usd.Prim: + asset: XFormPrimView = env.scene[asset_cfg.name] + + range_list = [ + orientation_distribution_params.get(key, (0.0, 0.0)) + for key in ["roll", "pitch", "yaw"] + ] + ranges = torch.tensor(range_list, device=asset._device) + rand_samples = math_utils.sample_uniform( + ranges[:, 0], ranges[:, 1], (1, 3), device=asset._device + ) + + orientations = math_utils.quat_from_euler_xyz( + rand_samples[:, 0], rand_samples[:, 1], rand_samples[:, 2] + ) + + asset.set_world_poses(orientations=orientations) + + +def follow_xform_orientation_linear_trajectory( + env: "BaseEnv", + env_ids: torch.Tensor, + orientation_step_params: Dict[str, float], + asset_cfg: SceneEntityCfg = SceneEntityCfg("object"), +) -> Usd.Prim: + asset: XFormPrimView = env.scene[asset_cfg.name] + + _, current_quat = asset.get_world_poses() + + steps = torch.tensor( + [orientation_step_params.get(key, 0.0) for key in ["roll", "pitch", "yaw"]], + device=asset._device, + ) + step_quat = math_utils.quat_from_euler_xyz(steps[0], steps[1], steps[2]).unsqueeze( + 0 + ) + + orientations = math_utils.quat_mul(current_quat, step_quat) + + asset.set_world_poses(orientations=orientations) + + +def reset_joints_by_offset( + env: "BaseEnv", + env_ids: torch.Tensor, + position_range: tuple[float, float], + velocity_range: tuple[float, float], + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +): + """Reset the robot joints with offsets around the default position and velocity by the given ranges. + + This function samples random values from the given ranges and biases the default joint positions and velocities + by these values. The biased values are then set into the physics simulation. + """ + # Extract the used quantities (to enable type-hinting) + asset: Articulation = env.scene[asset_cfg.name] + + # Get default joint state + joint_pos = asset.data.default_joint_pos[env_ids].clone() + joint_vel = asset.data.default_joint_vel[env_ids].clone() + + # Bias these values randomly + joint_pos += math_utils.sample_uniform( + *position_range, joint_pos.shape, joint_pos.device + ) + joint_vel += math_utils.sample_uniform( + *velocity_range, joint_vel.shape, joint_vel.device + ) + + # Clamp joint pos to limits + joint_pos_limits = asset.data.soft_joint_pos_limits[env_ids] + joint_pos = joint_pos.clamp_(joint_pos_limits[..., 0], joint_pos_limits[..., 1]) + # Clamp joint vel to limits + joint_vel_limits = asset.data.soft_joint_vel_limits[env_ids] + joint_vel = joint_vel.clamp_(-joint_vel_limits, joint_vel_limits) + + # Set into the physics simulation + joint_indices = asset.find_joints(asset_cfg.joint_names)[0] + asset.write_joint_state_to_sim( + joint_pos[:, joint_indices], + joint_vel[:, joint_indices], + joint_ids=joint_indices, + env_ids=env_ids, + ) + + +def reset_root_state_uniform_poisson_disk_2d( + env: "BaseEnv", + env_ids: torch.Tensor, + pose_range: dict[str, tuple[float, float]], + velocity_range: dict[str, tuple[float, float]], + radius: float, + asset_cfgs: List[SceneEntityCfg] = [ + SceneEntityCfg("robot"), + ], +): + # Extract the used quantities (to enable type-hinting) + assets: List[RigidObject | Articulation] = [ + env.scene[asset_cfg.name] for asset_cfg in asset_cfgs + ] + # Get default root state + root_states = torch.stack( + [asset.data.default_root_state[env_ids].clone() for asset in assets], + ).swapaxes(0, 1) + + # Poses + range_list = [ + pose_range.get(key, (0.0, 0.0)) + for key in ["x", "y", "z", "roll", "pitch", "yaw"] + ] + ranges = torch.tensor(range_list, dtype=torch.float32, device=assets[0].device) + samples_pos_xy = torch.tensor( + sampling_utils.sample_poisson_disk_2d_looped( + (len(env_ids), len(asset_cfgs)), + ( + (range_list[0][0], range_list[1][0]), + (range_list[0][1], range_list[1][1]), + ), + radius, + ), + device=assets[0].device, + ) + rand_samples = math_utils.sample_uniform( + ranges[2:, 0], + ranges[2:, 1], + (len(env_ids), len(asset_cfgs), 4), + device=assets[0].device, + ) + rand_samples = torch.cat([samples_pos_xy, rand_samples], dim=-1) + + positions = ( + root_states[:, :, 0:3] + + env.scene.env_origins[env_ids].repeat(len(asset_cfgs), 1, 1).swapaxes(0, 1) + + rand_samples[:, :, 0:3] + ) + orientations_delta = math_utils.quat_from_euler_xyz( + rand_samples[:, :, 3], rand_samples[:, :, 4], rand_samples[:, :, 5] + ) + orientations = math_utils.quat_mul(root_states[:, :, 3:7], orientations_delta) + + # Velocities + range_list = [ + velocity_range.get(key, (0.0, 0.0)) + for key in ["x", "y", "z", "roll", "pitch", "yaw"] + ] + ranges = torch.tensor(range_list, dtype=torch.float32, device=assets[0].device) + rand_samples = math_utils.sample_uniform( + ranges[:, 0], + ranges[:, 1], + (len(env_ids), len(asset_cfgs), 6), + device=assets[0].device, + ) + velocities = root_states[:, :, 7:13] + rand_samples + + # Set into the physics simulation + for asset, position, orientation, velocity in zip( + assets, + positions.unbind(1), + orientations.unbind(1), + velocities.unbind(1), + ): + asset.write_root_pose_to_sim( + torch.cat([position, orientation], dim=-1), env_ids=env_ids + ) + asset.write_root_velocity_to_sim(velocity, env_ids=env_ids) diff --git a/space_robotics_bench/core/mdp/observations.py b/space_robotics_bench/core/mdp/observations.py new file mode 100644 index 0000000..046ee49 --- /dev/null +++ b/space_robotics_bench/core/mdp/observations.py @@ -0,0 +1,24 @@ +from typing import TYPE_CHECKING + +import torch +from omni.isaac.lab.managers import SceneEntityCfg + +from space_robotics_bench.core.assets import Articulation + +if TYPE_CHECKING: + from space_robotics_bench.core.envs import BaseEnv + + +def body_incoming_wrench_mean( + env: "BaseEnv", asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") +) -> torch.Tensor: + """Incoming spatial wrench on bodies of an articulation in the simulation world frame. + + This is the 6-D wrench (force and torque) applied to the body link by the incoming joint force. + """ + + asset: Articulation = env.scene[asset_cfg.name] + link_incoming_forces = asset.root_physx_view.get_link_incoming_joint_force()[ + :, asset_cfg.body_ids + ] + return link_incoming_forces.mean(dim=1) diff --git a/space_robotics_bench/core/sim/__init__.py b/space_robotics_bench/core/sim/__init__.py new file mode 100644 index 0000000..81288d4 --- /dev/null +++ b/space_robotics_bench/core/sim/__init__.py @@ -0,0 +1,4 @@ +from omni.isaac.lab.sim import * # noqa: F403 + +from .schemas import * # noqa: F403 +from .spawners import * # noqa: F403 diff --git a/space_robotics_bench/core/sim/schemas/__init__.py b/space_robotics_bench/core/sim/schemas/__init__.py new file mode 100644 index 0000000..db6752a --- /dev/null +++ b/space_robotics_bench/core/sim/schemas/__init__.py @@ -0,0 +1,3 @@ +from omni.isaac.lab.sim.schemas import * # noqa: F403 + +from .cfg import * # noqa: F403 diff --git a/space_robotics_bench/core/sim/schemas/cfg.py b/space_robotics_bench/core/sim/schemas/cfg.py new file mode 100644 index 0000000..ff4043d --- /dev/null +++ b/space_robotics_bench/core/sim/schemas/cfg.py @@ -0,0 +1,30 @@ +from collections.abc import Callable +from typing import Literal + +from omni.isaac.lab.utils import configclass + +from . import impl + + +@configclass +class MeshCollisionPropertiesCfg: + func: Callable = impl.set_mesh_collision_properties + + mesh_approximation: ( + Literal[ + "none", + "convexHull", + "convexDecomposition", + "meshSimplification", + "convexMeshSimplification", + "boundingCube", + "boundingSphere", + "sphereFill", + "sdf", + ] + | None + ) = None + """Collision approximation to use for the collision shape.""" + + sdf_resolution: int = 160 + """Resolution of the SDF grid used for collision approximation.""" diff --git a/space_robotics_bench/core/sim/schemas/impl.py b/space_robotics_bench/core/sim/schemas/impl.py new file mode 100644 index 0000000..2f8a6a2 --- /dev/null +++ b/space_robotics_bench/core/sim/schemas/impl.py @@ -0,0 +1,77 @@ +from typing import TYPE_CHECKING + +import carb +import omni.isaac.core.utils.stage as stage_utils +from omni.physx.scripts import utils as physx_utils +from pxr import PhysxSchema, Usd, UsdPhysics + +from space_robotics_bench.core.sim import apply_nested + +if TYPE_CHECKING: + from . import cfg + + +@apply_nested +def set_mesh_collision_properties( + prim_path: str, + cfg: "cfg.MeshCollisionPropertiesCfg", + stage: Usd.Stage | None = None, +) -> bool: + """ + + Args: + prim_path: The prim path of parent. + cfg: The configuration for the collider. + stage: The stage where to find the prim. Defaults to None, in which case the + current stage is used. + + Returns: + True if the properties were successfully set, False otherwise. + """ + + # Apply mesh collision approximation + if cfg.mesh_approximation is not None: + if stage is None: + stage: Usd.Stage = stage_utils.get_current_stage() + prim: Usd.Prim = stage.GetPrimAtPath(prim_path) + + if physx_utils.hasSchema(prim, "CollisionAPI"): + carb.log_warn("CollisionAPI is already defined") + return + + def isPartOfRigidBody(currPrim): + if currPrim.HasAPI(UsdPhysics.RigidBodyAPI): + return True + + currPrim = currPrim.GetParent() + + return isPartOfRigidBody(currPrim) if currPrim.IsValid() else False + + if cfg.mesh_approximation == "none" and isPartOfRigidBody(prim): + carb.log_warn( + f"setCollider: {prim.GetPath()} is a part of a rigid body. Resetting approximation shape from none (trimesh) to convexHull" + ) + cfg.mesh_approximation = "convexHull" + + collisionAPI = UsdPhysics.CollisionAPI.Apply(prim) + PhysxSchema.PhysxCollisionAPI.Apply(prim) + collisionAPI.CreateCollisionEnabledAttr().Set(True) + + api = physx_utils.MESH_APPROXIMATIONS.get( + cfg.mesh_approximation, 0 + ) # None is a valid value + if api == 0: + carb.log_warn( + f"setCollider: invalid approximation type {cfg.mesh_approximation} provided for {prim.GetPath()}. Falling back to convexHull." + ) + cfg.mesh_approximation = "convexHull" + api = physx_utils.MESH_APPROXIMATIONS[cfg.mesh_approximation] + approximation_api = api.Apply(prim) if api is not None else None + if cfg.mesh_approximation == "sdf" and cfg.sdf_resolution: + approximation_api.CreateSdfResolutionAttr().Set(cfg.sdf_resolution) + + meshcollisionAPI = UsdPhysics.MeshCollisionAPI.Apply(prim) + + meshcollisionAPI.CreateApproximationAttr().Set(cfg.mesh_approximation) + + return True diff --git a/space_robotics_bench/core/sim/spawners/__init__.py b/space_robotics_bench/core/sim/spawners/__init__.py new file mode 100644 index 0000000..0156b63 --- /dev/null +++ b/space_robotics_bench/core/sim/spawners/__init__.py @@ -0,0 +1,5 @@ +from omni.isaac.lab.sim.spawners import * # noqa: F403 + +from .from_files import * # noqa: F403 +from .multi import * # noqa: F403 +from .procgen import * # noqa: F403 diff --git a/space_robotics_bench/core/sim/spawners/from_files/__init__.py b/space_robotics_bench/core/sim/spawners/from_files/__init__.py new file mode 100644 index 0000000..76c44b3 --- /dev/null +++ b/space_robotics_bench/core/sim/spawners/from_files/__init__.py @@ -0,0 +1,3 @@ +from omni.isaac.lab.sim.spawners.from_files import * # noqa: F403 + +from .cfg import * # noqa: F403 diff --git a/space_robotics_bench/core/sim/spawners/from_files/cfg.py b/space_robotics_bench/core/sim/spawners/from_files/cfg.py new file mode 100644 index 0000000..862c3b8 --- /dev/null +++ b/space_robotics_bench/core/sim/spawners/from_files/cfg.py @@ -0,0 +1,19 @@ +from collections.abc import Callable + +from omni.isaac.lab.sim import UsdFileCfg as __UsdFileCfg +from omni.isaac.lab.utils import configclass + +from space_robotics_bench.core.sim.schemas import MeshCollisionPropertiesCfg + +from . import impl + + +@configclass +class UsdFileCfg(__UsdFileCfg): + """ + Extended version of :class:`omni.isaac.lab.sim.UsdFileCfg`. + """ + + func: Callable = impl.spawn_from_usd + + mesh_collision_props: MeshCollisionPropertiesCfg | None = None diff --git a/space_robotics_bench/core/sim/spawners/from_files/impl.py b/space_robotics_bench/core/sim/spawners/from_files/impl.py new file mode 100644 index 0000000..3171695 --- /dev/null +++ b/space_robotics_bench/core/sim/spawners/from_files/impl.py @@ -0,0 +1,80 @@ +from typing import TYPE_CHECKING, Tuple + +import omni.isaac.core.utils.prims as prim_utils +import omni.isaac.core.utils.stage as stage_utils +from pxr import PhysxSchema, Usd, UsdPhysics + +from space_robotics_bench.core.sim import clone + +if TYPE_CHECKING: + from . import cfg + +from space_robotics_bench.core.sim.spawners.from_files import ( + spawn_from_usd as __spawn_from_usd, +) + + +@clone +def spawn_from_usd( + prim_path: str, + cfg: "cfg.UsdFileCfg", + translation: Tuple[float, float, float] | None = None, + orientation: Tuple[float, float, float, float] | None = None, +) -> Usd.Prim: + """ + Extended version of :class:`omni.isaac.lab.sim.spawners.from_files.spawn_from_usd`. + """ + + # Get stage + stage: Usd.Stage = stage_utils.get_current_stage() + if not stage.ResolveIdentifierToEditTarget(cfg.usd_path): + raise FileNotFoundError(f"USD file not found at path: '{cfg.usd_path}'.") + + # Create prim if it doesn't exist + if not prim_utils.is_prim_path_valid(prim_path): + prim_utils.create_prim( + prim_path, + usd_path=cfg.usd_path, + translation=translation, + orientation=orientation, + scale=cfg.scale, + ) + # Get prim + prim: Usd.Prim = stage.GetPrimAtPath(prim_path) + + # Define missing APIs + _define_missing_apis(prim, cfg) + + # Apply mesh collision properties + if cfg.mesh_collision_props is not None: + cfg.mesh_collision_props.func(prim_path, cfg.mesh_collision_props, stage) + + return __spawn_from_usd(prim_path, cfg, translation, orientation) + + +def _define_missing_apis(prim: Usd.Prim, cfg: "cfg.UsdFileCfg"): + if cfg.rigid_props is not None and not UsdPhysics.RigidBodyAPI(prim): + UsdPhysics.RigidBodyAPI.Apply(prim) + + if cfg.collision_props is not None and not UsdPhysics.CollisionAPI(prim): + UsdPhysics.CollisionAPI.Apply(prim) + + if cfg.mass_props is not None and not UsdPhysics.MassAPI(prim): + UsdPhysics.MassAPI.Apply(prim) + + if cfg.articulation_props is not None and not UsdPhysics.ArticulationRootAPI(prim): + UsdPhysics.ArticulationRootAPI.Apply(prim) + + if cfg.fixed_tendons_props is not None: + if not PhysxSchema.PhysxTendonAxisAPI(prim): + PhysxSchema.PhysxTendonAxisAPI.Apply(prim) + if not PhysxSchema.PhysxTendonAxisRootAPI(prim): + PhysxSchema.PhysxTendonAxisRootAPI.Apply(prim) + + if cfg.joint_drive_props is not None and not UsdPhysics.DriveAPI(prim): + UsdPhysics.DriveAPI.Apply(prim) + + if cfg.deformable_props is not None and not PhysxSchema.PhysxDeformableBodyAPI( + prim + ): + PhysxSchema.PhysxDeformableBodyAPI.Apply(prim) diff --git a/space_robotics_bench/core/sim/spawners/multi/__init__.py b/space_robotics_bench/core/sim/spawners/multi/__init__.py new file mode 100644 index 0000000..02f6a2c --- /dev/null +++ b/space_robotics_bench/core/sim/spawners/multi/__init__.py @@ -0,0 +1 @@ +from .cfg import * # noqa: F403 diff --git a/space_robotics_bench/core/sim/spawners/multi/cfg.py b/space_robotics_bench/core/sim/spawners/multi/cfg.py new file mode 100644 index 0000000..d4641a4 --- /dev/null +++ b/space_robotics_bench/core/sim/spawners/multi/cfg.py @@ -0,0 +1,76 @@ +from dataclasses import MISSING +from dataclasses import MISSING as DELAYED_CFG +from typing import Callable, List, Literal, Optional, Tuple + +from omni.isaac.lab.utils import configclass +from pxr import Usd + +import space_robotics_bench.core.sim as sim_utils + +from . import impl + + +@configclass +class MultiAssetCfg(sim_utils.SpawnerCfg): + """Configuration parameters for loading multiple assets randomly""" + + func: Callable[..., Usd.Prim] = impl.spawn_multi_asset_sequential + + assets_cfg: List[sim_utils.SpawnerCfg] = MISSING + """List of asset configurations to spawn""" + + +@configclass +class MultiShapeCfg(MultiAssetCfg): + assets_cfg: List[sim_utils.SpawnerCfg] = DELAYED_CFG + + size: Tuple[float, float, float] = MISSING + """Size of cuboid""" + + radius: Optional[float] = None + """Radius of sphere|cylinder|capsule|cone (default: self.size[0])""" + + height: Optional[float] = None + """Height of cylinder|capsule|cone (default: self.size[1])""" + + axis: Literal["X", "Y", "Z"] = "Z" + """Axis of cylinder|capsule|cone""" + + shape_cfg: sim_utils.ShapeCfg = sim_utils.ShapeCfg() + """Additional configuration applied to all shapes""" + + def __post_init__(self): + if self.radius is None: + self.radius = self.size[0] + if self.height is None: + self.height = self.size[1] + + # Extract ShapeCfg kwargs while ignoring private attributes and func + shape_cfg_kwargs = { + k: v + for k, v in self.shape_cfg.__dict__.items() + if not k.startswith("_") and k != "func" + } + + self.assets_cfg = [ + sim_utils.CuboidCfg(size=self.size, **shape_cfg_kwargs), + sim_utils.SphereCfg(radius=self.radius, **shape_cfg_kwargs), + sim_utils.CylinderCfg( + radius=self.radius, + height=self.height, + axis=self.axis, + **shape_cfg_kwargs, + ), + sim_utils.CapsuleCfg( + radius=self.radius, + height=self.height, + axis=self.axis, + **shape_cfg_kwargs, + ), + sim_utils.ConeCfg( + radius=self.radius, + height=self.height, + axis=self.axis, + **shape_cfg_kwargs, + ), + ] diff --git a/space_robotics_bench/core/sim/spawners/multi/impl.py b/space_robotics_bench/core/sim/spawners/multi/impl.py new file mode 100644 index 0000000..47aa81e --- /dev/null +++ b/space_robotics_bench/core/sim/spawners/multi/impl.py @@ -0,0 +1,152 @@ +import re +from typing import TYPE_CHECKING, Tuple + +import omni.isaac.core.utils.prims as prim_utils +import omni.usd +from pxr import Gf, Sdf, Semantics, Usd, UsdGeom, Vt + +import space_robotics_bench.core.sim as sim_utils + +if TYPE_CHECKING: + from . import cfg + + +def spawn_multi_asset_sequential( + prim_path: str, + cfg: "cfg.MultiAssetCfg", + translation: Tuple[float, float, float] | None = None, + orientation: Tuple[float, float, float, float] | None = None, +) -> Usd.Prim: + # Check if prim path contains a regex expression + # Note: A valid prim path can only contain alphanumeric characters, underscores, and forward slashes + root_path, asset_path = prim_path.rsplit("/", 1) + is_regex_expression = re.match(r"^[a-zA-Z0-9/_]+$", root_path) is None + + # Resolve matching prims for source prim path expression + if is_regex_expression and root_path != "": + source_prim_paths = sim_utils.find_matching_prim_paths(root_path) + if len(source_prim_paths) == 0: + raise RuntimeError( + f"Unable to find source prim path: '{root_path}'. Please create the prim before spawning." + ) + else: + source_prim_paths = [root_path] + + # Spawn everything first in a dataset prim + proto_root_prim_path = f"/World/db_{asset_path}" + prim_utils.create_prim(proto_root_prim_path, "Scope") + + proto_prim_paths = [] + n_assets = len(cfg.assets_cfg) + zfill_len = len(str(n_assets)) + for index, asset_cfg in enumerate(cfg.assets_cfg): + # Save the proto prim path + proto_prim_path = ( + f"{proto_root_prim_path}/{asset_path}_{str(index).zfill(zfill_len)}" + ) + proto_prim_paths.append(proto_prim_path) + + # Spawn single instance + prim = asset_cfg.func(proto_prim_path, asset_cfg) + + # Set the prim visibility + if hasattr(asset_cfg, "visible"): + imageable = UsdGeom.Imageable(prim) + if asset_cfg.visible: + imageable.MakeVisible() + else: + imageable.MakeInvisible() + + # Set the semantic annotations + if hasattr(asset_cfg, "semantic_tags") and asset_cfg.semantic_tags is not None: + # Note: Taken from replicator scripts.utils.utils.py + for semantic_type, semantic_value in asset_cfg.semantic_tags: + # Sanitize by replacing spaces with underscores + semantic_type_sanitized = semantic_type.replace(" ", "_") + semantic_value_sanitized = semantic_value.replace(" ", "_") + # Set the semantic API for the instance + instance_name = f"{semantic_type_sanitized}_{semantic_value_sanitized}" + sem = Semantics.SemanticsAPI.Apply(prim, instance_name) + # Create semantic type and data attributes + sem.CreateSemanticTypeAttr() + sem.CreateSemanticDataAttr() + sem.GetSemanticTypeAttr().Set(semantic_type) + sem.GetSemanticDataAttr().Set(semantic_value) + + # Activate rigid body contact sensors + if ( + hasattr(asset_cfg, "activate_contact_sensors") + and asset_cfg.activate_contact_sensors + ): + sim_utils.activate_contact_sensors( + proto_prim_path, asset_cfg.activate_contact_sensors + ) + + # Acquire stage + stage = omni.usd.get_context().get_stage() + + # Resolve prim paths for spawning and cloning + prim_paths = [ + f"{source_prim_path}/{asset_path}" for source_prim_path in source_prim_paths + ] + + # Convert orientation ordering (wxyz to xyzw) + orientation = (orientation[1], orientation[2], orientation[3], orientation[0]) + + # manually clone prims if the source prim path is a regex expression + with Sdf.ChangeBlock(): + for i, prim_path in enumerate(prim_paths): + # Spawn single instance + env_spec = Sdf.CreatePrimInLayer(stage.GetRootLayer(), prim_path) + + # Select assets in order for uniform distribution + proto_path = proto_prim_paths[i % n_assets] + + # Copy spec from the proto prim + Sdf.CopySpec( + env_spec.layer, + Sdf.Path(proto_path), + env_spec.layer, + Sdf.Path(prim_path), + ) + + ## Set the XformOp for the prim + _ = UsdGeom.Xform(stage.GetPrimAtPath(proto_path)).GetPrim().GetPrimStack() + # Set the order + op_order_spec = env_spec.GetAttributeAtPath(f"{prim_path}.xformOpOrder") + if op_order_spec is None: + op_order_spec = Sdf.AttributeSpec( + env_spec, UsdGeom.Tokens.xformOpOrder, Sdf.ValueTypeNames.TokenArray + ) + op_order_spec.default = Vt.TokenArray( + ["xformOp:translate", "xformOp:orient", "xformOp:scale"] + ) + # Translation + translate_spec = env_spec.GetAttributeAtPath( + f"{prim_path}.xformOp:translate" + ) + if translate_spec is None: + translate_spec = Sdf.AttributeSpec( + env_spec, "xformOp:translate", Sdf.ValueTypeNames.Double3 + ) + translate_spec.default = Gf.Vec3d(*translation) + # Orientation + orient_spec = env_spec.GetAttributeAtPath(f"{prim_path}.xformOp:orient") + if orient_spec is None: + orient_spec = Sdf.AttributeSpec( + env_spec, "xformOp:orient", Sdf.ValueTypeNames.Quatd + ) + orient_spec.default = Gf.Quatd(*orientation) + # Scale + scale_spec = env_spec.GetAttributeAtPath(f"{prim_path}.xformOp:scale") + if scale_spec is None: + scale_spec = Sdf.AttributeSpec( + env_spec, "xformOp:scale", Sdf.ValueTypeNames.Double3 + ) + scale_spec.default = Gf.Vec3d(1.0, 1.0, 1.0) + + # Delete the dataset prim after spawning + prim_utils.delete_prim(proto_root_prim_path) + + # Return the prim + return prim_utils.get_prim_at_path(prim_paths[0]) diff --git a/space_robotics_bench/core/sim/spawners/procgen/__init__.py b/space_robotics_bench/core/sim/spawners/procgen/__init__.py new file mode 100644 index 0000000..02f6a2c --- /dev/null +++ b/space_robotics_bench/core/sim/spawners/procgen/__init__.py @@ -0,0 +1 @@ +from .cfg import * # noqa: F403 diff --git a/space_robotics_bench/core/sim/spawners/procgen/cfg.py b/space_robotics_bench/core/sim/spawners/procgen/cfg.py new file mode 100644 index 0000000..832a5e2 --- /dev/null +++ b/space_robotics_bench/core/sim/spawners/procgen/cfg.py @@ -0,0 +1,88 @@ +from dataclasses import MISSING +from os import path +from typing import Any, Callable, Dict, List, Optional + +import platformdirs +from omni.isaac.lab.utils import configclass +from pxr import Usd + +import space_robotics_bench.core.sim as sim_utils +from space_robotics_bench.paths import SRB_SCRIPTS_DIR + +from . import impl + + +@configclass +class ProceduralAssetCfg(sim_utils.SpawnerCfg): + """ + Configuration for generating procedural assets. + """ + + func: Callable[..., Usd.Prim] = MISSING + + seed: int = 0 + """ + Initial seed used for generating the assets. + """ + + num_assets: int = 4 + """ + Number of assets to generate. + """ + + cache_dir: str = platformdirs.user_cache_dir("srb") + """ + Directory to cache the generated assets. + """ + + overwrite_min_age: int = -1 + """ + Number of seconds after which to overwrite the cached assets. Disabled if negative. + """ + + +@configclass +class BlenderNodesAssetCfg(ProceduralAssetCfg): + """ + Configuration for generating assets using Blender (geometry) nodes. + """ + + func: Callable[..., Usd.Prim] = impl.spawn_blender_procgen_assets + + blender_bin: str = "blender" + blender_args: List[str] = [ + "--factory-startup", + "--background", + "--offline-mode", + "--enable-autoexec", + "--python-exit-code", + "1", + ] + bpy_script: str = path.join(SRB_SCRIPTS_DIR, "blender", "procgen_assets.py") + + # Output + name: str = MISSING + autorun_scripts: List[str] = MISSING + + # Export + ext: str = ".usdz" + export_kwargs: Dict[str, Any] = {} + + # Geometry + geometry_nodes: Dict[str, Dict[str, Any]] = MISSING + decimate_angle_limit: Optional[float] = None + decimate_face_count: Optional[int] = None + + # Material + material: Optional[str] = None + texture_resolution: int = 1024 + detail: float = 1.0 + + # Prim + usd_file_cfg: sim_utils.UsdFileCfg = sim_utils.UsdFileCfg(usd_path="IGNORED") + + def __post_init__(self): + super().__post_init__() + + if not self.ext.startswith("."): + self.ext = f".{self.ext}" diff --git a/space_robotics_bench/core/sim/spawners/procgen/impl.py b/space_robotics_bench/core/sim/spawners/procgen/impl.py new file mode 100644 index 0000000..1a2de6d --- /dev/null +++ b/space_robotics_bench/core/sim/spawners/procgen/impl.py @@ -0,0 +1,163 @@ +import hashlib +import json +import math +import subprocess +import sys +import time +from os import path +from typing import TYPE_CHECKING, Any, Dict, Iterable, Tuple + +from pxr import Usd + +import space_robotics_bench.core.sim as sim_utils +from space_robotics_bench.core.sim.spawners.multi import MultiAssetCfg + +if TYPE_CHECKING: + from . import cfg + + +def spawn_blender_procgen_assets( + prim_path: str, + cfg: "cfg.BlenderNodesAssetCfg", + translation: Tuple[float, float, float] | None = None, + orientation: Tuple[float, float, float, float] | None = None, +) -> Usd.Prim: + # Scale the texture resolution by its factor + if cfg.detail == 0: + cfg.material = None + cfg.texture_resolution = 0 + else: + cfg.texture_resolution = max(16, math.ceil(cfg.detail * cfg.texture_resolution)) + if cfg.texture_resolution % 2 != 0: + cfg.texture_resolution += 1 + + # Extract configuration into script kwargs + script_kwargs = { + "autorun_scripts": cfg.autorun_scripts, + "name": cfg.name, + "ext": cfg.ext, + "overwrite_min_age": cfg.overwrite_min_age, + "seed": cfg.seed, + "num_assets": cfg.num_assets, + "export_kwargs": cfg.export_kwargs, + "geometry_nodes": cfg.geometry_nodes, + "decimate_angle_limit": cfg.decimate_angle_limit, + "decimate_face_count": cfg.decimate_face_count, + "material": cfg.material, + "texture_resolution": cfg.texture_resolution, + } + + # Derive the output directory based on the configuration + outdir = path.join( + cfg.cache_dir, + cfg.name, + _hash_filtered_dict( + script_kwargs, + filter_keys=[ + "autorun_scripts", + "name", + "ext", + "overwrite_min_age", + "seed", + "num_assets", + ], + ), + ) + script_kwargs["outdir"] = outdir + print(f"[TRACE] Caching procedural assets to '{outdir}'") + + # Skip generation if all the files already exist and they are too recent + skip_generation = True + for current_seed in range(cfg.seed, cfg.seed + cfg.num_assets): + filepath = path.join(outdir, f"{cfg.name}{current_seed}{cfg.ext}") + if path.exists(filepath) and ( + cfg.overwrite_min_age < 0 + or (cfg.overwrite_min_age > time.time() - path.getmtime(filepath)) + ): + continue + else: + skip_generation = False + break + + if skip_generation: + print( + f"[INFO] {cfg.num_assets} procedural '{cfg.name}' asset(s) for '{prim_path}' are already cached. Skipping generation." + ) + else: + # Convert kwargs to args + script_args = [] + for key, value in script_kwargs.items(): + if value: + if value is not False: + script_args.append(f"--{key}") + + if isinstance(value, str): + script_args.append(value) + elif isinstance(value, Dict): + script_args.append(json.dumps(value)) + elif isinstance(value, Iterable): + script_args.extend(str(v) for v in value) + elif value is None or value is True: + continue + else: + script_args.append(str(value)) + + # Run the Blender script + print( + f"[INFO] Generating {cfg.num_assets} procedural '{cfg.name}' asset(s) for '{prim_path}' (this might take a while)" + ) + cmd = [ + cfg.blender_bin, + *cfg.blender_args, + "--python", + cfg.bpy_script, + "--", + *script_args, + ] + print(f"[TRACE] Call: {' '.join(cmd)}") + result = subprocess.run( + cmd, + capture_output=True, + text=True, + ) + if result.returncode != 0: + print(result.stdout, file=sys.stdout) + print(result.stderr, file=sys.stderr) + raise RuntimeError( + f"Failed to generate procedural assets using Blender: {' '.join(result.args)}" + ) + + # Leverage MultiAssetCfg to spawn the generated assets + multi_asset_cfg = MultiAssetCfg( + assets_cfg=[ + sim_utils.UsdFileCfg( + usd_path=path.join( + outdir, + asset_basename, + ), + **{ + k: v + for k, v in cfg.usd_file_cfg.__dict__.items() + if not k.startswith("_") and k not in ["func", "usd_path"] + }, + ) + for asset_basename in [ + f"{cfg.name}{current_seed}{cfg.ext}" + for current_seed in range(cfg.seed, cfg.seed + cfg.num_assets) + ] + ] + ) + return multi_asset_cfg.func( + prim_path=prim_path, + cfg=multi_asset_cfg, + translation=translation, + orientation=orientation, + ) + + +def _hash_filtered_dict(input: Dict[str, Any], filter_keys: Iterable[str] = []) -> str: + dhash = hashlib.md5() + filtered = {k: v for k, v in input.items() if k not in filter_keys} + encoded = json.dumps(filtered, sort_keys=True).encode() + dhash.update(encoded) + return dhash.hexdigest() diff --git a/space_robotics_bench/core/teleop_devices/__init__.py b/space_robotics_bench/core/teleop_devices/__init__.py new file mode 100644 index 0000000..2205efd --- /dev/null +++ b/space_robotics_bench/core/teleop_devices/__init__.py @@ -0,0 +1,8 @@ +from omni.isaac.lab.devices import * # noqa: F403 + +from .se3_keyboard import Se3Keyboard # noqa: F401 +from .se3_ros2 import Se3ROS2 # noqa: F401 +from .se3_spacemouse import Se3SpaceMouse # noqa: F401 +from .se3_touch import Se3Touch # noqa: F401 + +from .combined import CombinedInterface # noqa: F401 isort:skip diff --git a/space_robotics_bench/core/teleop_devices/combined.py b/space_robotics_bench/core/teleop_devices/combined.py new file mode 100644 index 0000000..d400751 --- /dev/null +++ b/space_robotics_bench/core/teleop_devices/combined.py @@ -0,0 +1,236 @@ +from collections.abc import Callable +from typing import List, Optional, Union + +import numpy as np +from omni.isaac.lab.devices import DeviceBase + +from space_robotics_bench.core.actions import ( + ManipulatorTaskSpaceActionCfg, + MultiCopterActionGroupCfg, + WheeledRoverActionGroupCfg, +) +from space_robotics_bench.core.teleop_devices import ( + Se3Gamepad, + Se3Keyboard, + Se3ROS2, + Se3SpaceMouse, + Se3Touch, +) +from space_robotics_bench.utils.ros import enable_ros2_bridge + + +class CombinedInterface(DeviceBase): + def __init__( + self, + devices: List[str], + pos_sensitivity: float = 1.0, + rot_sensitivity: float = 1.0, + action_cfg: Optional[ + Union[ + ManipulatorTaskSpaceActionCfg, + MultiCopterActionGroupCfg, + WheeledRoverActionGroupCfg, + ] + ] = None, + node: Optional[object] = None, + ): + enable_ros2_bridge() + import rclpy + from rclpy.node import Node + + if node is None: + rclpy.init(args=None) + self._node = Node("srb_teleop_combined") + else: + self._node = node + + self._action_cfg = action_cfg + self.interfaces = [] + for device in devices: + if device.lower() == "keyboard": + self.interfaces.append( + Se3Keyboard( + pos_sensitivity=0.05 * pos_sensitivity, + rot_sensitivity=10.0 * rot_sensitivity, + ) + ) + elif device.lower() == "ros2": + self.interfaces.append( + Se3ROS2( + node=self._node, + pos_sensitivity=1.0 * pos_sensitivity, + rot_sensitivity=1.0 * rot_sensitivity, + ) + ) + elif device.lower() == "touch": + self.interfaces.append( + Se3Touch( + node=self._node, + pos_sensitivity=1.0 * pos_sensitivity, + rot_sensitivity=0.15 * rot_sensitivity, + ) + ) + elif device.lower() == "spacemouse": + self.interfaces.append( + Se3SpaceMouse( + pos_sensitivity=0.1 * pos_sensitivity, + rot_sensitivity=0.05 * rot_sensitivity, + ) + ) + elif device.lower() == "gamepad": + self.interfaces.append( + Se3Gamepad( + pos_sensitivity=0.1 * pos_sensitivity, + rot_sensitivity=0.1 * rot_sensitivity, + ) + ) + else: + raise ValueError(f"Invalid device interface '{device}'.") + + self.gain = 1.0 + + def cb_gain_decrease(): + self.gain *= 0.75 + print(f"Gain: {self.gain}") + + self.add_callback("O", cb_gain_decrease) + + def cb_gain_increase(): + self.gain *= 1.25 + print(f"Gain: {self.gain}") + + self.add_callback("P", cb_gain_increase) + + def __del__(self): + for interface in self.interfaces: + interface.__del__() + + def __str__(self) -> str: + msg = "Combined Interface\n" + msg += f"Devices: {', '.join([interface.__class__.__name__ for interface in self.interfaces])}\n" + + for interface in self.interfaces: + if isinstance(interface, Se3Keyboard) and self._action_cfg is not None: + msg += self._keyboard_control_scheme(self._action_cfg) + continue + msg += "\n" + msg += interface.__str__() + + return msg + + """ + Operations + """ + + def reset(self): + for interface in self.interfaces: + interface.reset() + + self._close_gripper = False + self._prev_gripper_cmds = [False] * len(self.interfaces) + + def add_callback(self, key: str, func: Callable): + for interface in self.interfaces: + if isinstance(interface, Se3Keyboard): + interface.add_callback(key=key, func=func) + if isinstance(interface, Se3SpaceMouse) and key in ["L", "R", "LR"]: + interface.add_callback(key=key, func=func) + + def advance(self) -> tuple[np.ndarray, bool]: + raw_actions = [interface.advance() for interface in self.interfaces] + + twist = self.gain * np.sum( + np.stack([a[0] for a in raw_actions], axis=0), axis=0 + ) + + for i, prev_gripper_cmd in enumerate(self._prev_gripper_cmds): + if prev_gripper_cmd != raw_actions[i][1]: + self._close_gripper = not self._close_gripper + break + self._prev_gripper_cmds = [a[1] for a in raw_actions] + + return twist, self._close_gripper + + def set_ft_feedback(self, ft_feedback: np.ndarray): + for interface in self.interfaces: + if isinstance(interface, Se3Touch): + interface.set_ft_feedback(ft_feedback) + + @staticmethod + def _keyboard_control_scheme( + action_cfg: Union[ + ManipulatorTaskSpaceActionCfg, + MultiCopterActionGroupCfg, + WheeledRoverActionGroupCfg, + ], + ) -> str: + if isinstance(action_cfg, ManipulatorTaskSpaceActionCfg): + return """ ++------------------------------------------------+ +| Keyboard Scheme (focus the Isaac Sim window) | ++------------------------------------------------+ ++------------------------------------------------+ +| Reset: [ L ] | +| Decrease Gain [ O ] | Increase Gain: [ P ] | +| Toggle Gripper: [ R / K ] | ++------------------------------------------------+ +| Translation | +| [ W ] (+X) [ Q ] (+Z) | +| ↑ ↑ | +| | | | +| (-Y) [ A ] ← + → [ D ] (+Y) + | +| | | | +| ↓ ↓ | +| [ S ] (-X) [ E ] (-Z) | +|------------------------------------------------| +| Rotation | +| [ Z ] ←--------(±X)--------→ [ X ] | +| | +| [ T ] ↻--------(±Y)--------↺ [ G ] | +| | +| [ C ] ↺--------(±Z)--------↻ [ V ] | ++------------------------------------------------+ + """ + elif isinstance(action_cfg, MultiCopterActionGroupCfg): + return """ ++------------------------------------------------+ +| Keyboard Scheme (focus the Isaac Sim window) | ++------------------------------------------------+ ++------------------------------------------------+ +| Decrease Gain [ O ] | Increase Gain: [ P ] | +| Reset: [ L ] | ++------------------------------------------------+ +| Translation | +| [ W ] (+X) [ Q ] (+Z) | +| ↑ ↑ | +| | | | +| (-Y) [ A ] ← + → [ D ] (+Y) + | +| | | | +| ↓ ↓ | +| [ S ] (-X) [ E ] (-Z) | +|------------------------------------------------| +| Rotation | +| [ C ] ↺--------(±Z)--------↻ [ V ] | ++------------------------------------------------+ + """ + elif isinstance(action_cfg, WheeledRoverActionGroupCfg): + return """ ++------------------------------------------------+ +| Keyboard Scheme (focus the Isaac Sim window) | ++------------------------------------------------+ ++------------------------------------------------+ +| Decrease Gain [ O ] | Increase Gain: [ P ] | +| Reset: [ L ] | ++------------------------------------------------+ +| Planar Motion | +| [ W ] (+X) | +| ↑ | +| | | +| (-Y) [ A ] ← + → [ D ] (+Y) | +| | | +| ↓ | +| [ S ] (-X) | ++------------------------------------------------+ + """ + else: + raise ValueError(f"Invalid action configuration '{action_cfg}'.") diff --git a/space_robotics_bench/core/teleop_devices/se3_keyboard.py b/space_robotics_bench/core/teleop_devices/se3_keyboard.py new file mode 100644 index 0000000..e82d083 --- /dev/null +++ b/space_robotics_bench/core/teleop_devices/se3_keyboard.py @@ -0,0 +1,22 @@ +import carb +from omni.isaac.lab.devices import Se3Keyboard as __Se3Keyboard + + +class Se3Keyboard(__Se3Keyboard): + def __str__(self) -> str: + msg = super().__str__() + + msg += "\n" + msg += "\t----------------------------------------------\n" + msg += "\tAdditional controls:\n" + msg += "\tToggle gripper (alternative): R\n" + return msg + + def _on_keyboard_event(self, event, *args, **kwargs) -> bool: + ret = super()._on_keyboard_event(event, *args, **kwargs) + + if event.type == carb.input.KeyboardEventType.KEY_PRESS: + if event.input.name == "R": + self._close_gripper = not self._close_gripper + + return ret diff --git a/space_robotics_bench/core/teleop_devices/se3_ros2.py b/space_robotics_bench/core/teleop_devices/se3_ros2.py new file mode 100644 index 0000000..2f0b8c6 --- /dev/null +++ b/space_robotics_bench/core/teleop_devices/se3_ros2.py @@ -0,0 +1,130 @@ +import os +import threading +from collections.abc import Callable +from typing import Optional + +import numpy as np +from rclpy.node import Node + +from space_robotics_bench.core.teleop_devices import DeviceBase +from space_robotics_bench.utils.ros import enable_ros2_bridge + + +class Se3ROS2(DeviceBase, Node): + def __init__( + self, + node: Optional[object] = None, + pos_sensitivity: float = 1.0, + rot_sensitivity: float = 1.0, + ): + enable_ros2_bridge() + import rclpy + from geometry_msgs.msg import Twist + from rclpy.node import Node + from std_msgs.msg import Bool, Float64 + + if node is None: + rclpy.init(args=None) + self._node = Node("srb_teleop_ros2") + else: + self._node = node + + # Store inputs + self.pos_sensitivity = pos_sensitivity + self.rot_sensitivity = rot_sensitivity + + self.sub_cmd_bel = self._node.create_subscription( + Twist, "cmd_vel", self.cb_twist, 1 + ) + self.sub_gripper = self._node.create_subscription( + Bool, "gripper", self.cb_event, 1 + ) + self.sub_latency = self._node.create_subscription( + Float64, "gui/latency", self.cb_latency, 1 + ) + + self.latency = 0.0 + self.command_queue = [] + self.last_command = None + + # Command buffers + self._close_gripper = False + self._delta_pos = np.zeros(3) # (x, y, z) + self._delta_rot = np.zeros(3) # (roll, pitch, yaw) + + # Run a thread for listening to device + if node is None: + self._thread = threading.Thread(target=rclpy.spin, args=(self._node,)) + self._thread.daemon = True + self._thread.start() + + def cb_twist(self, msg): + self._delta_pos[0] = self.pos_sensitivity * msg.linear.x + self._delta_pos[1] = self.pos_sensitivity * msg.linear.y + self._delta_pos[2] = self.pos_sensitivity * msg.linear.z + + self._delta_rot[0] = self.rot_sensitivity * msg.angular.x + self._delta_rot[1] = self.rot_sensitivity * msg.angular.y + self._delta_rot[2] = self.rot_sensitivity * msg.angular.z + + def cb_event(self, msg): + if msg.data: + self._close_gripper = not self._close_gripper + + def cb_latency(self, msg): + if msg.data != self.latency: + self.latency = msg.data + self.command_queue = [] + self.feedback_queue = [] + self.last_command = None + self.last_feedback = None + + def __del__(self): + self._thread.join() + + def __str__(self) -> str: + msg = f"ROS 2 Interface ({self.__class__.__name__})\n" + msg += f"Listenining on ROS_DOMAIN_ID: {os.environ.get('ROS_DOMAIN_ID', 0)}\n" + return msg + + def reset(self): + self._close_gripper = False + self._delta_pos = np.zeros(3) + self._delta_rot = np.zeros(3) + self.command_queue = [] + self.feedback_queue = [] + self.last_command = None + self.last_feedback = None + + def add_callback(self, key: str, func: Callable): + raise NotImplementedError + + def advance(self) -> tuple[np.ndarray, bool]: + commands = ( + np.concatenate([self._delta_pos, self._delta_rot]), + self._close_gripper, + ) + + if self.latency == 0.0: + return commands + else: + system_time = self._node.get_clock().now() + self.command_queue.append((system_time, commands)) + + # Find the last viable command + last_viable_command = None + for i, (t, _) in enumerate(self.command_queue): + if (system_time - t).nanoseconds / 1e9 > self.latency: + last_viable_command = i + else: + break + + if last_viable_command is not None: + _, self.last_command = self.command_queue[last_viable_command] + self.command_queue = self.command_queue[last_viable_command + 1 :] + return self.last_command + else: + if self.last_command is not None: + return self.last_command + else: + return np.zeros(6), commands[1] diff --git a/space_robotics_bench/core/teleop_devices/se3_spacemouse.py b/space_robotics_bench/core/teleop_devices/se3_spacemouse.py new file mode 100644 index 0000000..82c5e23 --- /dev/null +++ b/space_robotics_bench/core/teleop_devices/se3_spacemouse.py @@ -0,0 +1,112 @@ +import sys +import threading +import time +from collections.abc import Callable +from typing import List + +import numpy as np +import pyspacemouse +from scipy.spatial.transform import Rotation + +from space_robotics_bench.core.teleop_devices import DeviceBase + + +class Se3SpaceMouse(DeviceBase): + def __init__( + self, + pos_sensitivity: float = 0.4, + rot_sensitivity: float = 0.8, + rate: float = 1000.0, + ): + # Store inputs + self.pos_sensitivity = pos_sensitivity + self.rot_sensitivity = rot_sensitivity + self.sleep_rate = 1.0 / rate + + # Command buffers + self._close_gripper = False + self._delta_pos = np.zeros(3) # (x, y, z) + self._delta_rot = np.zeros(3) # (roll, pitch, yaw) + self._additional_callbacks = {} + + # Open the device + try: + success = pyspacemouse.open( + dof_callback=self._cb_dof, + button_callback=self._cb_button, + ) + if success: + # Run a background thread for the device + self._thread = threading.Thread(target=self._run_device) + self._thread.daemon = True + self._thread.start() + else: + print( + "[ERROR] Failed to open a SpaceMouse device. Is it connected?", + file=sys.stderr, + ) + except Exception as e: + print( + f"[ERROR] Failed to open a SpaceMouse device. Is it connected?\n{e}", + file=sys.stderr, + ) + + def __del__(self): + self._thread.join() + + def __str__(self) -> str: + msg = f"Spacemouse Controller ({self.__class__.__name__})\n" + msg += "\tToggle gripper (alternative): Right button\n" + msg += "\tReset: Left button\n" + return msg + + def reset(self): + # Default flags + self._close_gripper = False + self._delta_pos = np.zeros(3) + self._delta_rot = np.zeros(3) + + def add_callback(self, key: str, func: Callable): + if key not in ["L", "R", "LR"]: + raise ValueError( + f"Only left (L), right (R), and right-left (LR) buttons supported. Provided: {key}." + ) + self._additional_callbacks[key] = func + + def advance(self) -> tuple[np.ndarray, bool]: + rot_vec = Rotation.from_euler("XYZ", self._delta_rot).as_rotvec() + return np.concatenate([self._delta_pos, rot_vec]), self._close_gripper + + def _run_device(self): + while True: + _state = pyspacemouse.read() + time.sleep(self.sleep_rate) + + def _cb_dof(self, state: pyspacemouse.SpaceNavigator): + self._delta_pos = np.array( + [ + state.y * self.pos_sensitivity, + -state.x * self.pos_sensitivity, + state.z * self.pos_sensitivity, + ] + ) + self._delta_rot = np.array( + [ + -state.roll * self.rot_sensitivity, + -state.pitch * self.rot_sensitivity, + -state.yaw * self.rot_sensitivity, + ] + ) + + def _cb_button(self, state: pyspacemouse.SpaceNavigator, buttons: List[bool]): + if buttons[0]: + self.reset() + if "L" in self._additional_callbacks.keys(): + self._additional_callbacks["L"]() + if buttons[1]: + self._close_gripper = not self._close_gripper + if "R" in self._additional_callbacks.keys(): + self._additional_callbacks["R"]() + if all(buttons): + if "LR" in self._additional_callbacks.keys(): + self._additional_callbacks["LR"]() diff --git a/space_robotics_bench/core/teleop_devices/se3_touch.py b/space_robotics_bench/core/teleop_devices/se3_touch.py new file mode 100644 index 0000000..3c722b8 --- /dev/null +++ b/space_robotics_bench/core/teleop_devices/se3_touch.py @@ -0,0 +1,226 @@ +import os +import threading +from collections.abc import Callable +from typing import Optional + +import numpy as np + +from space_robotics_bench.core.teleop_devices import DeviceBase +from space_robotics_bench.utils.ros import enable_ros2_bridge + + +class Se3Touch(DeviceBase): + def __init__( + self, + node: Optional[object] = None, + pos_sensitivity: float = 1.0, + rot_sensitivity: float = 1.0, + ): + enable_ros2_bridge() + import rclpy + from geometry_msgs.msg import Twist, Vector3 + from rclpy.node import Node + from std_msgs.msg import Bool, Float64 + + if node is None: + rclpy.init(args=None) + self._node = Node("se3_teleop_touch") + else: + self._node = node + + # Store inputs + self.pos_sensitivity = pos_sensitivity + self.rot_sensitivity = rot_sensitivity + + self.sub_twist = self._node.create_subscription( + Twist, "touch/twist_stylus_current", self.cb_twist, 1 + ) + self.sub_button = self._node.create_subscription( + Bool, "touch/button", self.cb_button, 1 + ) + self.sub_event = self._node.create_subscription( + Bool, "touch/event", self.cb_event, 1 + ) + self.pub_force_feedback = self._node.create_publisher( + Vector3, "touch/force_feedback", 1 + ) + + self.sub_latency = self._node.create_subscription( + Float64, "gui/latency", self.cb_latency, 1 + ) + self.sub_motion_sensitivity = self._node.create_subscription( + Float64, "gui/motion_sensitivity", self.cb_motion_sensitivity, 1 + ) + self.sub_force_feedback_sensitivity = self._node.create_subscription( + Float64, + "gui/force_feedback_sensitivity", + self.cb_force_feedback_sensitivity, + 1, + ) + + # Detects if the button is pressed, which in turn enables the motion + self.is_motion_enabled = False + self.gain = 1.0 + self.inactive_feedback_decay = 1.0 + + self.latency = 0.0 + self.command_queue = [] + self.feedback_queue = [] + self.last_command = None + self.last_feedback = None + + self.force_feedback_sensitivity = 1.0 + + # Command buffers + self._close_gripper = False + self._delta_pos = np.zeros(3) # (x, y, z) + self._delta_rot = np.zeros(3) # (roll, pitch, yaw) + + # Run a thread for listening to device + if node is None: + self._thread = threading.Thread(target=rclpy.spin, args=(self._node,)) + self._thread.daemon = True + self._thread.start() + + def cb_twist(self, msg): + if self.is_motion_enabled: + self._delta_pos[0] = -self.pos_sensitivity * self.gain * msg.linear.z + self._delta_pos[1] = -self.pos_sensitivity * self.gain * msg.linear.x + self._delta_pos[2] = self.pos_sensitivity * self.gain * msg.linear.y + + self._delta_rot[0] = self.rot_sensitivity * self.gain * msg.angular.z + self._delta_rot[1] = self.rot_sensitivity * self.gain * msg.angular.x + self._delta_rot[2] = self.rot_sensitivity * self.gain * msg.angular.y + + def cb_button(self, msg): + self.is_motion_enabled = msg.data + if not self.is_motion_enabled: + self._delta_pos = np.zeros(3) + self._delta_rot = np.zeros(3) + + def cb_event(self, msg): + if msg.data: + self._close_gripper = not self._close_gripper + + def cb_latency(self, msg): + if msg.data != self.latency: + self.latency = msg.data + self.command_queue = [] + self.feedback_queue = [] + self.last_command = None + self.last_feedback = None + + def cb_motion_sensitivity(self, msg): + self.gain = msg.data + + def cb_force_feedback_sensitivity(self, msg): + self.force_feedback_sensitivity = msg.data + + def __del__(self): + self._thread.join() + + def __str__(self) -> str: + msg = f'Haptic "Touch" Interface ({self.__class__.__name__})\n' + msg += f"Listenining on ROS_DOMAIN_ID: {os.environ.get('ROS_DOMAIN_ID', 0)}\n" + msg += "Hold to enable motion: Dark grey button\n" + msg += "Press to toggle the gripper: Light grey button\n" + return msg + + def reset(self): + self._close_gripper = False + self._delta_pos = np.zeros(3) + self._delta_rot = np.zeros(3) + self.command_queue = [] + self.feedback_queue = [] + self.last_command = None + self.last_feedback = None + + def add_callback(self, key: str, func: Callable): + raise NotImplementedError + + def advance(self) -> tuple[np.ndarray, bool]: + commands = ( + np.concatenate([self._delta_pos, self._delta_rot]), + self._close_gripper, + ) + + if self.latency == 0.0: + return commands + else: + system_time = self._node.get_clock().now() + self.command_queue.append((system_time, commands)) + + # Find the last viable command + last_viable_command = None + for i, (t, _) in enumerate(self.command_queue): + if (system_time - t).nanoseconds / 1e9 > self.latency: + last_viable_command = i + else: + break + + if last_viable_command is not None: + _, self.last_command = self.command_queue[last_viable_command] + self.command_queue = self.command_queue[last_viable_command + 1 :] + return self.last_command + else: + if self.last_command is not None: + return self.last_command + else: + return np.zeros(6), commands[1] + + def set_ft_feedback(self, ft_feedback): + from geometry_msgs.msg import Vector3 + + if self.latency == 0.0: + if self.is_motion_enabled: + self.pub_force_feedback.publish( + Vector3( + x=self.force_feedback_sensitivity * float(ft_feedback[0]), + y=self.force_feedback_sensitivity * float(ft_feedback[1]), + z=self.force_feedback_sensitivity * float(ft_feedback[2]), + ) + ) + self.inactive_feedback_decay = 1.0 + else: + if self.inactive_feedback_decay > 0.0: + self.inactive_feedback_decay = max( + 0.0, 0.95 * self.inactive_feedback_decay + ) + self.pub_force_feedback.publish( + Vector3( + x=self.inactive_feedback_decay + * self.force_feedback_sensitivity + * float(ft_feedback[0]), + y=self.inactive_feedback_decay + * self.force_feedback_sensitivity + * float(ft_feedback[1]), + z=self.inactive_feedback_decay + * self.force_feedback_sensitivity + * float(ft_feedback[2]), + ) + ) + else: + system_time = self._node.get_clock().now() + self.feedback_queue.append((system_time, ft_feedback)) + + # Find the last viable feedback + last_viable_feedback = None + for i, (t, _) in enumerate(self.feedback_queue): + if (system_time - t).nanoseconds / 1e9 > self.latency: + last_viable_feedback = i + else: + break + + if last_viable_feedback is not None: + _, self.last_feedback = self.feedback_queue[last_viable_feedback] + self.feedback_queue = self.feedback_queue[last_viable_feedback + 1 :] + self.pub_force_feedback.publish( + Vector3( + x=self.force_feedback_sensitivity + * float(self.last_feedback[0]), + y=self.force_feedback_sensitivity + * float(self.last_feedback[1]), + z=self.force_feedback_sensitivity + * float(self.last_feedback[2]), + ) + ) diff --git a/space_robotics_bench/core/wrappers/__init__.py b/space_robotics_bench/core/wrappers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/space_robotics_bench/core/wrappers/dreamerv3.py b/space_robotics_bench/core/wrappers/dreamerv3.py new file mode 100644 index 0000000..b500c78 --- /dev/null +++ b/space_robotics_bench/core/wrappers/dreamerv3.py @@ -0,0 +1,350 @@ +import functools +from typing import Literal, Union + +import dreamerv3 +import embodied +import gymnasium +import numpy as np +import torch + +from space_robotics_bench.core.envs import BaseEnv, BaseEnvManaged + + +def process_dreamerv3_cfg( + cfg: dict, + logdir: str, + num_envs: int, + task_name: str = "unknown_task", + model_size: Literal[ + "debug", "size12m", "size25m", "size50m", "size100m", "size200m", "size400m" + ] = "debug", +) -> embodied.Config: + config = embodied.Config(dreamerv3.Agent.configs["defaults"]) + + config = config.update( + dreamerv3.Agent.configs[model_size], + ) + config = config.update(cfg) + config = config.update( + { + "logdir": logdir, + "task": task_name, + "run.num_envs": num_envs, + } + ) + config = embodied.Flags(config).parse(argv=[]) + config = config.update( + replay_length=config.replay_length or config.batch_length, + replay_length_eval=config.replay_length_eval or config.batch_length_eval, + ) + + return config + + +class EmbodiedEnvWrapper(embodied.Env): + def __init__( + self, + env: Union[BaseEnv, BaseEnvManaged], + obs_key="image", + act_key="action", + ): + # check that input is valid + if not isinstance(env.unwrapped, (BaseEnv, BaseEnvManaged)): + raise ValueError( + f"The environment must be inherited from ManagerBasedRLEnv. Environment type: {type(env)}" + ) + # initialize the wrapper + self._env = env + # collect common information + self.num_envs = self.unwrapped.num_envs + self.sim_device = self.unwrapped.device + self.render_mode = self.unwrapped.render_mode + # add buffer for logging episodic information + self._ep_rew_buf = torch.zeros(self.num_envs, device=self.sim_device) + self._ep_len_buf = torch.zeros(self.num_envs, device=self.sim_device) + + self._action_space = self.unwrapped.single_action_space + if ( + isinstance(self._action_space, gymnasium.spaces.Box) + and self._action_space.is_bounded() != "both" + ): + self._action_space = gymnasium.spaces.Box( + low=-100.0, high=100.0, shape=self._action_space.shape + ) + + self._obs_dict = hasattr(self.unwrapped.single_observation_space, "spaces") + self._act_dict = hasattr(self._action_space, "spaces") + self._obs_key = obs_key + self._act_key = act_key + + self._done = np.ones(self.num_envs, dtype=bool) + # self._is_first = np.zeros(self.num_envs, dtype=bool) + self._info = [None for _ in range(self.num_envs)] + + def __len__(self): + return self.num_envs + + def __str__(self): + return f"<{type(self).__name__}{self.env}>" + + def __repr__(self): + return str(self) + + @classmethod + def class_name(cls) -> str: + return cls.__name__ + + @property + def unwrapped(self) -> Union[BaseEnv, BaseEnvManaged]: + return self.env.unwrapped + + def get_episode_rewards(self) -> list[float]: + return self._ep_rew_buf.cpu().tolist() + + def get_episode_lengths(self) -> list[int]: + return self._ep_len_buf.cpu().tolist() + + @property + def env(self): + return self._env + + @property + def info(self): + return self._info + + @functools.cached_property + def obs_space(self): + if self._obs_dict: + spaces = self._flatten(self.unwrapped.single_observation_space.spaces) + else: + spaces = {self._obs_key: self.unwrapped.single_observation_space} + spaces = {k: self._convert(v) for k, v in spaces.items()} + return { + **spaces, + "reward": embodied.Space(np.float32), + "is_first": embodied.Space(bool), + "is_last": embodied.Space(bool), + "is_terminal": embodied.Space(bool), + } + + @functools.cached_property + def act_space(self): + if self._act_dict: + spaces = self._flatten(self._action_space.spaces) + else: + spaces = {self._act_key: self._action_space} + spaces = {k: self._convert(v) for k, v in spaces.items()} + spaces["reset"] = embodied.Space(bool) + return spaces + + def _convert_space(self, space): + if hasattr(space, "n"): + return embodied.Space(np.int32, (), 0, space.n) + else: + return embodied.Space(space.dtype, space.shape, space.low, space.high) + + def seed(self, seed: int | None = None) -> list[int | None]: + return [self.unwrapped.seed(seed)] * self.unwrapped.num_envs + + def reset(self): # noqa: D102 + obs, self._info = self._env.reset() + self._done = np.zeros(self.num_envs, dtype=bool) + # self._is_first = np.zeros(self.num_envs, dtype=bool) + return self._obs( + obs=obs, + reward=np.zeros(self.num_envs, dtype=np.float32), + is_first=np.ones(self.num_envs, dtype=bool), + is_last=np.zeros(self.num_envs, dtype=bool), + is_terminal=np.zeros(self.num_envs, dtype=bool), + ) + + def step(self, action): + if action["reset"].all() or self._done.all(): + return self.reset() + + if self._act_dict: + action = self._unflatten(action) + else: + action = action[self._act_key] + + if not isinstance(action, torch.Tensor): + action = np.asarray(action) + action = torch.from_numpy(action.copy()).to( + device=self.sim_device, dtype=torch.float32 + ) + else: + action = action.to(device=self.sim_device, dtype=torch.float32) + + obs, reward, terminated, truncated, self._info = self._env.step(action) + + # update episode un-discounted return and length + self._ep_rew_buf += reward + self._ep_len_buf += 1 + + self._done = (terminated | truncated).detach().cpu().numpy() + reset_ids = (self._done > 0).nonzero() + + reward = reward.detach().cpu().numpy() + terminated = terminated.detach().cpu().numpy() + truncated = truncated.detach().cpu().numpy() + + # reset info for terminated environments + self._ep_rew_buf[reset_ids] = 0 + self._ep_len_buf[reset_ids] = 0 + + # is_first = self._is_first.copy() + # self._is_first = self._done.copy() + + return self._obs( + obs=obs, + reward=reward, + # is_first=is_first, + is_first=np.zeros(self.num_envs, dtype=bool), + is_last=self._done, + is_terminal=terminated, + ) + + def render(self): + image = self._env.render() + assert image is not None + return image + + def close(self): + try: + self.env.close() + except Exception: + pass + + def get_attr(self, attr_name, indices=None): + # resolve indices + if indices is None: + indices = slice(None) + num_indices = self.num_envs + else: + num_indices = len(indices) + # obtain attribute value + attr_val = getattr(self.env, attr_name) + # return the value + if not isinstance(attr_val, torch.Tensor): + return [attr_val] * num_indices + else: + return attr_val[indices].detach().cpu().numpy() + + def set_attr(self, attr_name, value, indices=None): + raise NotImplementedError("Setting attributes is not supported.") + + def env_method(self, method_name: str, *method_args, indices=None, **method_kwargs): + if method_name == "render": + # gymnasium does not support changing render mode at runtime + return self.env.render() + else: + # this isn't properly implemented but it is not necessary. + # mostly done for completeness. + env_method = getattr(self.env, method_name) + return env_method(*method_args, indices=indices, **method_kwargs) + + def _obs( + self, + obs: torch.Tensor | dict[str, torch.Tensor], + reward: np.ndarray, + is_first: np.ndarray, + is_last: np.ndarray, + is_terminal: np.ndarray, + ): + if not self._obs_dict: + obs = {self._obs_key: obs} + obs = self._flatten(obs) + obs = {k: v.detach().cpu().numpy() for k, v in obs.items()} + obs.update( + reward=reward, + is_first=is_first, + is_last=is_last, + is_terminal=is_terminal, + ) + return obs + + def _flatten(self, nest, prefix=None): + result = {} + for key, value in nest.items(): + key = prefix + "/" + key if prefix else key + if isinstance(value, gymnasium.spaces.Dict): + value = value.spaces + if isinstance(value, dict): + result.update(self._flatten(value, key)) + else: + result[key] = value + return result + + def _unflatten(self, flat): + result = {} + for key, value in flat.items(): + parts = key.split("/") + node = result + for part in parts[:-1]: + if part not in node: + node[part] = {} + node = node[part] + node[parts[-1]] = value + return result + + def _convert(self, space): + if hasattr(space, "n"): + return embodied.Space(np.int32, (), 0, space.n) + return embodied.Space(space.dtype, space.shape, space.low, space.high) + + +class DriverParallelEnv: + def __init__(self, env, num_envs: int, **kwargs): + self.kwargs = kwargs + self.length = num_envs + self.env = env + self.act_space = self.env.act_space + self.callbacks = [] + self.acts = None + self.carry = None + self.reset() + + def reset(self, init_policy=None): + self.acts = { + k: np.zeros((self.length,) + v.shape, v.dtype) + for k, v in self.act_space.items() + } + self.acts["reset"] = np.ones(self.length, bool) + self.carry = init_policy and init_policy(self.length) + + def close(self): + self.env.close() + + def on_step(self, callback): + self.callbacks.append(callback) + + def __call__(self, policy, steps=0, episodes=0): + step, episode = 0, 0 + while step < steps or episode < episodes: + step, episode = self._step(policy, step, episode) + + def _step(self, policy, step, episode): + acts = self.acts + assert all(len(x) == self.length for x in acts.values()) + assert all(isinstance(v, np.ndarray) for v in acts.values()) + obs = self.env.step(acts) + assert all(len(x) == self.length for x in obs.values()), obs + acts, outs, self.carry = policy(obs, self.carry, **self.kwargs) + assert all(k not in acts for k in outs), (list(outs.keys()), list(acts.keys())) + if obs["is_last"].any(): + mask = ~obs["is_last"] + acts = {k: self._mask(v, mask) for k, v in acts.items()} + acts["reset"] = obs["is_last"].copy() + self.acts = acts + trans = {**obs, **acts, **outs} + for i in range(self.length): + trn = {k: v[i] for k, v in trans.items()} + [fn(trn, i, **self.kwargs) for fn in self.callbacks] + step += len(obs["is_first"]) + episode += obs["is_last"].sum() + return step, episode + + def _mask(self, value, mask): + while mask.ndim < value.ndim: + mask = mask[..., None] + return value * mask.astype(value.dtype) diff --git a/space_robotics_bench/core/wrappers/robomimic.py b/space_robotics_bench/core/wrappers/robomimic.py new file mode 100644 index 0000000..22b5195 --- /dev/null +++ b/space_robotics_bench/core/wrappers/robomimic.py @@ -0,0 +1,297 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Interface to collect and store data from the environment using format from `robomimic`.""" + +# needed to import for allowing type-hinting: np.ndarray | torch.Tensor + +import json +import os +from collections.abc import Iterable + +import carb +import h5py +import numpy as np +import torch + + +class RobomimicDataCollector: + """Data collection interface for robomimic. + + This class implements a data collector interface for saving simulation states to disk. + The data is stored in `HDF5`_ binary data format. The class is useful for collecting + demonstrations. The collected data follows the `structure`_ from robomimic. + + All datasets in `robomimic` require the observations and next observations obtained + from before and after the environment step. These are stored as a dictionary of + observations in the keys "obs" and "next_obs" respectively. + + For certain agents in `robomimic`, the episode data should have the following + additional keys: "actions", "rewards", "dones". This behavior can be altered by changing + the dataset keys required in the training configuration for the respective learning agent. + + For reference on datasets, please check the robomimic `documentation`. + + .. _HDF5: https://www.h5py.org/ + .. _structure: https://robomimic.github.io/docs/datasets/overview.html#dataset-structure + .. _documentation: https://github.com/ARISE-Initiative/robomimic/blob/8c23c696260ffb2709b95cd868cf3fc6480afeba/robomimic/config/base_config.py#L167-L173 + """ + + def __init__( + self, + env_name: str, + directory_path: str, + filename: str = "test", + num_demos: int = 1, + flush_freq: int = 1, + env_config: dict | None = None, + ): + """Initializes the data collection wrapper. + + Args: + env_name: The name of the environment. + directory_path: The path to store collected data. + filename: The basename of the saved file. Defaults to "test". + num_demos: Number of demonstrations to record until stopping. Defaults to 1. + flush_freq: Frequency to dump data to disk. Defaults to 1. + env_config: The configuration for the environment. Defaults to None. + """ + # save input arguments + self._env_name = env_name + self._env_config = env_config + self._directory = os.path.abspath(directory_path) + self._filename = filename + self._num_demos = num_demos + self._flush_freq = flush_freq + # print info + print(self.__str__()) + + # create directory it doesn't exist + if not os.path.isdir(self._directory): + os.makedirs(self._directory) + + # placeholder for current hdf5 file object + self._h5_file_stream = None + self._h5_data_group = None + self._h5_episode_group = None + + # store count of demos within episode + self._demo_count = 0 + # flags for setting up + self._is_first_interaction = True + self._is_stop = False + # create buffers to store data + self._dataset = {} + + def __del__(self): + """Destructor for data collector.""" + if not self._is_stop: + self.close() + + def __str__(self) -> str: + """Represents the data collector as a string.""" + msg = "Dataset collector object" + msg += f"\tStoring trajectories in directory: {self._directory}\n" + msg += f"\tNumber of demos for collection : {self._num_demos}\n" + msg += f"\tFrequency for saving data to disk: {self._flush_freq}\n" + + return msg + + """ + Properties + """ + + @property + def demo_count(self) -> int: + """The number of demos collected so far.""" + return self._demo_count + + """ + Operations. + """ + + def is_stopped(self) -> bool: + """Whether data collection is stopped or not. + + Returns: + True if data collection has stopped. + """ + return self._is_stop + + def reset(self): + """Reset the internals of data logger.""" + # setup the file to store data in + if self._is_first_interaction: + self._demo_count = 0 + self._create_new_file(self._filename) + self._is_first_interaction = False + # clear out existing buffers + self._dataset = {} + + def add(self, key: str, value: np.ndarray | torch.Tensor): + """Add a key-value pair to the dataset. + + The key can be nested by using the "/" character. For example: + "obs/joint_pos". Currently only two-level nesting is supported. + + Args: + key: The key name. + value: The corresponding value + of shape (N, ...), where `N` is number of environments. + + Raises: + ValueError: When provided key has sub-keys more than 2. Example: "obs/joints/pos", instead + of "obs/joint_pos". + """ + # check if data should be recorded + if self._is_first_interaction: + carb.log_warn("Please call reset before adding new data. Calling reset...") + self.reset() + if self._is_stop: + carb.log_warn( + f"Desired number of demonstrations collected: {self._demo_count} >= {self._num_demos}." + ) + return + # check datatype + if isinstance(value, torch.Tensor): + value = value.cpu().numpy() + else: + value = np.asarray(value) + # check if there are sub-keys + sub_keys = key.split("/") + num_sub_keys = len(sub_keys) + if len(sub_keys) > 2: + raise ValueError( + f"Input key '{key}' has elements {num_sub_keys} which is more than two." + ) + # add key to dictionary if it doesn't exist + for i in range(value.shape[0]): + # demo index + if f"env_{i}" not in self._dataset: + self._dataset[f"env_{i}"] = {} + # key index + if num_sub_keys == 2: + # create keys + if sub_keys[0] not in self._dataset[f"env_{i}"]: + self._dataset[f"env_{i}"][sub_keys[0]] = {} + if sub_keys[1] not in self._dataset[f"env_{i}"][sub_keys[0]]: + self._dataset[f"env_{i}"][sub_keys[0]][sub_keys[1]] = [] + # add data to key + self._dataset[f"env_{i}"][sub_keys[0]][sub_keys[1]].append(value[i]) + else: + # create keys + if sub_keys[0] not in self._dataset[f"env_{i}"]: + self._dataset[f"env_{i}"][sub_keys[0]] = [] + # add data to key + self._dataset[f"env_{i}"][sub_keys[0]].append(value[i]) + + def flush(self, env_ids: Iterable[int] = (0,)): + """Flush the episode data based on environment indices. + + Args: + env_ids: Environment indices to write data for. Defaults to (0). + """ + # check that data is being recorded + if self._h5_file_stream is None or self._h5_data_group is None: + carb.log_error( + "No file stream has been opened. Please call reset before flushing data." + ) + return + + # iterate over each environment and add their data + for index in env_ids: + # data corresponding to demo + env_dataset = self._dataset[f"env_{index}"] + + # create episode group based on demo count + h5_episode_group = self._h5_data_group.create_group( + f"demo_{self._demo_count}" + ) + # store number of steps taken + h5_episode_group.attrs["num_samples"] = len(env_dataset["actions"]) + # store other data from dictionary + for key, value in env_dataset.items(): + if isinstance(value, dict): + # create group + key_group = h5_episode_group.create_group(key) + # add sub-keys values + for sub_key, sub_value in value.items(): + key_group.create_dataset(sub_key, data=np.array(sub_value)) + else: + h5_episode_group.create_dataset(key, data=np.array(value)) + # increment total step counts + self._h5_data_group.attrs["total"] += h5_episode_group.attrs["num_samples"] + + # increment total demo counts + self._demo_count += 1 + # reset buffer for environment + self._dataset[f"env_{index}"] = {} + + # dump at desired frequency + if self._demo_count % self._flush_freq == 0: + self._h5_file_stream.flush() + print( + f">>> Flushing data to disk. Collected demos: {self._demo_count} / {self._num_demos}" + ) + + # if demos collected then stop + if self._demo_count >= self._num_demos: + print( + f">>> Desired number of demonstrations collected: {self._demo_count} >= {self._num_demos}." + ) + self.close() + # break out of loop + break + + def close(self): + """Stop recording and save the file at its current state.""" + if not self._is_stop: + print( + f">>> Closing recording of data. Collected demos: {self._demo_count} / {self._num_demos}" + ) + # close the file safely + if self._h5_file_stream is not None: + self._h5_file_stream.close() + # mark that data collection is stopped + self._is_stop = True + + """ + Helper functions. + """ + + def _create_new_file(self, fname: str): + """Create a new HDF5 file for writing episode info into. + + Reference: + https://robomimic.github.io/docs/datasets/overview.html + + Args: + fname: The base name of the file. + """ + if not fname.endswith(".hdf5"): + fname += ".hdf5" + # define path to file + hdf5_path = os.path.join(self._directory, fname) + # construct the stream object + self._h5_file_stream = h5py.File(hdf5_path, "w") + # create group to store data + self._h5_data_group = self._h5_file_stream.create_group("data") + # stores total number of samples accumulated across demonstrations + self._h5_data_group.attrs["total"] = 0 + # store the environment meta-info + # -- we use gym environment type + # Ref: https://github.com/ARISE-Initiative/robomimic/blob/8c23c696260ffb2709b95cd868cf3fc6480afeba/robomimic/envs/env_base.py#L15 + env_type = 2 + # -- check if env config provided + if self._env_config is None: + self._env_config = {} + # -- add info + self._h5_data_group.attrs["env_args"] = json.dumps( + { + "env_name": self._env_name, + "type": env_type, + "env_kwargs": self._env_config, + } + ) diff --git a/space_robotics_bench/core/wrappers/sb3.py b/space_robotics_bench/core/wrappers/sb3.py new file mode 100644 index 0000000..4a68d80 --- /dev/null +++ b/space_robotics_bench/core/wrappers/sb3.py @@ -0,0 +1,350 @@ +from typing import Any, Union + +import gymnasium as gym +import numpy as np +import torch +from stable_baselines3.common.utils import constant_fn +from stable_baselines3.common.vec_env.base_vec_env import ( + VecEnv, + VecEnvObs, + VecEnvStepReturn, +) + +from space_robotics_bench.core.envs import BaseEnv, BaseEnvManaged + + +def process_sb3_cfg(cfg: dict) -> dict: + """Convert simple YAML types to Stable-Baselines classes/components. + + Args: + cfg: A configuration dictionary. + + Returns: + A dictionary containing the converted configuration. + + Reference: + https://github.com/DLR-RM/rl-baselines3-zoo/blob/0e5eb145faefa33e7d79c7f8c179788574b20da5/utils/exp_manager.py#L358 + """ + + def update_dict(hyperparams: dict[str, Any]) -> dict[str, Any]: + for key, value in hyperparams.items(): + if isinstance(value, dict): + update_dict(value) + else: + if key in [ + "policy_kwargs", + "replay_buffer_class", + "replay_buffer_kwargs", + ]: + hyperparams[key] = eval(value) + elif key in [ + "learning_rate", + "clip_range", + "clip_range_vf", + "delta_std", + ]: + if isinstance(value, str): + _, initial_value = value.split("_") + initial_value = float(initial_value) + hyperparams[key] = ( + lambda progress_remaining: progress_remaining + * initial_value + ) + elif isinstance(value, (float, int)): + # Negative value: ignore (ex: for clipping) + if value < 0: + continue + hyperparams[key] = constant_fn(float(value)) + else: + raise ValueError(f"Invalid value for {key}: {hyperparams[key]}") + + return hyperparams + + # parse agent configuration and convert to classes + return update_dict(cfg) + + +class Sb3VecEnvWrapper(VecEnv): + """Wraps around Isaac Lab environment for Stable Baselines3. + + Isaac Sim internally implements a vectorized environment. However, since it is + still considered a single environment instance, Stable Baselines tries to wrap + around it using the :class:`DummyVecEnv`. This is only done if the environment + is not inheriting from their :class:`VecEnv`. Thus, this class thinly wraps + over the environment from :class:`ManagerBasedRLEnv`. + + Note: + While Stable-Baselines3 supports Gym 0.26+ API, their vectorized environment + still uses the old API (i.e. it is closer to Gym 0.21). Thus, we implement + the old API for the vectorized environment. + + We also add monitoring functionality that computes the un-discounted episode + return and length. This information is added to the info dicts under key `episode`. + + In contrast to the Isaac Lab environment, stable-baselines expect the following: + + 1. numpy datatype for MDP signals + 2. a list of info dicts for each sub-environment (instead of a dict) + 3. when environment has terminated, the observations from the environment should correspond + to the one after reset. The "real" final observation is passed using the info dicts + under the key ``terminal_observation``. + + .. warning:: + + By the nature of physics stepping in Isaac Sim, it is not possible to forward the + simulation buffers without performing a physics step. Thus, reset is performed + inside the :meth:`step()` function after the actual physics step is taken. + Thus, the returned observations for terminated environments is the one after the reset. + + .. caution:: + + This class must be the last wrapper in the wrapper chain. This is because the wrapper does not follow + the :class:`gym.Wrapper` interface. Any subsequent wrappers will need to be modified to work with this + wrapper. + + Reference: + + 1. https://stable-baselines3.readthedocs.io/en/master/guide/vec_envs.html + 2. https://stable-baselines3.readthedocs.io/en/master/common/monitor.html + + """ + + def __init__(self, env: Union[BaseEnv, BaseEnvManaged]): + """Initialize the wrapper. + + Args: + env: The environment to wrap around. + + Raises: + ValueError: When the environment is not an instance of :class:`ManagerBasedRLEnv`. + """ + # check that input is valid- + if not isinstance(env.unwrapped, (BaseEnv, BaseEnvManaged)): + raise ValueError( + "The environment must be inherited from ManagerBasedRLEnv or DirectRLEnv. Environment type:" + f" {type(env)}" + ) + # initialize the wrapper + self.env = env + # collect common information + self.num_envs = self.unwrapped.num_envs + self.sim_device = self.unwrapped.device + self.render_mode = self.unwrapped.render_mode + + # obtain gym spaces + # note: stable-baselines3 does not like when we have unbounded action space so + # we set it to some high value here. Maybe this is not general but something to think about. + observation_space = self.unwrapped.single_observation_space["policy"] + action_space = self.unwrapped.single_action_space + if isinstance(action_space, gym.spaces.Box) and not action_space.is_bounded( + "both" + ): + action_space = gym.spaces.Box(low=-100, high=100, shape=action_space.shape) + + # initialize vec-env + VecEnv.__init__(self, self.num_envs, observation_space, action_space) + # add buffer for logging episodic information + self._ep_rew_buf = torch.zeros(self.num_envs, device=self.sim_device) + self._ep_len_buf = torch.zeros(self.num_envs, device=self.sim_device) + + def __str__(self): + """Returns the wrapper name and the :attr:`env` representation string.""" + return f"<{type(self).__name__}{self.env}>" + + def __repr__(self): + """Returns the string representation of the wrapper.""" + return str(self) + + """ + Properties -- Gym.Wrapper + """ + + @classmethod + def class_name(cls) -> str: + """Returns the class name of the wrapper.""" + return cls.__name__ + + @property + def unwrapped(self) -> Union[BaseEnv, BaseEnvManaged]: + """Returns the base environment of the wrapper. + + This will be the bare :class:`gymnasium.Env` environment, underneath all layers of wrappers. + """ + return self.env.unwrapped + + """ + Properties + """ + + def get_episode_rewards(self) -> list[float]: + """Returns the rewards of all the episodes.""" + return self._ep_rew_buf.cpu().tolist() + + def get_episode_lengths(self) -> list[int]: + """Returns the number of time-steps of all the episodes.""" + return self._ep_len_buf.cpu().tolist() + + """ + Operations - MDP + """ + + def seed(self, seed: int | None = None) -> list[int | None]: # noqa: D102 + return [self.unwrapped.seed(seed)] * self.unwrapped.num_envs + + def reset(self) -> VecEnvObs: # noqa: D102 + obs_dict, _ = self.env.reset() + # convert data types to numpy depending on backend + return self._process_obs(obs_dict) + + def step_async(self, actions): # noqa: D102 + # convert input to numpy array + if not isinstance(actions, torch.Tensor): + actions = np.asarray(actions) + actions = torch.from_numpy(actions).to( + device=self.sim_device, dtype=torch.float32 + ) + else: + actions = actions.to(device=self.sim_device, dtype=torch.float32) + # convert to tensor + self._async_actions = actions + + def step_wait(self) -> VecEnvStepReturn: # noqa: D102 + # record step information + obs_dict, rew, terminated, truncated, extras = self.env.step( + self._async_actions + ) + # update episode un-discounted return and length + self._ep_rew_buf += rew + self._ep_len_buf += 1 + # compute reset ids + dones = terminated | truncated + reset_ids = (dones > 0).nonzero(as_tuple=False) + + # convert data types to numpy depending on backend + # note: ManagerBasedRLEnv uses torch backend (by default). + obs = self._process_obs(obs_dict) + rew = rew.detach().cpu().numpy() + terminated = terminated.detach().cpu().numpy() + truncated = truncated.detach().cpu().numpy() + dones = dones.detach().cpu().numpy() + # convert extra information to list of dicts + infos = self._process_extras(obs, terminated, truncated, extras, reset_ids) + + # reset info for terminated environments + self._ep_rew_buf[reset_ids] = 0 + self._ep_len_buf[reset_ids] = 0 + + return obs, rew, dones, infos + + def close(self): # noqa: D102 + self.env.close() + + def get_attr(self, attr_name, indices=None): # noqa: D102 + # resolve indices + if indices is None: + indices = slice(None) + num_indices = self.num_envs + else: + num_indices = len(indices) + # obtain attribute value + attr_val = getattr(self.env, attr_name) + # return the value + if not isinstance(attr_val, torch.Tensor): + return [attr_val] * num_indices + else: + return attr_val[indices].detach().cpu().numpy() + + def set_attr(self, attr_name, value, indices=None): # noqa: D102 + raise NotImplementedError("Setting attributes is not supported.") + + def env_method( + self, method_name: str, *method_args, indices=None, **method_kwargs + ): # noqa: D102 + if method_name == "render": + # gymnasium does not support changing render mode at runtime + return self.env.render() + else: + # this isn't properly implemented but it is not necessary. + # mostly done for completeness. + env_method = getattr(self.env, method_name) + return env_method(*method_args, indices=indices, **method_kwargs) + + def env_is_wrapped(self, wrapper_class, indices=None): # noqa: D102 + raise NotImplementedError( + "Checking if environment is wrapped is not supported." + ) + + def get_images(self): # noqa: D102 + raise NotImplementedError("Getting images is not supported.") + + """ + Helper functions. + """ + + def _process_obs( + self, obs_dict: torch.Tensor | dict[str, torch.Tensor] + ) -> np.ndarray | dict[str, np.ndarray]: + """Convert observations into NumPy data type.""" + # Sb3 doesn't support asymmetric observation spaces, so we only use "policy" + obs = obs_dict["policy"] + # note: ManagerBasedRLEnv uses torch backend (by default). + if isinstance(obs, dict): + for key, value in obs.items(): + obs[key] = value.detach().cpu().numpy() + elif isinstance(obs, torch.Tensor): + obs = obs.detach().cpu().numpy() + else: + raise NotImplementedError(f"Unsupported data type: {type(obs)}") + return obs + + def _process_extras( + self, + obs: np.ndarray, + terminated: np.ndarray, + truncated: np.ndarray, + extras: dict, + reset_ids: np.ndarray, + ) -> list[dict[str, Any]]: + """Convert miscellaneous information into dictionary for each sub-environment.""" + # create empty list of dictionaries to fill + infos: list[dict[str, Any]] = [ + dict.fromkeys(extras.keys()) for _ in range(self.num_envs) + ] + # fill-in information for each sub-environment + # note: This loop becomes slow when number of environments is large. + for idx in range(self.num_envs): + # fill-in episode monitoring info + if idx in reset_ids: + infos[idx]["episode"] = {} + infos[idx]["episode"]["r"] = float(self._ep_rew_buf[idx]) + infos[idx]["episode"]["l"] = float(self._ep_len_buf[idx]) + else: + infos[idx]["episode"] = None + # fill-in bootstrap information + infos[idx]["TimeLimit.truncated"] = truncated[idx] and not terminated[idx] + # fill-in information from extras + for key, value in extras.items(): + # 1. remap extra episodes information safely + # 2. for others just store their values + if key == "log": + # only log this data for episodes that are terminated + if infos[idx]["episode"] is not None: + for sub_key, sub_value in value.items(): + infos[idx]["episode"][sub_key] = sub_value + else: + infos[idx][key] = value[idx] + # add information about terminal observation separately + if idx in reset_ids: + # extract terminal observations + if isinstance(obs, dict): + terminal_obs = dict.fromkeys(obs.keys()) + for key, value in obs.items(): + terminal_obs[key] = value[idx] + else: + terminal_obs = obs[idx] + # add info to dict + infos[idx]["terminal_observation"] = terminal_obs + else: + infos[idx]["terminal_observation"] = None + # return list of dictionaries + return infos diff --git a/space_robotics_bench/envs/__init__.py b/space_robotics_bench/envs/__init__.py new file mode 100644 index 0000000..fadafcc --- /dev/null +++ b/space_robotics_bench/envs/__init__.py @@ -0,0 +1,3 @@ +from .aerial_robotics import * # noqa: F403 +from .manipulation import * # noqa: F403 +from .mobile_robotics import * # noqa: F403 diff --git a/space_robotics_bench/envs/aerial_robotics/__init__.py b/space_robotics_bench/envs/aerial_robotics/__init__.py new file mode 100644 index 0000000..c1a49b6 --- /dev/null +++ b/space_robotics_bench/envs/aerial_robotics/__init__.py @@ -0,0 +1,4 @@ +from .cfg import * # noqa: F403 +from .impl import * # noqa: F403 + +from .extensions import * # noqa: F403 isort:skip diff --git a/space_robotics_bench/envs/aerial_robotics/cfg.py b/space_robotics_bench/envs/aerial_robotics/cfg.py new file mode 100644 index 0000000..90a8a6e --- /dev/null +++ b/space_robotics_bench/envs/aerial_robotics/cfg.py @@ -0,0 +1,167 @@ +import math + +import torch +from omni.isaac.lab.envs import ViewerCfg +from omni.isaac.lab.managers import EventTermCfg, SceneEntityCfg +from omni.isaac.lab.scene import InteractiveSceneCfg +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.sim as sim_utils +from space_robotics_bench import assets +from space_robotics_bench.core import mdp +from space_robotics_bench.core.envs import BaseEnvCfg +from space_robotics_bench.core.sim import SimulationCfg + + +@configclass +class BaseAerialRoboticsEnvEventCfg: + ## Default scene reset + reset_all = EventTermCfg(func=mdp.reset_scene_to_default, mode="reset") + + ## Light + reset_rand_light_rot = EventTermCfg( + func=mdp.reset_xform_orientation_uniform, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("light"), + "orientation_distribution_params": { + "roll": ( + -75.0 * torch.pi / 180.0, + 75.0 * torch.pi / 180.0, + ), + "pitch": ( + -75.0 * torch.pi / 180.0, + 75.0 * torch.pi / 180.0, + ), + }, + }, + ) + # reset_rand_light_rot = EventTermCfg( + # func=mdp.follow_xform_orientation_linear_trajectory, + # mode="interval", + # interval_range_s=(0.1, 0.1), + # is_global_time=True, + # params={ + # "asset_cfg": SceneEntityCfg("light"), + # "orientation_step_params": { + # "roll": 0.25 * torch.pi / 180.0, + # "pitch": 0.5 * torch.pi / 180.0, + # }, + # }, + # ) + + ## Robot + reset_rand_robot_state = EventTermCfg( + func=mdp.reset_root_state_uniform, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("robot"), + "pose_range": { + "x": (-5.0, 5.0), + "y": (-5.0, 5.0), + "z": (10.0, 10.0), + "yaw": ( + -torch.pi, + torch.pi, + ), + }, + "velocity_range": {}, + }, + ) + + +@configclass +class BaseAerialRoboticsEnvCfg(BaseEnvCfg): + ## Environment + episode_length_s: float = 50.0 + env_rate: float = 1.0 / 50.0 + + ## Agent + agent_rate: float = 1.0 / 50.0 + + ## Simulation + sim = SimulationCfg( + disable_contact_processing=True, + physx=sim_utils.PhysxCfg( + enable_ccd=False, + enable_stabilization=False, + bounce_threshold_velocity=0.0, + friction_correlation_distance=0.02, + min_velocity_iteration_count=1, + # GPU settings + gpu_temp_buffer_capacity=2 ** (24 - 5), + gpu_max_rigid_contact_count=2 ** (22 - 4), + gpu_max_rigid_patch_count=2 ** (13 - 1), + gpu_heap_capacity=2 ** (26 - 6), + gpu_found_lost_pairs_capacity=2 ** (18 - 2), + gpu_found_lost_aggregate_pairs_capacity=2 ** (10 - 1), + gpu_total_aggregate_pairs_capacity=2 ** (10 - 1), + gpu_max_soft_body_contacts=2 ** (20 - 3), + gpu_max_particle_contacts=2 ** (20 - 3), + gpu_collision_stack_size=2 ** (26 - 4), + gpu_max_num_partitions=8, + ), + physics_material=sim_utils.RigidBodyMaterialCfg( + static_friction=1.0, + dynamic_friction=1.0, + restitution=0.0, + friction_combine_mode="multiply", + restitution_combine_mode="multiply", + ), + ) + + ## Viewer + viewer = ViewerCfg( + lookat=(0.0, 0.0, 10.0), + eye=(-10.0, 0.0, 20.0), + origin_type="env", + env_index=0, + ) + + ## Scene + scene = InteractiveSceneCfg(num_envs=1, env_spacing=97.0, replicate_physics=False) + + ## Events + events = BaseAerialRoboticsEnvEventCfg() + + def __post_init__(self): + super().__post_init__() + + ## Simulation + self.decimation = int(self.agent_rate / self.env_rate) + self.sim.dt = self.env_rate + self.sim.render_interval = self.decimation + self.sim.gravity = (0.0, 0.0, -self.env_cfg.scenario.gravity_magnitude) + # Increase GPU settings based on the number of environments + gpu_capacity_factor = self.scene.num_envs + self.sim.physx.gpu_heap_capacity *= gpu_capacity_factor + self.sim.physx.gpu_collision_stack_size *= gpu_capacity_factor + self.sim.physx.gpu_temp_buffer_capacity *= gpu_capacity_factor + self.sim.physx.gpu_max_rigid_contact_count *= gpu_capacity_factor + self.sim.physx.gpu_max_rigid_patch_count *= gpu_capacity_factor + self.sim.physx.gpu_found_lost_pairs_capacity *= gpu_capacity_factor + self.sim.physx.gpu_found_lost_aggregate_pairs_capacity *= gpu_capacity_factor + self.sim.physx.gpu_total_aggregate_pairs_capacity *= gpu_capacity_factor + self.sim.physx.gpu_max_soft_body_contacts *= gpu_capacity_factor + self.sim.physx.gpu_max_particle_contacts *= gpu_capacity_factor + self.sim.physx.gpu_max_num_partitions = min( + 2 ** math.floor(1.0 + math.pow(self.scene.num_envs, 0.2)), 32 + ) + + ## Scene + self.scene.light = assets.sunlight_from_env_cfg(self.env_cfg) + self.scene.sky = assets.sky_from_env_cfg(self.env_cfg) + self.scene.terrain = assets.terrain_from_env_cfg( + self.env_cfg, + num_assets=self.scene.num_envs, + size=(self.scene.env_spacing - 1,) * 2, + procgen_kwargs={ + "density": 0.24, + "texture_resolution": 6144, + }, + ) + self.robot_cfg = assets.aerial_robot_from_env_cfg(self.env_cfg) + self.scene.robot = self.robot_cfg.asset_cfg + + ## Actions + self.actions = self.robot_cfg.action_cfg diff --git a/space_robotics_bench/envs/aerial_robotics/extensions/__init__.py b/space_robotics_bench/envs/aerial_robotics/extensions/__init__.py new file mode 100644 index 0000000..74ff8c6 --- /dev/null +++ b/space_robotics_bench/envs/aerial_robotics/extensions/__init__.py @@ -0,0 +1 @@ +from .visual import * # noqa: F403 diff --git a/space_robotics_bench/envs/aerial_robotics/extensions/visual/__init__.py b/space_robotics_bench/envs/aerial_robotics/extensions/visual/__init__.py new file mode 100644 index 0000000..623a285 --- /dev/null +++ b/space_robotics_bench/envs/aerial_robotics/extensions/visual/__init__.py @@ -0,0 +1,2 @@ +from .cfg import * # noqa: F403 +from .impl import * # noqa: F403 diff --git a/space_robotics_bench/envs/aerial_robotics/extensions/visual/cfg.py b/space_robotics_bench/envs/aerial_robotics/extensions/visual/cfg.py new file mode 100644 index 0000000..bb26468 --- /dev/null +++ b/space_robotics_bench/envs/aerial_robotics/extensions/visual/cfg.py @@ -0,0 +1,74 @@ +from dataclasses import MISSING +from typing import Tuple + +from omni.isaac.lab.envs import ViewerCfg +from omni.isaac.lab.scene import InteractiveSceneCfg +from omni.isaac.lab.sensors import CameraCfg +from omni.isaac.lab.sensors.camera.camera_cfg import PinholeCameraCfg +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.assets as asset_utils +import space_robotics_bench.utils.math as math_utils + + +@configclass +class VisualAerialRoboticsEnvExtCfg: + ## Subclass requirements + agent_rate: int = MISSING + scene: InteractiveSceneCfg = MISSING + viewer: ViewerCfg = MISSING + robot_cfg: asset_utils.AerialRobotCfg = MISSING + + ## Enabling flags + enable_camera_scene: bool = True + enable_camera_bottom: bool = True + + ## Resolution + camera_resolution: Tuple[int, int] = (64, 64) + camera_framerate: int = 0 # 0 matches the agent rate + + def __post_init__(self): + ## Scene + # self.scene.env_spacing += 4.0 + + ## Sensors + framerate = ( + self.camera_framerate if self.camera_framerate > 0 else self.agent_rate + ) + # Scene camera + if self.enable_camera_scene: + self.scene.camera_scene = CameraCfg( + prim_path=f"{self.scene.robot.prim_path}/{self.robot_cfg.frame_base.prim_relpath}/camera_scene", + offset=CameraCfg.OffsetCfg( + convention="world", + pos=(-2.5, 0.0, 2.5), + rot=math_utils.quat_from_rpy(0.0, 45.0, 0.0), + ), + spawn=PinholeCameraCfg( + clipping_range=(0.05, 75.0 - 0.05), + ), + width=self.camera_resolution[0], + height=self.camera_resolution[1], + update_period=framerate, + data_types=["rgb", "distance_to_camera"], + ) + + # Bottom camera + if self.enable_camera_bottom: + self.scene.camera_bottom = CameraCfg( + prim_path=f"{self.scene.robot.prim_path}/{self.robot_cfg.frame_camera_bottom.prim_relpath}", + offset=CameraCfg.OffsetCfg( + convention="world", + pos=self.robot_cfg.frame_camera_bottom.offset.translation, + rot=self.robot_cfg.frame_camera_bottom.offset.rotation, + ), + spawn=PinholeCameraCfg( + focal_length=5.0, + horizontal_aperture=12.0, + clipping_range=(0.05, 50.0 - 0.05), + ), + width=self.camera_resolution[0], + height=self.camera_resolution[1], + update_period=framerate, + data_types=["rgb", "distance_to_camera"], + ) diff --git a/space_robotics_bench/envs/aerial_robotics/extensions/visual/impl.py b/space_robotics_bench/envs/aerial_robotics/extensions/visual/impl.py new file mode 100644 index 0000000..e24df98 --- /dev/null +++ b/space_robotics_bench/envs/aerial_robotics/extensions/visual/impl.py @@ -0,0 +1,46 @@ +from typing import Dict, Tuple + +import torch +from omni.isaac.lab.scene import InteractiveScene +from omni.isaac.lab.sensors import Camera + +from space_robotics_bench.utils import image_proc +from space_robotics_bench.utils import string as string_utils + +from .cfg import VisualAerialRoboticsEnvExtCfg + + +class VisualAerialRoboticsEnvExt: + ## Subclass requirements + common_step_counter: int + scene: InteractiveScene + cfg: VisualAerialRoboticsEnvExtCfg + + def __init__(self, cfg: VisualAerialRoboticsEnvExtCfg, **kwargs): + ## Extract camera sensors from the scene + self.__cameras: Dict[ + str, # Name of the output image + Tuple[ + Camera, # Camera sensor + Tuple[float, float], # Depth range + ], + ] = { + f"image_{string_utils.sanitize_camera_name(key)}": ( + sensor, + getattr(cfg.scene, key).spawn.clipping_range, + ) + for key, sensor in self.scene._sensors.items() + if type(sensor) == Camera + } + + def _get_observations(self) -> Dict[str, torch.Tensor]: + observation = {} + for image_name, (sensor, depth_range) in self.__cameras.items(): + observation.update( + image_proc.construct_observation( + **image_proc.extract_images(sensor), + depth_range=depth_range, + image_name=image_name, + ) + ) + return observation diff --git a/space_robotics_bench/envs/aerial_robotics/impl.py b/space_robotics_bench/envs/aerial_robotics/impl.py new file mode 100644 index 0000000..907b912 --- /dev/null +++ b/space_robotics_bench/envs/aerial_robotics/impl.py @@ -0,0 +1,14 @@ +from space_robotics_bench.core.assets import Articulation +from space_robotics_bench.core.envs import BaseEnv + +from .cfg import BaseAerialRoboticsEnvCfg + + +class BaseAerialRoboticsEnv(BaseEnv): + cfg: BaseAerialRoboticsEnvCfg + + def __init__(self, cfg: BaseAerialRoboticsEnvCfg, **kwargs): + super().__init__(cfg, **kwargs) + + ## Get handles to scene assets + self._robot: Articulation = self.scene["robot"] diff --git a/space_robotics_bench/envs/manipulation/__init__.py b/space_robotics_bench/envs/manipulation/__init__.py new file mode 100644 index 0000000..c1a49b6 --- /dev/null +++ b/space_robotics_bench/envs/manipulation/__init__.py @@ -0,0 +1,4 @@ +from .cfg import * # noqa: F403 +from .impl import * # noqa: F403 + +from .extensions import * # noqa: F403 isort:skip diff --git a/space_robotics_bench/envs/manipulation/cfg.py b/space_robotics_bench/envs/manipulation/cfg.py new file mode 100644 index 0000000..79e2a00 --- /dev/null +++ b/space_robotics_bench/envs/manipulation/cfg.py @@ -0,0 +1,206 @@ +import math + +import torch +from omni.isaac.lab.envs import ViewerCfg +from omni.isaac.lab.managers import EventTermCfg, SceneEntityCfg +from omni.isaac.lab.scene import InteractiveSceneCfg +from omni.isaac.lab.sensors import ContactSensorCfg +from omni.isaac.lab.sensors.frame_transformer.frame_transformer_cfg import ( + FrameTransformerCfg, + OffsetCfg, +) +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.envs as env_utils +import space_robotics_bench.core.sim as sim_utils +from space_robotics_bench import assets +from space_robotics_bench.core import mdp +from space_robotics_bench.core.envs import BaseEnvCfg +from space_robotics_bench.core.markers import FRAME_MARKER_SMALL_CFG +from space_robotics_bench.core.sim import SimulationCfg + + +@configclass +class BaseManipulationEnvEventCfg: + ## Default scene reset + reset_all = EventTermCfg(func=mdp.reset_scene_to_default, mode="reset") + + ## Light + reset_rand_light_rot = EventTermCfg( + func=mdp.reset_xform_orientation_uniform, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("light"), + "orientation_distribution_params": { + "roll": ( + -75.0 * torch.pi / 180.0, + 75.0 * torch.pi / 180.0, + ), + "pitch": ( + 0.0, + 75.0 * torch.pi / 180.0, + ), + }, + }, + ) + + ## Robot + reset_rand_robot_state = EventTermCfg( + func=mdp.reset_joints_by_offset, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("robot"), + "position_range": (-torch.pi / 32, torch.pi / 32), + "velocity_range": (0.0, 0.0), + }, + ) + + +@configclass +class BaseManipulationEnvCfg(BaseEnvCfg): + ## Environment + episode_length_s: float = 50.0 + env_rate: float = 1.0 / 200.0 + + ## Agent + agent_rate: float = 1.0 / 50.0 + + ## Simulation + sim = SimulationCfg( + disable_contact_processing=True, + physx=sim_utils.PhysxCfg( + enable_ccd=True, + enable_stabilization=True, + bounce_threshold_velocity=0.0, + friction_correlation_distance=0.005, + min_velocity_iteration_count=2, + # GPU settings + gpu_temp_buffer_capacity=2 ** (24 - 2), + gpu_max_rigid_contact_count=2 ** (22 - 1), + gpu_max_rigid_patch_count=2 ** (13 - 0), + gpu_heap_capacity=2 ** (26 - 3), + gpu_found_lost_pairs_capacity=2 ** (18 - 1), + gpu_found_lost_aggregate_pairs_capacity=2 ** (10 - 0), + gpu_total_aggregate_pairs_capacity=2 ** (10 - 0), + gpu_max_soft_body_contacts=2 ** (20 - 1), + gpu_max_particle_contacts=2 ** (20 - 1), + gpu_collision_stack_size=2 ** (26 - 3), + gpu_max_num_partitions=8, + ), + physics_material=sim_utils.RigidBodyMaterialCfg( + static_friction=1.0, + dynamic_friction=1.0, + restitution=0.0, + friction_combine_mode="multiply", + restitution_combine_mode="multiply", + ), + ) + + ## Viewer + viewer = ViewerCfg( + lookat=(0.0, 0.0, 0.25), + eye=(2.0, 0.0, 1.75), + origin_type="env", + env_index=0, + ) + + ## Scene + scene = InteractiveSceneCfg(num_envs=1, env_spacing=9.0, replicate_physics=False) + + ## Events + events = BaseManipulationEnvEventCfg() + + def __post_init__(self): + super().__post_init__() + + ## Simulation + self.decimation = int(self.agent_rate / self.env_rate) + self.sim.dt = self.env_rate + self.sim.render_interval = self.decimation + self.sim.gravity = (0.0, 0.0, -self.env_cfg.scenario.gravity_magnitude) + # Increase GPU settings based on the number of environments + gpu_capacity_factor = math.pow(self.scene.num_envs, 0.2) + self.sim.physx.gpu_heap_capacity *= gpu_capacity_factor + self.sim.physx.gpu_collision_stack_size *= gpu_capacity_factor + self.sim.physx.gpu_temp_buffer_capacity *= gpu_capacity_factor + self.sim.physx.gpu_max_rigid_contact_count *= gpu_capacity_factor + self.sim.physx.gpu_max_rigid_patch_count *= gpu_capacity_factor + self.sim.physx.gpu_found_lost_pairs_capacity *= gpu_capacity_factor + self.sim.physx.gpu_found_lost_aggregate_pairs_capacity *= gpu_capacity_factor + self.sim.physx.gpu_total_aggregate_pairs_capacity *= gpu_capacity_factor + self.sim.physx.gpu_max_soft_body_contacts *= gpu_capacity_factor + self.sim.physx.gpu_max_particle_contacts *= gpu_capacity_factor + self.sim.physx.gpu_max_num_partitions = min( + 2 ** math.floor(1.0 + gpu_capacity_factor), 32 + ) + + ## Scene + if self.env_cfg.scenario == env_utils.Scenario.ORBIT: + self.scene.env_spacing = 42.0 + self.scene.light = assets.sunlight_from_env_cfg(self.env_cfg) + self.scene.sky = assets.sky_from_env_cfg(self.env_cfg) + self.robot_cfg = assets.manipulator_from_env_cfg(self.env_cfg) + self.scene.robot = self.robot_cfg.asset_cfg + self.vehicle_cfg = assets.vehicle_from_env_cfg(self.env_cfg) + self.scene.terrain = assets.terrain_from_env_cfg( + self.env_cfg, + num_assets=self.scene.num_envs, + size=(self.scene.env_spacing - 1,) * 2, + procgen_kwargs={ + "density": 0.04, + "flat_area_size": 2.0, + "texture_resolution": 2048, + }, + ) + if self.vehicle_cfg: + # Add vehicle to scene + self.scene.vehicle = self.vehicle_cfg.asset_cfg + self.scene.vehicle.init_state.pos = ( + self.vehicle_cfg.frame_manipulator_base.offset.translation + ) + + # Update the robot based on the vehicle + self.scene.robot.init_state.pos = ( + self.vehicle_cfg.frame_manipulator_base.offset.translation + ) + self.scene.robot.init_state.rot = ( + self.vehicle_cfg.frame_manipulator_base.offset.rotation + ) + + ## Actions + self.actions = self.robot_cfg.action_cfg + + ## Sensors + self.scene.tf_robot_ee = FrameTransformerCfg( + prim_path=f"{self.scene.robot.prim_path}/{self.robot_cfg.frame_base.prim_relpath}", + target_frames=[ + FrameTransformerCfg.FrameCfg( + name="robot_ee", + prim_path=f"{self.scene.robot.prim_path}/{self.robot_cfg.frame_ee.prim_relpath}", + offset=OffsetCfg( + pos=self.robot_cfg.frame_ee.offset.translation, + rot=self.robot_cfg.frame_ee.offset.rotation, + ), + ), + ], + visualizer_cfg=FRAME_MARKER_SMALL_CFG.replace( + prim_path="/Visuals/robot_ee" + ), + ) + self.scene.contacts_robot = ContactSensorCfg( + prim_path=f"{self.scene.robot.prim_path}/.*", + update_period=0.0, + ) + + ## Events + self.events.reset_rand_robot_state.params["asset_cfg"].joint_names = ( + self.robot_cfg.regex_joints_arm + ) + if self.env_cfg.scenario == env_utils.Scenario.ORBIT: + # Fix the orientation of the light such that it fits with the orbital HDR + self.events.reset_rand_light_rot.params[ + "orientation_distribution_params" + ] = { + "roll": (-20.0 * torch.pi / 180.0,) * 2, + "pitch": (50.0 * torch.pi / 180.0,) * 2, + } diff --git a/space_robotics_bench/envs/manipulation/extensions/__init__.py b/space_robotics_bench/envs/manipulation/extensions/__init__.py new file mode 100644 index 0000000..74ff8c6 --- /dev/null +++ b/space_robotics_bench/envs/manipulation/extensions/__init__.py @@ -0,0 +1 @@ +from .visual import * # noqa: F403 diff --git a/space_robotics_bench/envs/manipulation/extensions/visual/__init__.py b/space_robotics_bench/envs/manipulation/extensions/visual/__init__.py new file mode 100644 index 0000000..623a285 --- /dev/null +++ b/space_robotics_bench/envs/manipulation/extensions/visual/__init__.py @@ -0,0 +1,2 @@ +from .cfg import * # noqa: F403 +from .impl import * # noqa: F403 diff --git a/space_robotics_bench/envs/manipulation/extensions/visual/cfg.py b/space_robotics_bench/envs/manipulation/extensions/visual/cfg.py new file mode 100644 index 0000000..c56e384 --- /dev/null +++ b/space_robotics_bench/envs/manipulation/extensions/visual/cfg.py @@ -0,0 +1,110 @@ +from dataclasses import MISSING +from typing import Optional, Tuple + +from omni.isaac.lab.envs import ViewerCfg +from omni.isaac.lab.scene import InteractiveSceneCfg +from omni.isaac.lab.sensors import CameraCfg +from omni.isaac.lab.sensors.camera.camera_cfg import PinholeCameraCfg +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.assets as asset_utils +import space_robotics_bench.utils.math as math_utils + + +@configclass +class VisualManipulationEnvExtCfg: + ## Subclass requirements + agent_rate: int = MISSING + scene: InteractiveSceneCfg = MISSING + viewer: ViewerCfg = MISSING + robot_cfg: asset_utils.ManipulatorCfg = MISSING + vehicle_cfg: Optional[asset_utils.VehicleCfg] = None + + ## Enabling flags + enable_camera_scene: bool = False + enable_camera_base: bool = True + enable_camera_wrist: bool = True + + ## Resolution + camera_resolution: Tuple[int, int] = (64, 64) + camera_framerate: int = 0 # 0 matches the agent rate + + def __post_init__(self): + ## Scene + # self.scene.env_spacing += 4.0 + + ## Sensors + framerate = ( + self.camera_framerate if self.camera_framerate > 0 else self.agent_rate + ) + # Scene camera + if self.enable_camera_scene: + self.scene.camera_scene = CameraCfg( + prim_path="{ENV_REGEX_NS}/camera_scene", + offset=CameraCfg.OffsetCfg( + convention="world", + pos=(1.2, 0.0, 0.8), + rot=math_utils.quat_from_rpy(0.0, 30.0, 180.0), + ), + spawn=PinholeCameraCfg( + clipping_range=(0.01, 4.0 - 0.01), + ), + width=self.camera_resolution[0], + height=self.camera_resolution[1], + update_period=framerate, + data_types=["rgb", "distance_to_camera"], + ) + + # Robot base camera + if self.enable_camera_base: + camera_base_kwargs = { + "spawn": PinholeCameraCfg( + focal_length=5.0, + horizontal_aperture=12.0, + clipping_range=(0.001, 2.5 - 0.001), + ), + "width": self.camera_resolution[0], + "height": self.camera_resolution[1], + "update_period": framerate, + "data_types": ["rgb", "distance_to_camera"], + } + if self.vehicle_cfg and self.vehicle_cfg.frame_camera_base: + self.scene.camera_base = CameraCfg( + prim_path=f"{self.scene.vehicle.prim_path}/{self.vehicle_cfg.frame_camera_base.prim_relpath}", + offset=CameraCfg.OffsetCfg( + convention="world", + pos=self.vehicle_cfg.frame_camera_base.offset.translation, + rot=self.vehicle_cfg.frame_camera_base.offset.rotation, + ), + **camera_base_kwargs, + ) + else: + self.scene.camera_base = CameraCfg( + prim_path=f"{self.scene.robot.prim_path}/{self.robot_cfg.frame_camera_base.prim_relpath}", + offset=CameraCfg.OffsetCfg( + convention="world", + pos=self.robot_cfg.frame_camera_base.offset.translation, + rot=self.robot_cfg.frame_camera_base.offset.rotation, + ), + **camera_base_kwargs, + ) + + # Wrist camera + if self.enable_camera_wrist: + self.scene.camera_wrist = CameraCfg( + prim_path=f"{self.scene.robot.prim_path}/{self.robot_cfg.frame_camera_wrist.prim_relpath}", + offset=CameraCfg.OffsetCfg( + convention="world", + pos=self.robot_cfg.frame_camera_wrist.offset.translation, + rot=self.robot_cfg.frame_camera_wrist.offset.rotation, + ), + spawn=PinholeCameraCfg( + focal_length=10.0, + horizontal_aperture=16.0, + clipping_range=(0.001, 1.5 - 0.001), + ), + width=self.camera_resolution[0], + height=self.camera_resolution[1], + update_period=framerate, + data_types=["rgb", "distance_to_camera"], + ) diff --git a/space_robotics_bench/envs/manipulation/extensions/visual/impl.py b/space_robotics_bench/envs/manipulation/extensions/visual/impl.py new file mode 100644 index 0000000..17525a7 --- /dev/null +++ b/space_robotics_bench/envs/manipulation/extensions/visual/impl.py @@ -0,0 +1,46 @@ +from typing import Dict, Tuple + +import torch +from omni.isaac.lab.scene import InteractiveScene +from omni.isaac.lab.sensors import Camera + +from space_robotics_bench.utils import image_proc +from space_robotics_bench.utils import string as string_utils + +from .cfg import VisualManipulationEnvExtCfg + + +class VisualManipulationEnvExt: + ## Subclass requirements + common_step_counter: int + scene: InteractiveScene + cfg: VisualManipulationEnvExtCfg + + def __init__(self, cfg: VisualManipulationEnvExtCfg, **kwargs): + ## Extract camera sensors from the scene + self.__cameras: Dict[ + str, # Name of the output image + Tuple[ + Camera, # Camera sensor + Tuple[float, float], # Depth range + ], + ] = { + f"image_{string_utils.sanitize_camera_name(key)}": ( + sensor, + getattr(cfg.scene, key).spawn.clipping_range, + ) + for key, sensor in self.scene._sensors.items() + if type(sensor) == Camera + } + + def _get_observations(self) -> Dict[str, torch.Tensor]: + observation = {} + for image_name, (sensor, depth_range) in self.__cameras.items(): + observation.update( + image_proc.construct_observation( + **image_proc.extract_images(sensor), + depth_range=depth_range, + image_name=image_name, + ) + ) + return observation diff --git a/space_robotics_bench/envs/manipulation/impl.py b/space_robotics_bench/envs/manipulation/impl.py new file mode 100644 index 0000000..ef1b957 --- /dev/null +++ b/space_robotics_bench/envs/manipulation/impl.py @@ -0,0 +1,21 @@ +from omni.isaac.lab.sensors import ContactSensor +from omni.isaac.lab.sensors.frame_transformer.frame_transformer_cfg import ( + FrameTransformer, +) + +from space_robotics_bench.core.assets import Articulation +from space_robotics_bench.core.envs import BaseEnv + +from .cfg import BaseManipulationEnvCfg + + +class BaseManipulationEnv(BaseEnv): + cfg: BaseManipulationEnvCfg + + def __init__(self, cfg: BaseManipulationEnvCfg, **kwargs): + super().__init__(cfg, **kwargs) + + ## Get handles to scene assets + self._robot: Articulation = self.scene["robot"] + self._tf_robot_ee: FrameTransformer = self.scene["tf_robot_ee"] + self._contacts_robot: ContactSensor = self.scene["contacts_robot"] diff --git a/space_robotics_bench/envs/mobile_robotics/__init__.py b/space_robotics_bench/envs/mobile_robotics/__init__.py new file mode 100644 index 0000000..c1a49b6 --- /dev/null +++ b/space_robotics_bench/envs/mobile_robotics/__init__.py @@ -0,0 +1,4 @@ +from .cfg import * # noqa: F403 +from .impl import * # noqa: F403 + +from .extensions import * # noqa: F403 isort:skip diff --git a/space_robotics_bench/envs/mobile_robotics/cfg.py b/space_robotics_bench/envs/mobile_robotics/cfg.py new file mode 100644 index 0000000..19a24ed --- /dev/null +++ b/space_robotics_bench/envs/mobile_robotics/cfg.py @@ -0,0 +1,152 @@ +import math + +import torch +from omni.isaac.lab.envs import ViewerCfg +from omni.isaac.lab.managers import EventTermCfg, SceneEntityCfg +from omni.isaac.lab.scene import InteractiveSceneCfg +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.sim as sim_utils +from space_robotics_bench import assets +from space_robotics_bench.core import mdp +from space_robotics_bench.core.envs import BaseEnvCfg +from space_robotics_bench.core.sim import SimulationCfg + + +@configclass +class BaseMobileRoboticsEnvEventCfg: + ## Default scene reset + reset_all = EventTermCfg(func=mdp.reset_scene_to_default, mode="reset") + + ## Light + reset_rand_light_rot = EventTermCfg( + func=mdp.reset_xform_orientation_uniform, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("light"), + "orientation_distribution_params": { + "roll": ( + -75.0 * torch.pi / 180.0, + 75.0 * torch.pi / 180.0, + ), + "pitch": ( + -75.0 * torch.pi / 180.0, + 75.0 * torch.pi / 180.0, + ), + }, + }, + ) + + ## Robot + reset_rand_robot_state = EventTermCfg( + func=mdp.reset_root_state_uniform, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("robot"), + "pose_range": { + "yaw": ( + -torch.pi, + torch.pi, + ), + }, + "velocity_range": {}, + }, + ) + + +@configclass +class BaseMobileRoboticsEnvCfg(BaseEnvCfg): + ## Environment + episode_length_s: float = 50.0 + env_rate: float = 1.0 / 100.0 + + ## Agent + agent_rate: float = 1.0 / 50.0 + + ## Simulation + sim = SimulationCfg( + disable_contact_processing=True, + physx=sim_utils.PhysxCfg( + enable_ccd=False, + enable_stabilization=False, + bounce_threshold_velocity=0.0, + friction_correlation_distance=0.01, + min_velocity_iteration_count=1, + # GPU settings + gpu_temp_buffer_capacity=2 ** (24 - 5), + gpu_max_rigid_contact_count=2 ** (22 - 4), + gpu_max_rigid_patch_count=2 ** (13 - 1), + gpu_heap_capacity=2 ** (26 - 6), + gpu_found_lost_pairs_capacity=2 ** (18 - 2), + gpu_found_lost_aggregate_pairs_capacity=2 ** (10 - 1), + gpu_total_aggregate_pairs_capacity=2 ** (10 - 1), + gpu_max_soft_body_contacts=2 ** (20 - 3), + gpu_max_particle_contacts=2 ** (20 - 3), + gpu_collision_stack_size=2 ** (26 - 4), + gpu_max_num_partitions=8, + ), + physics_material=sim_utils.RigidBodyMaterialCfg( + static_friction=1.0, + dynamic_friction=1.0, + restitution=0.0, + friction_combine_mode="multiply", + restitution_combine_mode="multiply", + ), + ) + + ## Viewer + viewer = ViewerCfg( + lookat=(0.0, 0.0, 0.0), + eye=(-7.5, 0.0, 10.0), + origin_type="env", + env_index=0, + ) + + ## Scene + scene = InteractiveSceneCfg(num_envs=1, env_spacing=65.0, replicate_physics=False) + + ## Events + events = BaseMobileRoboticsEnvEventCfg() + + def __post_init__(self): + super().__post_init__() + + ## Simulation + self.decimation = int(self.agent_rate / self.env_rate) + self.sim.dt = self.env_rate + self.sim.render_interval = self.decimation + self.sim.gravity = (0.0, 0.0, -self.env_cfg.scenario.gravity_magnitude) + # Increase GPU settings based on the number of environments + gpu_capacity_factor = self.scene.num_envs + self.sim.physx.gpu_heap_capacity *= gpu_capacity_factor + self.sim.physx.gpu_collision_stack_size *= gpu_capacity_factor + self.sim.physx.gpu_temp_buffer_capacity *= gpu_capacity_factor + self.sim.physx.gpu_max_rigid_contact_count *= gpu_capacity_factor + self.sim.physx.gpu_max_rigid_patch_count *= gpu_capacity_factor + self.sim.physx.gpu_found_lost_pairs_capacity *= gpu_capacity_factor + self.sim.physx.gpu_found_lost_aggregate_pairs_capacity *= gpu_capacity_factor + self.sim.physx.gpu_total_aggregate_pairs_capacity *= gpu_capacity_factor + self.sim.physx.gpu_max_soft_body_contacts *= gpu_capacity_factor + self.sim.physx.gpu_max_particle_contacts *= gpu_capacity_factor + self.sim.physx.gpu_max_num_partitions = min( + 2 ** math.floor(1.0 + math.pow(self.scene.num_envs, 0.2)), 32 + ) + + ## Scene + self.scene.light = assets.sunlight_from_env_cfg(self.env_cfg) + self.scene.sky = assets.sky_from_env_cfg(self.env_cfg) + self.scene.terrain = assets.terrain_from_env_cfg( + self.env_cfg, + num_assets=self.scene.num_envs, + size=(self.scene.env_spacing - 1,) * 2, + procgen_kwargs={ + "density": 0.16, + "flat_area_size": 4.0, + "texture_resolution": 4096, + }, + ) + self.robot_cfg = assets.rover_from_env_cfg(self.env_cfg) + self.scene.robot = self.robot_cfg.asset_cfg + + ## Actions + self.actions = self.robot_cfg.action_cfg diff --git a/space_robotics_bench/envs/mobile_robotics/extensions/__init__.py b/space_robotics_bench/envs/mobile_robotics/extensions/__init__.py new file mode 100644 index 0000000..74ff8c6 --- /dev/null +++ b/space_robotics_bench/envs/mobile_robotics/extensions/__init__.py @@ -0,0 +1 @@ +from .visual import * # noqa: F403 diff --git a/space_robotics_bench/envs/mobile_robotics/extensions/visual/__init__.py b/space_robotics_bench/envs/mobile_robotics/extensions/visual/__init__.py new file mode 100644 index 0000000..623a285 --- /dev/null +++ b/space_robotics_bench/envs/mobile_robotics/extensions/visual/__init__.py @@ -0,0 +1,2 @@ +from .cfg import * # noqa: F403 +from .impl import * # noqa: F403 diff --git a/space_robotics_bench/envs/mobile_robotics/extensions/visual/cfg.py b/space_robotics_bench/envs/mobile_robotics/extensions/visual/cfg.py new file mode 100644 index 0000000..8872f62 --- /dev/null +++ b/space_robotics_bench/envs/mobile_robotics/extensions/visual/cfg.py @@ -0,0 +1,74 @@ +from dataclasses import MISSING +from typing import Tuple + +from omni.isaac.lab.envs import ViewerCfg +from omni.isaac.lab.scene import InteractiveSceneCfg +from omni.isaac.lab.sensors import CameraCfg +from omni.isaac.lab.sensors.camera.camera_cfg import PinholeCameraCfg +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.assets as asset_utils +import space_robotics_bench.utils.math as math_utils + + +@configclass +class VisualMobileRoboticsEnvExtCfg: + ## Subclass requirements + agent_rate: int = MISSING + scene: InteractiveSceneCfg = MISSING + viewer: ViewerCfg = MISSING + robot_cfg: asset_utils.WheeledRoverCfg = MISSING + + ## Enabling flags + enable_camera_scene: bool = True + enable_camera_front: bool = True + + ## Resolution + camera_resolution: Tuple[int, int] = (64, 64) + camera_framerate: int = 0 # 0 matches the agent rate + + def __post_init__(self): + ## Scene + # self.scene.env_spacing += 4.0 + + ## Sensors + framerate = ( + self.camera_framerate if self.camera_framerate > 0 else self.agent_rate + ) + # Scene camera + if self.enable_camera_scene: + self.scene.camera_scene = CameraCfg( + prim_path=f"{self.scene.robot.prim_path}/{self.robot_cfg.frame_base.prim_relpath}/camera_scene", + offset=CameraCfg.OffsetCfg( + convention="world", + pos=(0.0, 7.5, 5.0), + rot=math_utils.quat_from_rpy(0.0, 30.0, -90.0), + ), + spawn=PinholeCameraCfg( + clipping_range=(0.01, 20.0 - 0.01), + ), + width=self.camera_resolution[0], + height=self.camera_resolution[1], + update_period=framerate, + data_types=["rgb", "distance_to_camera"], + ) + + # Front camera + if self.enable_camera_front: + self.scene.camera_front = CameraCfg( + prim_path=f"{self.scene.robot.prim_path}/{self.robot_cfg.frame_camera_front.prim_relpath}", + offset=CameraCfg.OffsetCfg( + convention="world", + pos=self.robot_cfg.frame_camera_front.offset.translation, + rot=self.robot_cfg.frame_camera_front.offset.rotation, + ), + spawn=PinholeCameraCfg( + focal_length=5.0, + horizontal_aperture=12.0, + clipping_range=(0.01, 20.0 - 0.01), + ), + width=self.camera_resolution[0], + height=self.camera_resolution[1], + update_period=framerate, + data_types=["rgb", "distance_to_camera"], + ) diff --git a/space_robotics_bench/envs/mobile_robotics/extensions/visual/impl.py b/space_robotics_bench/envs/mobile_robotics/extensions/visual/impl.py new file mode 100644 index 0000000..e1191a2 --- /dev/null +++ b/space_robotics_bench/envs/mobile_robotics/extensions/visual/impl.py @@ -0,0 +1,46 @@ +from typing import Dict, Tuple + +import torch +from omni.isaac.lab.scene import InteractiveScene +from omni.isaac.lab.sensors import Camera + +from space_robotics_bench.utils import image_proc +from space_robotics_bench.utils import string as string_utils + +from .cfg import VisualMobileRoboticsEnvExtCfg + + +class VisualMobileRoboticsEnvExt: + ## Subclass requirements + common_step_counter: int + scene: InteractiveScene + cfg: VisualMobileRoboticsEnvExtCfg + + def __init__(self, cfg: VisualMobileRoboticsEnvExtCfg, **kwargs): + ## Extract camera sensors from the scene + self.__cameras: Dict[ + str, # Name of the output image + Tuple[ + Camera, # Camera sensor + Tuple[float, float], # Depth range + ], + ] = { + f"image_{string_utils.sanitize_camera_name(key)}": ( + sensor, + getattr(cfg.scene, key).spawn.clipping_range, + ) + for key, sensor in self.scene._sensors.items() + if type(sensor) == Camera + } + + def _get_observations(self) -> Dict[str, torch.Tensor]: + observation = {} + for image_name, (sensor, depth_range) in self.__cameras.items(): + observation.update( + image_proc.construct_observation( + **image_proc.extract_images(sensor), + depth_range=depth_range, + image_name=image_name, + ) + ) + return observation diff --git a/space_robotics_bench/envs/mobile_robotics/impl.py b/space_robotics_bench/envs/mobile_robotics/impl.py new file mode 100644 index 0000000..30d2b15 --- /dev/null +++ b/space_robotics_bench/envs/mobile_robotics/impl.py @@ -0,0 +1,14 @@ +from space_robotics_bench.core.assets import Articulation +from space_robotics_bench.core.envs import BaseEnv + +from .cfg import BaseMobileRoboticsEnvCfg + + +class BaseMobileRoboticsEnv(BaseEnv): + cfg: BaseMobileRoboticsEnvCfg + + def __init__(self, cfg: BaseMobileRoboticsEnvCfg, **kwargs): + super().__init__(cfg, **kwargs) + + ## Get handles to scene assets + self._robot: Articulation = self.scene["robot"] diff --git a/space_robotics_bench/paths.py b/space_robotics_bench/paths.py new file mode 100644 index 0000000..d496e51 --- /dev/null +++ b/space_robotics_bench/paths.py @@ -0,0 +1,23 @@ +from os import path + +# Path to repository root directory +SRB_DIR = path.dirname(path.dirname(path.realpath(__file__))) + +# Path to assets directory +SRB_ASSETS_DIR = path.join(SRB_DIR, "assets") +SRB_ASSETS_DIR_SRB = path.join(SRB_ASSETS_DIR, "srb_assets") +SRB_ASSETS_DIR_SRB_HDRI = path.join(SRB_ASSETS_DIR_SRB, "hdri") +SRB_ASSETS_DIR_SRB_MODEL = path.join(SRB_ASSETS_DIR_SRB, "model") +SRB_ASSETS_DIR_SRB_OBJECT = path.join(SRB_ASSETS_DIR_SRB_MODEL, "object") +SRB_ASSETS_DIR_SRB_ROBOT = path.join(SRB_ASSETS_DIR_SRB_MODEL, "robot") +SRB_ASSETS_DIR_SRB_TERRAIN = path.join(SRB_ASSETS_DIR_SRB_MODEL, "terrain") +SRB_ASSETS_DIR_SRB_VEHICLE = path.join(SRB_ASSETS_DIR_SRB_MODEL, "vehicle") + +# Path to confdig directory +SRB_CONFIG_DIR = path.join(SRB_DIR, "config") + +# Path to hyperparameters directory +SRB_HYPERPARAMS_DIR = path.join(SRB_DIR, "hyperparams") + +# Path to scripts directory +SRB_SCRIPTS_DIR = path.join(SRB_DIR, "scripts") diff --git a/space_robotics_bench/py.typed b/space_robotics_bench/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/space_robotics_bench/tasks/__init__.py b/space_robotics_bench/tasks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/space_robotics_bench/tasks/_demos/__init__.py b/space_robotics_bench/tasks/_demos/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/space_robotics_bench/tasks/_demos/gateway/__init__.py b/space_robotics_bench/tasks/_demos/gateway/__init__.py new file mode 100644 index 0000000..b380c63 --- /dev/null +++ b/space_robotics_bench/tasks/_demos/gateway/__init__.py @@ -0,0 +1,17 @@ +from space_robotics_bench.utils.registry import register_tasks + +from .task import Task, TaskCfg +from .task_visual import VisualTask, VisualTaskCfg + +BASE_TASK_NAME = __name__.split(".")[-1] +register_tasks( + { + BASE_TASK_NAME: {}, + f"{BASE_TASK_NAME}_visual": { + "entry_point": VisualTask, + "task_cfg": VisualTaskCfg, + }, + }, + default_entry_point=Task, + default_task_cfg=TaskCfg, +) diff --git a/space_robotics_bench/tasks/_demos/gateway/task.py b/space_robotics_bench/tasks/_demos/gateway/task.py new file mode 100644 index 0000000..9cbe907 --- /dev/null +++ b/space_robotics_bench/tasks/_demos/gateway/task.py @@ -0,0 +1,171 @@ +import sys +from typing import Dict, Sequence, Tuple + +import torch +from omni.isaac.lab.envs import ViewerCfg +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.envs as env_utils +from space_robotics_bench import assets +from space_robotics_bench.core.assets import Articulation +from space_robotics_bench.core.envs import BaseEnv +from space_robotics_bench.envs import BaseManipulationEnv, BaseManipulationEnvCfg + +############## +### Config ### +############## + + +@configclass +class TaskCfg(BaseManipulationEnvCfg): + viewer = ViewerCfg( + lookat=(0.0, 0.0, 2.5), + eye=(15.0, 0.0, 12.5), + origin_type="env", + env_index=0, + ) + + def __post_init__(self): + if self.env_cfg.scenario != env_utils.Scenario.ORBIT: + print( + f"[WARN] Environment requires ORBIT scenario ({self.env_cfg.scenario} ignored)", + file=sys.stderr, + ) + self.env_cfg.scenario = env_utils.Scenario.ORBIT + if self.env_cfg.assets.terrain.variant != env_utils.AssetVariant.NONE: + print( + f"[WARN] Environment requires NONE terrain ({self.env_cfg.assets.terrain.variant} ignored)", + file=sys.stderr, + ) + self.env_cfg.assets.terrain.variant = env_utils.AssetVariant.NONE + + super().__post_init__() + + ## Scene + self.scene.env_spacing = 42.0 + self.robot_cfg = assets.canadarm3_large_cfg() + self.scene.robot = self.robot_cfg.asset_cfg + + ## Actions + self.actions = self.robot_cfg.action_cfg + + ## Sensors + self.scene.tf_robot_ee = None + self.scene.contacts_robot = None + + ## Events + self.events.reset_rand_robot_state.params["asset_cfg"].joint_names = ( + self.robot_cfg.regex_joints_arm + ) + + +############ +### Task ### +############ + + +class Task(BaseManipulationEnv): + cfg: TaskCfg + + def __init__(self, cfg: TaskCfg, **kwargs): + # super().__init__(cfg, **kwargs) + BaseEnv.__init__(self, cfg, **kwargs) + + ## Get handles to scene assets + self._robot: Articulation = self.scene["robot"] + + ## Pre-compute metrics used in hot loops + self._max_episode_length = self.max_episode_length + + ## Initialize the intermediate state + self._update_intermediate_state() + + def _reset_idx(self, env_ids: Sequence[int]): + super()._reset_idx(env_ids) + + def _get_dones(self) -> Tuple[torch.Tensor, torch.Tensor]: + # Note: This assumes that `_get_dones()` is called before `_get_rewards()` and `_get_observations()` in `step()` + self._update_intermediate_state() + + if not self.cfg.enable_truncation: + self._truncations = torch.zeros_like(self._truncations) + + return self._terminations, self._truncations + + def _get_rewards(self) -> torch.Tensor: + return self._rewards + + def _get_observations(self) -> Dict[str, torch.Tensor]: + return { + "robot_joint_pos": self._robot.data.joint_pos, + } + + ######################## + ### Helper Functions ### + ######################## + + def _update_intermediate_state(self): + ## Compute other intermediate states + ( + self._remaining_time, + self._rewards, + self._terminations, + self._truncations, + ) = _compute_intermediate_state( + current_action=self.action_manager.action, + previous_action=self.action_manager.prev_action, + episode_length_buf=self.episode_length_buf, + max_episode_length=self._max_episode_length, + ) + + +############################# +### TorchScript functions ### +############################# + + +@torch.jit.script +def _compute_intermediate_state( + *, + current_action: torch.Tensor, + previous_action: torch.Tensor, + episode_length_buf: torch.Tensor, + max_episode_length: int, +) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + ## Intermediate states + # Time + remaining_time = 1 - (episode_length_buf / max_episode_length).unsqueeze(-1) + + ## Rewards + # Penalty: Action rate + WEIGHT_ACTION_RATE = -0.05 + penalty_action_rate = WEIGHT_ACTION_RATE * torch.sum( + torch.square(current_action - previous_action), dim=1 + ) + + # Total reward + rewards = torch.sum( + torch.stack( + [ + penalty_action_rate, + ], + dim=-1, + ), + dim=-1, + ) + + ## Termination and truncation + truncations = episode_length_buf > (max_episode_length - 1) + terminations = torch.zeros_like(truncations) + + return ( + remaining_time, + rewards, + terminations, + truncations, + ) diff --git a/space_robotics_bench/tasks/_demos/gateway/task_visual.py b/space_robotics_bench/tasks/_demos/gateway/task_visual.py new file mode 100644 index 0000000..9b765bc --- /dev/null +++ b/space_robotics_bench/tasks/_demos/gateway/task_visual.py @@ -0,0 +1,32 @@ +from typing import Dict + +import torch +from omni.isaac.lab.utils import configclass + +from space_robotics_bench.envs import ( + VisualManipulationEnvExt, + VisualManipulationEnvExtCfg, +) + +from .task import Task, TaskCfg + + +@configclass +class VisualTaskCfg(TaskCfg, VisualManipulationEnvExtCfg): + def __post_init__(self): + TaskCfg.__post_init__(self) + VisualManipulationEnvExtCfg.__post_init__(self) + + +class VisualTask(Task, VisualManipulationEnvExt): + cfg: VisualTaskCfg + + def __init__(self, cfg: VisualTaskCfg, **kwargs): + Task.__init__(self, cfg, **kwargs) + VisualManipulationEnvExt.__init__(self, cfg, **kwargs) + + def _get_observations(self) -> Dict[str, torch.Tensor]: + return { + **Task._get_observations(self), + **VisualManipulationEnvExt._get_observations(self), + } diff --git a/space_robotics_bench/tasks/_demos/ingenuity/__init__.py b/space_robotics_bench/tasks/_demos/ingenuity/__init__.py new file mode 100644 index 0000000..b380c63 --- /dev/null +++ b/space_robotics_bench/tasks/_demos/ingenuity/__init__.py @@ -0,0 +1,17 @@ +from space_robotics_bench.utils.registry import register_tasks + +from .task import Task, TaskCfg +from .task_visual import VisualTask, VisualTaskCfg + +BASE_TASK_NAME = __name__.split(".")[-1] +register_tasks( + { + BASE_TASK_NAME: {}, + f"{BASE_TASK_NAME}_visual": { + "entry_point": VisualTask, + "task_cfg": VisualTaskCfg, + }, + }, + default_entry_point=Task, + default_task_cfg=TaskCfg, +) diff --git a/space_robotics_bench/tasks/_demos/ingenuity/task.py b/space_robotics_bench/tasks/_demos/ingenuity/task.py new file mode 100644 index 0000000..19351ee --- /dev/null +++ b/space_robotics_bench/tasks/_demos/ingenuity/task.py @@ -0,0 +1,130 @@ +from typing import Dict, Sequence, Tuple + +import torch +from omni.isaac.lab.utils import configclass + +from space_robotics_bench.envs import BaseAerialRoboticsEnv, BaseAerialRoboticsEnvCfg + +############## +### Config ### +############## + + +@configclass +class TaskCfg(BaseAerialRoboticsEnvCfg): + def __post_init__(self): + super().__post_init__() + + +############ +### Task ### +############ + + +class Task(BaseAerialRoboticsEnv): + cfg: TaskCfg + + def __init__(self, cfg: TaskCfg, **kwargs): + super().__init__(cfg, **kwargs) + + ## Pre-compute metrics used in hot loops + self._max_episode_length = self.max_episode_length + + ## Initialize the intermediate state + self._update_intermediate_state() + + def _reset_idx(self, env_ids: Sequence[int]): + super()._reset_idx(env_ids) + + def _get_dones(self) -> Tuple[torch.Tensor, torch.Tensor]: + # Note: This assumes that `_get_dones()` is called before `_get_rewards()` and `_get_observations()` in `step()` + self._update_intermediate_state() + + if not self.cfg.enable_truncation: + self._truncations = torch.zeros_like(self._truncations) + + return self._terminations, self._truncations + + def _get_rewards(self) -> torch.Tensor: + return self._rewards + + def _get_observations(self) -> Dict[str, torch.Tensor]: + return { + "robot_base_pose_w": torch.cat( + [ + self._robot.data.body_pos_w[:, 0], + self._robot.data.body_quat_w[:, 0], + ], + dim=-1, + ) + } + + ######################## + ### Helper Functions ### + ######################## + + def _update_intermediate_state(self): + ## Compute other intermediate states + ( + self._remaining_time, + self._rewards, + self._terminations, + self._truncations, + ) = _compute_intermediate_state( + current_action=self.action_manager.action, + previous_action=self.action_manager.prev_action, + episode_length_buf=self.episode_length_buf, + max_episode_length=self._max_episode_length, + ) + + +############################# +### TorchScript functions ### +############################# + + +@torch.jit.script +def _compute_intermediate_state( + *, + current_action: torch.Tensor, + previous_action: torch.Tensor, + episode_length_buf: torch.Tensor, + max_episode_length: int, +) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + ## Intermediate states + # Time + remaining_time = 1 - (episode_length_buf / max_episode_length).unsqueeze(-1) + + ## Rewards + # Penalty: Action rate + WEIGHT_ACTION_RATE = -0.05 + penalty_action_rate = WEIGHT_ACTION_RATE * torch.sum( + torch.square(current_action - previous_action), dim=1 + ) + + # Total reward + rewards = torch.sum( + torch.stack( + [ + penalty_action_rate, + ], + dim=-1, + ), + dim=-1, + ) + + ## Termination and truncation + truncations = episode_length_buf > (max_episode_length - 1) + terminations = torch.zeros_like(truncations) + + return ( + remaining_time, + rewards, + terminations, + truncations, + ) diff --git a/space_robotics_bench/tasks/_demos/ingenuity/task_visual.py b/space_robotics_bench/tasks/_demos/ingenuity/task_visual.py new file mode 100644 index 0000000..2356faf --- /dev/null +++ b/space_robotics_bench/tasks/_demos/ingenuity/task_visual.py @@ -0,0 +1,32 @@ +from typing import Dict + +import torch +from omni.isaac.lab.utils import configclass + +from space_robotics_bench.envs import ( + VisualAerialRoboticsEnvExt, + VisualAerialRoboticsEnvExtCfg, +) + +from .task import Task, TaskCfg + + +@configclass +class VisualTaskCfg(TaskCfg, VisualAerialRoboticsEnvExtCfg): + def __post_init__(self): + TaskCfg.__post_init__(self) + VisualAerialRoboticsEnvExtCfg.__post_init__(self) + + +class VisualTask(Task, VisualAerialRoboticsEnvExt): + cfg: VisualTaskCfg + + def __init__(self, cfg: VisualTaskCfg, **kwargs): + Task.__init__(self, cfg, **kwargs) + VisualAerialRoboticsEnvExt.__init__(self, cfg, **kwargs) + + def _get_observations(self) -> Dict[str, torch.Tensor]: + return { + **Task._get_observations(self), + **VisualAerialRoboticsEnvExt._get_observations(self), + } diff --git a/space_robotics_bench/tasks/_demos/perseverance/__init__.py b/space_robotics_bench/tasks/_demos/perseverance/__init__.py new file mode 100644 index 0000000..b380c63 --- /dev/null +++ b/space_robotics_bench/tasks/_demos/perseverance/__init__.py @@ -0,0 +1,17 @@ +from space_robotics_bench.utils.registry import register_tasks + +from .task import Task, TaskCfg +from .task_visual import VisualTask, VisualTaskCfg + +BASE_TASK_NAME = __name__.split(".")[-1] +register_tasks( + { + BASE_TASK_NAME: {}, + f"{BASE_TASK_NAME}_visual": { + "entry_point": VisualTask, + "task_cfg": VisualTaskCfg, + }, + }, + default_entry_point=Task, + default_task_cfg=TaskCfg, +) diff --git a/space_robotics_bench/tasks/_demos/perseverance/task.py b/space_robotics_bench/tasks/_demos/perseverance/task.py new file mode 100644 index 0000000..3c12055 --- /dev/null +++ b/space_robotics_bench/tasks/_demos/perseverance/task.py @@ -0,0 +1,130 @@ +from typing import Dict, Sequence, Tuple + +import torch +from omni.isaac.lab.utils import configclass + +from space_robotics_bench.envs import BaseMobileRoboticsEnv, BaseMobileRoboticsEnvCfg + +############## +### Config ### +############## + + +@configclass +class TaskCfg(BaseMobileRoboticsEnvCfg): + def __post_init__(self): + super().__post_init__() + + +############ +### Task ### +############ + + +class Task(BaseMobileRoboticsEnv): + cfg: TaskCfg + + def __init__(self, cfg: TaskCfg, **kwargs): + super().__init__(cfg, **kwargs) + + ## Pre-compute metrics used in hot loops + self._max_episode_length = self.max_episode_length + + ## Initialize the intermediate state + self._update_intermediate_state() + + def _reset_idx(self, env_ids: Sequence[int]): + super()._reset_idx(env_ids) + + def _get_dones(self) -> Tuple[torch.Tensor, torch.Tensor]: + # Note: This assumes that `_get_dones()` is called before `_get_rewards()` and `_get_observations()` in `step()` + self._update_intermediate_state() + + if not self.cfg.enable_truncation: + self._truncations = torch.zeros_like(self._truncations) + + return self._terminations, self._truncations + + def _get_rewards(self) -> torch.Tensor: + return self._rewards + + def _get_observations(self) -> Dict[str, torch.Tensor]: + return { + "robot_base_pose_w": torch.cat( + [ + self._robot.data.body_pos_w[:, 0], + self._robot.data.body_quat_w[:, 0], + ], + dim=-1, + ) + } + + ######################## + ### Helper Functions ### + ######################## + + def _update_intermediate_state(self): + ## Compute other intermediate states + ( + self._remaining_time, + self._rewards, + self._terminations, + self._truncations, + ) = _compute_intermediate_state( + current_action=self.action_manager.action, + previous_action=self.action_manager.prev_action, + episode_length_buf=self.episode_length_buf, + max_episode_length=self._max_episode_length, + ) + + +############################# +### TorchScript functions ### +############################# + + +@torch.jit.script +def _compute_intermediate_state( + *, + current_action: torch.Tensor, + previous_action: torch.Tensor, + episode_length_buf: torch.Tensor, + max_episode_length: int, +) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + ## Intermediate states + # Time + remaining_time = 1 - (episode_length_buf / max_episode_length).unsqueeze(-1) + + ## Rewards + # Penalty: Action rate + WEIGHT_ACTION_RATE = -0.05 + penalty_action_rate = WEIGHT_ACTION_RATE * torch.sum( + torch.square(current_action - previous_action), dim=1 + ) + + # Total reward + rewards = torch.sum( + torch.stack( + [ + penalty_action_rate, + ], + dim=-1, + ), + dim=-1, + ) + + ## Termination and truncation + truncations = episode_length_buf > (max_episode_length - 1) + terminations = torch.zeros_like(truncations) + + return ( + remaining_time, + rewards, + terminations, + truncations, + ) diff --git a/space_robotics_bench/tasks/_demos/perseverance/task_visual.py b/space_robotics_bench/tasks/_demos/perseverance/task_visual.py new file mode 100644 index 0000000..07b8758 --- /dev/null +++ b/space_robotics_bench/tasks/_demos/perseverance/task_visual.py @@ -0,0 +1,32 @@ +from typing import Dict + +import torch +from omni.isaac.lab.utils import configclass + +from space_robotics_bench.envs import ( + VisualMobileRoboticsEnvExt, + VisualMobileRoboticsEnvExtCfg, +) + +from .task import Task, TaskCfg + + +@configclass +class VisualTaskCfg(TaskCfg, VisualMobileRoboticsEnvExtCfg): + def __post_init__(self): + TaskCfg.__post_init__(self) + VisualMobileRoboticsEnvExtCfg.__post_init__(self) + + +class VisualTask(Task, VisualMobileRoboticsEnvExt): + cfg: VisualTaskCfg + + def __init__(self, cfg: VisualTaskCfg, **kwargs): + Task.__init__(self, cfg, **kwargs) + VisualMobileRoboticsEnvExt.__init__(self, cfg, **kwargs) + + def _get_observations(self) -> Dict[str, torch.Tensor]: + return { + **Task._get_observations(self), + **VisualMobileRoboticsEnvExt._get_observations(self), + } diff --git a/space_robotics_bench/tasks/debris_capture/__init__.py b/space_robotics_bench/tasks/debris_capture/__init__.py new file mode 100644 index 0000000..b380c63 --- /dev/null +++ b/space_robotics_bench/tasks/debris_capture/__init__.py @@ -0,0 +1,17 @@ +from space_robotics_bench.utils.registry import register_tasks + +from .task import Task, TaskCfg +from .task_visual import VisualTask, VisualTaskCfg + +BASE_TASK_NAME = __name__.split(".")[-1] +register_tasks( + { + BASE_TASK_NAME: {}, + f"{BASE_TASK_NAME}_visual": { + "entry_point": VisualTask, + "task_cfg": VisualTaskCfg, + }, + }, + default_entry_point=Task, + default_task_cfg=TaskCfg, +) diff --git a/space_robotics_bench/tasks/debris_capture/task.py b/space_robotics_bench/tasks/debris_capture/task.py new file mode 100644 index 0000000..6295159 --- /dev/null +++ b/space_robotics_bench/tasks/debris_capture/task.py @@ -0,0 +1,469 @@ +import sys +from typing import Any, Dict, List, Optional, Sequence, Tuple + +import torch +from omni.isaac.lab.managers import EventTermCfg, SceneEntityCfg +from omni.isaac.lab.sensors import ContactSensor, ContactSensorCfg +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.envs as env_utils +import space_robotics_bench.utils.math as math_utils +from space_robotics_bench import assets +from space_robotics_bench.core.assets import AssetCfg, RigidObject, RigidObjectCfg +from space_robotics_bench.envs import ( + BaseManipulationEnv, + BaseManipulationEnvCfg, + BaseManipulationEnvEventCfg, + mdp, +) + +############## +### Config ### +############## + + +class DebrisCfg(AssetCfg): + class Config: + arbitrary_types_allowed = True # Due to EventTermCfg + + ## Model + asset_cfg: RigidObjectCfg + + ## Randomization + state_randomizer: EventTermCfg + + +@configclass +class TaskCfg(BaseManipulationEnvCfg): + ## Environment + episode_length_s: float = 10.0 + + ## Task + is_finite_horizon: bool = False + + ## Events + @configclass + class EventCfg(BaseManipulationEnvEventCfg): + ## Object + reset_rand_object_state: Optional[EventTermCfg] = None + + events = EventCfg() + + def __post_init__(self): + if self.env_cfg.scenario != env_utils.Scenario.ORBIT: + print( + f"[WARN] Environment requires ORBIT scenario ({self.env_cfg.scenario} ignored)", + file=sys.stderr, + ) + self.env_cfg.scenario = env_utils.Scenario.ORBIT + if self.env_cfg.assets.terrain.variant != env_utils.AssetVariant.NONE: + print( + f"[WARN] Environment requires NONE terrain ({self.env_cfg.assets.terrain.variant} ignored)", + file=sys.stderr, + ) + self.env_cfg.assets.terrain.variant = env_utils.AssetVariant.NONE + + super().__post_init__() + + ## Simulation + self.sim.gravity = (0.0, 0.0, 0.0) + + ## Scene + self.object_cfg = self._object_cfg( + self.env_cfg, + num_assets=self.scene.num_envs, + init_state=RigidObjectCfg.InitialStateCfg(pos=(1.0, 0.0, 0.5)), + spawn_kwargs={ + "activate_contact_sensors": True, + }, + ) + self.scene.object = self.object_cfg.asset_cfg + + ## Sensors + self.scene.contacts_robot_hand_obj = ContactSensorCfg( + prim_path=f"{self.scene.robot.prim_path}/{self.robot_cfg.regex_links_hand}", + update_period=0.0, + # Note: This causes error 'Filter pattern did not match the correct number of entries' + # However, it seems to function properly anyway... + filter_prim_paths_expr=[self.scene.object.prim_path], + ) + + ## Events + self.events.reset_rand_object_state = self.object_cfg.state_randomizer + + ######################## + ### Helper Functions ### + ######################## + + @staticmethod + def _object_cfg( + env_cfg: env_utils.EnvironmentConfig, + *, + num_assets: int, + prim_path: str = "{ENV_REGEX_NS}/sample", + asset_cfg: SceneEntityCfg = SceneEntityCfg("object"), + spawn_kwargs: Dict[str, Any] = {}, + **kwargs, + ) -> DebrisCfg: + return DebrisCfg( + asset_cfg=assets.object_of_interest_from_env_cfg( + env_cfg, + num_assets=num_assets, + prim_path=prim_path, + spawn_kwargs=spawn_kwargs, + **kwargs, + ), + state_randomizer=EventTermCfg( + func=mdp.reset_root_state_uniform, + mode="reset", + params={ + "asset_cfg": asset_cfg, + "pose_range": { + "x": (-0.25, 0.25), + "y": (-0.25, 0.25), + "z": (-0.25, 0.25), + "roll": (-torch.pi, torch.pi), + "pitch": (-torch.pi, torch.pi), + "yaw": (-torch.pi, torch.pi), + }, + "velocity_range": { + "x": (-0.2 - 0.05, -0.2 + 0.05), + "y": (-0.05, 0.05), + "z": (-0.05, 0.05), + "roll": (-torch.pi, torch.pi), + "pitch": (-torch.pi, torch.pi), + "yaw": (-torch.pi, torch.pi), + }, + }, + ), + ) + + +############ +### Task ### +############ + + +class Task(BaseManipulationEnv): + cfg: TaskCfg + + def __init__(self, cfg: TaskCfg, **kwargs): + super().__init__(cfg, **kwargs) + + ## Get handles to scene assets + self._contacts_robot_hand_obj: ContactSensor = self.scene[ + "contacts_robot_hand_obj" + ] + self._object: RigidObject = self.scene["object"] + + ## Pre-compute metrics used in hot loops + self._robot_arm_joint_indices, _ = self._robot.find_joints( + self.cfg.robot_cfg.regex_joints_arm + ) + self._robot_hand_joint_indices, _ = self._robot.find_joints( + self.cfg.robot_cfg.regex_joints_hand + ) + self._max_episode_length = self.max_episode_length + self._obj_com_offset = self._object.data._root_physx_view.get_coms().to( + self.device + ) + + ## Initialize the intermediate state + self._update_intermediate_state() + + def _reset_idx(self, env_ids: Sequence[int]): + super()._reset_idx(env_ids) + + def _get_dones(self) -> Tuple[torch.Tensor, torch.Tensor]: + # Note: This assumes that `_get_dones()` is called before `_get_rewards()` and `_get_observations()` in `step()` + self._update_intermediate_state() + + if not self.cfg.enable_truncation: + self._truncations = torch.zeros_like(self._truncations) + + return self._terminations, self._truncations + + def _get_rewards(self) -> torch.Tensor: + return self._rewards + + def _get_observations(self) -> Dict[str, torch.Tensor]: + return _construct_observations( + remaining_time=self._remaining_time, + robot_joint_pos_arm=self._robot_joint_pos_arm, + robot_joint_pos_hand=self._robot_joint_pos_hand, + robot_ee_pos_wrt_base=self._robot_ee_pos_wrt_base, + robot_ee_rotmat_wrt_base=self._robot_ee_rotmat_wrt_base, + robot_hand_wrench=self._robot_hand_wrench, + obj_com_pos_wrt_robot_ee=self._obj_com_pos_wrt_robot_ee, + obj_com_rotmat_wrt_robot_ee=self._obj_com_rotmat_wrt_robot_ee, + obj_vel_w=self._obj_vel_w, + ) + + ######################## + ### Helper Functions ### + ######################## + + def _update_intermediate_state(self): + ## Extract intermediate states + self._robot_ee_pos_wrt_base = self._tf_robot_ee.data.target_pos_source[:, 0, :] + self._robot_hand_wrench = ( + self._robot.root_physx_view.get_link_incoming_joint_force()[ + :, self._robot_hand_joint_indices + ] + ) + self._obj_vel_w = self._object.data.root_vel_w + + ## Compute other intermediate states + ( + self._remaining_time, + self._robot_joint_pos_arm, + self._robot_joint_pos_hand, + self._robot_ee_rotmat_wrt_base, + self._obj_com_pos_wrt_robot_ee, + self._obj_com_rotmat_wrt_robot_ee, + self._rewards, + self._terminations, + self._truncations, + ) = _compute_intermediate_state( + current_action=self.action_manager.action, + previous_action=self.action_manager.prev_action, + episode_length_buf=self.episode_length_buf, + max_episode_length=self._max_episode_length, + robot_arm_joint_indices=self._robot_arm_joint_indices, + robot_hand_joint_indices=self._robot_hand_joint_indices, + joint_pos=self._robot.data.joint_pos, + soft_joint_pos_limits=self._robot.data.soft_joint_pos_limits, + robot_ee_pos_w=self._tf_robot_ee.data.target_pos_w[:, 0, :], + robot_ee_quat_w=self._tf_robot_ee.data.target_quat_w[:, 0, :], + robot_ee_quat_wrt_base=self._tf_robot_ee.data.target_quat_source[:, 0, :], + robot_arm_contact_net_forces=self._contacts_robot.data.net_forces_w, + robot_hand_obj_contact_force_matrix=self._contacts_robot_hand_obj.data.force_matrix_w, + obj_pos_w=self._object.data.root_pos_w, + obj_quat_w=self._object.data.root_quat_w, + obj_com_offset=self._obj_com_offset, + obj_vel_w=self._obj_vel_w, + ) + + +############################# +### TorchScript functions ### +############################# + + +@torch.jit.script +def _compute_intermediate_state( + *, + current_action: torch.Tensor, + previous_action: torch.Tensor, + episode_length_buf: torch.Tensor, + max_episode_length: int, + robot_arm_joint_indices: List[int], + robot_hand_joint_indices: List[int], + joint_pos: torch.Tensor, + soft_joint_pos_limits: torch.Tensor, + robot_ee_pos_w: torch.Tensor, + robot_ee_quat_w: torch.Tensor, + robot_ee_quat_wrt_base: torch.Tensor, + robot_arm_contact_net_forces: torch.Tensor, + robot_hand_obj_contact_force_matrix: torch.Tensor, + obj_pos_w: torch.Tensor, + obj_quat_w: torch.Tensor, + obj_com_offset: torch.Tensor, + obj_vel_w: torch.Tensor, +) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + ## Intermediate states + # Time + remaining_time = 1 - (episode_length_buf / max_episode_length).unsqueeze(-1) + + # Robot joint positions + joint_pos_normalized = math_utils.scale_transform( + joint_pos, + soft_joint_pos_limits[:, :, 0], + soft_joint_pos_limits[:, :, 1], + ) + robot_joint_pos_arm, robot_joint_pos_hand = ( + joint_pos_normalized[:, robot_arm_joint_indices], + joint_pos_normalized[:, robot_hand_joint_indices], + ) + + # End-effector pose (position and '6D' rotation) + robot_ee_rotmat_wrt_base = math_utils.matrix_from_quat(robot_ee_quat_wrt_base) + + # Transformation | Object origin -> Object CoM + obj_com_pos_w, obj_com_quat_w = math_utils.combine_frame_transforms( + t01=obj_pos_w, + q01=obj_quat_w, + t12=obj_com_offset[:, :3], + q12=obj_com_offset[:, 3:], + ) + + # Transformation | End-effector -> Object CoM + obj_com_pos_wrt_robot_ee, obj_com_quat_wrt_robot_ee = ( + math_utils.subtract_frame_transforms( + t01=robot_ee_pos_w, + q01=robot_ee_quat_w, + t02=obj_com_pos_w, + q02=obj_com_quat_w, + ) + ) + obj_com_rotmat_wrt_robot_ee = math_utils.matrix_from_quat(obj_com_quat_wrt_robot_ee) + + ## Rewards + # Penalty: Action rate + WEIGHT_ACTION_RATE = -0.05 + penalty_action_rate = WEIGHT_ACTION_RATE * torch.sum( + torch.square(current_action - previous_action), dim=1 + ) + + # Penalty: Undesired robot arm contacts + WEIGHT_UNDERSIRED_ROBOT_ARM_CONTACTS = -0.1 + THRESHOLD_UNDERSIRED_ROBOT_ARM_CONTACTS = 10.0 + penalty_undersired_robot_arm_contacts = WEIGHT_UNDERSIRED_ROBOT_ARM_CONTACTS * ( + torch.max(torch.norm(robot_arm_contact_net_forces, dim=-1), dim=1)[0] + > THRESHOLD_UNDERSIRED_ROBOT_ARM_CONTACTS + ) + + # Reward: Distance | End-effector <--> Object + WEIGHT_DISTANCE_EE_TO_OBJ = 1.0 + TANH_STD_DISTANCE_EE_TO_OBJ = 0.25 + reward_distance_ee_to_obj = WEIGHT_DISTANCE_EE_TO_OBJ * ( + 1.0 + - torch.tanh( + torch.norm(obj_com_pos_wrt_robot_ee, dim=-1) / TANH_STD_DISTANCE_EE_TO_OBJ + ) + ) + + # Reward: Object grasped + WEIGHT_OBJ_GRASPED = 4.0 + THRESHOLD_OBJ_GRASPED = 5.0 + reward_obj_grasped = WEIGHT_OBJ_GRASPED * ( + torch.mean( + torch.max(torch.norm(robot_hand_obj_contact_force_matrix, dim=-1), dim=-1)[ + 0 + ], + dim=1, + ) + > THRESHOLD_OBJ_GRASPED + ) + + # Penalty: Object velocity (linear) + WEIGHT_OBJ_VEL_LINEAR = -2.0 + penalty_obj_vel_linear = WEIGHT_OBJ_VEL_LINEAR * torch.norm( + obj_vel_w[:, :3], dim=-1 + ) + + # Penalty: Object velocity (angular) + WEIGHT_OBJ_VEL_ANGULAR = -1.0 / (2.0 * torch.pi) + penalty_obj_vel_angular = WEIGHT_OBJ_VEL_ANGULAR * torch.norm( + obj_vel_w[:, 3:], dim=-1 + ) + + # Total reward + rewards = torch.sum( + torch.stack( + [ + penalty_action_rate, + penalty_undersired_robot_arm_contacts, + reward_distance_ee_to_obj, + reward_obj_grasped, + penalty_obj_vel_linear, + penalty_obj_vel_angular, + ], + dim=-1, + ), + dim=-1, + ) + + ## Termination and truncation + truncations = episode_length_buf > (max_episode_length - 1) + terminations = torch.zeros_like(truncations) + + # print( + # f""" + # penalty | action_rate: {float(penalty_action_rate[0])} + # penalty | undersired_robot_arm_contacts: {float(penalty_undersired_robot_arm_contacts[0])} + # reward | distance_ee_to_obj: {float(reward_distance_ee_to_obj[0])} + # reward | obj_grasped: {float(reward_obj_grasped[0])} + # penalty | obj_vel_linear: {float(penalty_obj_vel_linear[0])} + # penalty | obj_vel_angular: {float(penalty_obj_vel_angular[0])} + # total: {float(rewards[0])} + # """ + # ) + + return ( + remaining_time, + robot_joint_pos_arm, + robot_joint_pos_hand, + robot_ee_rotmat_wrt_base, + obj_com_pos_wrt_robot_ee, + obj_com_rotmat_wrt_robot_ee, + rewards, + terminations, + truncations, + ) + + +@torch.jit.script +def _construct_observations( + *, + remaining_time: torch.Tensor, + robot_joint_pos_arm: torch.Tensor, + robot_joint_pos_hand: torch.Tensor, + robot_ee_pos_wrt_base: torch.Tensor, + robot_ee_rotmat_wrt_base: torch.Tensor, + robot_hand_wrench: torch.Tensor, + obj_com_pos_wrt_robot_ee: torch.Tensor, + obj_com_rotmat_wrt_robot_ee: torch.Tensor, + obj_vel_w: torch.Tensor, +) -> Dict[str, torch.Tensor]: + """ + Note: The `robot_hand_wrench` is considered as state (robot without force-torque sensors) + """ + + num_envs = remaining_time.size(0) + + # Robot joint positions + robot_joint_pos_hand_mean = robot_joint_pos_hand.mean(dim=-1, keepdim=True) + + # End-effector pose (position and '6D' rotation) + robot_ee_rot6d = math_utils.rotmat_to_rot6d(robot_ee_rotmat_wrt_base) + + # Wrench + robot_hand_wrench_full = robot_hand_wrench.view(num_envs, -1) + robot_hand_wrench_mean = robot_hand_wrench.mean(dim=1) + + # Transformation | End-effector -> Object CoM + obj_com_rot6d_wrt_robot_ee = math_utils.rotmat_to_rot6d(obj_com_rotmat_wrt_robot_ee) + + return { + "state": torch.cat( + [ + obj_com_pos_wrt_robot_ee, + obj_com_rot6d_wrt_robot_ee, + obj_vel_w, + robot_hand_wrench_mean, + ], + dim=-1, + ), + "state_dyn": torch.cat([robot_hand_wrench_full], dim=-1), + "proprio": torch.cat( + [ + remaining_time, + robot_ee_pos_wrt_base, + robot_ee_rot6d, + robot_joint_pos_hand_mean, + ], + dim=-1, + ), + "proprio_dyn": torch.cat([robot_joint_pos_arm, robot_joint_pos_hand], dim=-1), + } diff --git a/space_robotics_bench/tasks/debris_capture/task_visual.py b/space_robotics_bench/tasks/debris_capture/task_visual.py new file mode 100644 index 0000000..9b765bc --- /dev/null +++ b/space_robotics_bench/tasks/debris_capture/task_visual.py @@ -0,0 +1,32 @@ +from typing import Dict + +import torch +from omni.isaac.lab.utils import configclass + +from space_robotics_bench.envs import ( + VisualManipulationEnvExt, + VisualManipulationEnvExtCfg, +) + +from .task import Task, TaskCfg + + +@configclass +class VisualTaskCfg(TaskCfg, VisualManipulationEnvExtCfg): + def __post_init__(self): + TaskCfg.__post_init__(self) + VisualManipulationEnvExtCfg.__post_init__(self) + + +class VisualTask(Task, VisualManipulationEnvExt): + cfg: VisualTaskCfg + + def __init__(self, cfg: VisualTaskCfg, **kwargs): + Task.__init__(self, cfg, **kwargs) + VisualManipulationEnvExt.__init__(self, cfg, **kwargs) + + def _get_observations(self) -> Dict[str, torch.Tensor]: + return { + **Task._get_observations(self), + **VisualManipulationEnvExt._get_observations(self), + } diff --git a/space_robotics_bench/tasks/peg_in_hole/__init__.py b/space_robotics_bench/tasks/peg_in_hole/__init__.py new file mode 100644 index 0000000..b380c63 --- /dev/null +++ b/space_robotics_bench/tasks/peg_in_hole/__init__.py @@ -0,0 +1,17 @@ +from space_robotics_bench.utils.registry import register_tasks + +from .task import Task, TaskCfg +from .task_visual import VisualTask, VisualTaskCfg + +BASE_TASK_NAME = __name__.split(".")[-1] +register_tasks( + { + BASE_TASK_NAME: {}, + f"{BASE_TASK_NAME}_visual": { + "entry_point": VisualTask, + "task_cfg": VisualTaskCfg, + }, + }, + default_entry_point=Task, + default_task_cfg=TaskCfg, +) diff --git a/space_robotics_bench/tasks/peg_in_hole/task.py b/space_robotics_bench/tasks/peg_in_hole/task.py new file mode 100644 index 0000000..aa54e11 --- /dev/null +++ b/space_robotics_bench/tasks/peg_in_hole/task.py @@ -0,0 +1,673 @@ +from typing import Any, Dict, List, Optional, Sequence, Tuple + +import torch +from omni.isaac.core.prims.xform_prim_view import XFormPrimView +from omni.isaac.lab.managers import EventTermCfg, SceneEntityCfg +from omni.isaac.lab.sensors import ContactSensor, ContactSensorCfg +from omni.isaac.lab.utils import configclass +from pydantic import NonNegativeInt + +import space_robotics_bench.core.envs as env_utils +import space_robotics_bench.utils.math as math_utils +from space_robotics_bench import assets +from space_robotics_bench.core.assets import AssetCfg, RigidObject, RigidObjectCfg +from space_robotics_bench.envs import ( + BaseManipulationEnv, + BaseManipulationEnvCfg, + BaseManipulationEnvEventCfg, + mdp, +) + +############## +### Config ### +############## + + +class PegCfg(AssetCfg): + class Config: + arbitrary_types_allowed = True # Due to EventTermCfg + + ## Model + asset_cfg: RigidObjectCfg + + ## Geometry + offset_pos_ends: Tuple[ + Tuple[float, float, float], + Tuple[float, float, float], + ] + + ## Rotational symmetry of the peg represented as integer + # 0: Circle (infinite symmetry) + # 1: No symmetry (exactly one fit) + # n: n-fold symmetry (360/n deg between each symmetry) + rot_symmetry_n: NonNegativeInt = 1 + + ## Randomization + state_randomizer: Optional[EventTermCfg] = None + + +class HoleCfg(AssetCfg): + ## Geometry + offset_pos_bottom: Tuple[float, float, float] = (0.0, 0.0, 0.0) + offset_pos_entrance: Tuple[float, float, float] + + +def peg_and_hole_cfg( + env_cfg: env_utils.EnvironmentConfig, + *, + prim_path_peg: str = "{ENV_REGEX_NS}/peg", + prim_path_hole: str = "{ENV_REGEX_NS}/hole", + asset_cfg_peg: SceneEntityCfg = SceneEntityCfg("object"), + num_assets: int = 1, + size: Tuple[float, float, float] = (0.05, 0.05, 0.05), + spawn_kwargs_peg: Dict[str, Any] = {}, + spawn_kwargs_hole: Dict[str, Any] = {}, + procgen_seed_offset: int = 0, + procgen_kwargs_peg: Dict[str, Any] = {}, + procgen_kwargs_hole: Dict[str, Any] = {}, + **kwargs, +) -> Tuple[PegCfg, HoleCfg]: + pose_range_peg = { + "x": (-0.25 - 0.025, -0.25 + 0.0125), + "y": (-0.05, 0.05), + "roll": (torch.pi / 2, torch.pi / 2), + "yaw": ( + torch.pi / 2 - torch.pi / 16, + torch.pi / 2 + torch.pi / 16, + ), + } + rot_symmetry_n = 4 + offset_pos_ends = ((0.0, 0.0, 0.0), (0.0, 0.0, 0.2)) + offset_pos_entrance = (0.0, 0.0, 0.02) + + peg_cfg, hole_cfg = assets.peg_in_hole_from_env_cfg( + env_cfg, + prim_path_peg=prim_path_peg, + prim_path_hole=prim_path_hole, + num_assets=num_assets, + size=size, + spawn_kwargs_peg=spawn_kwargs_peg, + spawn_kwargs_hole=spawn_kwargs_hole, + procgen_seed_offset=procgen_seed_offset, + procgen_kwargs_peg=procgen_kwargs_peg, + procgen_kwargs_hole=procgen_kwargs_hole, + **kwargs, + ) + + return PegCfg( + asset_cfg=peg_cfg, + offset_pos_ends=offset_pos_ends, + rot_symmetry_n=rot_symmetry_n, + state_randomizer=EventTermCfg( + func=mdp.reset_root_state_uniform, + mode="reset", + params={ + "asset_cfg": asset_cfg_peg, + "pose_range": pose_range_peg, + "velocity_range": {}, + }, + ), + ), HoleCfg( + asset_cfg=hole_cfg, + offset_pos_entrance=offset_pos_entrance, + ) + + +@configclass +class TaskCfg(BaseManipulationEnvCfg): + ## Environment + episode_length_s: float = 10.0 + + ## Task + is_finite_horizon: bool = False + + ## Events + @configclass + class EventCfg(BaseManipulationEnvEventCfg): + ## Object + reset_rand_object_state: Optional[EventTermCfg] = None + + events = EventCfg() + + def __post_init__(self): + super().__post_init__() + + ## Scene + self.object_cfg, self.target_cfg = peg_and_hole_cfg( + self.env_cfg, + num_assets=self.scene.num_envs, + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.5, 0.0, 0.02)), + spawn_kwargs_peg={ + "activate_contact_sensors": True, + }, + ) + self.scene.object = self.object_cfg.asset_cfg + self.scene.target = self.target_cfg.asset_cfg + + ## Sensors + self.scene.contacts_robot_hand_obj = ContactSensorCfg( + prim_path=f"{self.scene.robot.prim_path}/{self.robot_cfg.regex_links_hand}", + update_period=0.0, + # Note: This causes error 'Filter pattern did not match the correct number of entries' + # However, it seems to function properly anyway... + filter_prim_paths_expr=[self.scene.object.prim_path], + ) + + ## Events + self.events.reset_rand_object_state = self.object_cfg.state_randomizer + + +############ +### Task ### +############ + + +class Task(BaseManipulationEnv): + cfg: TaskCfg + + def __init__(self, cfg: TaskCfg, **kwargs): + super().__init__(cfg, **kwargs) + + ## Get handles to scene assets + self._contacts_robot_hand_obj: ContactSensor = self.scene[ + "contacts_robot_hand_obj" + ] + self._object: RigidObject = self.scene["object"] + self._target: XFormPrimView = self.scene["target"] + + ## Pre-compute metrics used in hot loops + self._robot_arm_joint_indices, _ = self._robot.find_joints( + self.cfg.robot_cfg.regex_joints_arm + ) + self._robot_hand_joint_indices, _ = self._robot.find_joints( + self.cfg.robot_cfg.regex_joints_hand + ) + self._max_episode_length = self.max_episode_length + self._obj_com_offset = self._object.data._root_physx_view.get_coms().to( + self.device + ) + + ## Initialize buffers + self._initial_obj_height_w = torch.zeros( + self.num_envs, dtype=torch.float32, device=self.device + ) + self._peg_offset_pos_ends = torch.tensor( + self.cfg.object_cfg.offset_pos_ends, dtype=torch.float32, device=self.device + ).repeat(self.num_envs, 1, 1) + self._peg_rot_symmetry_n = torch.tensor( + self.cfg.object_cfg.rot_symmetry_n, dtype=torch.int32, device=self.device + ).repeat(self.num_envs) + self._hole_offset_pos_bottom = torch.tensor( + self.cfg.target_cfg.offset_pos_bottom, + dtype=torch.float32, + device=self.device, + ).repeat(self.num_envs, 1) + self._hole_offset_pos_entrance = torch.tensor( + self.cfg.target_cfg.offset_pos_entrance, + dtype=torch.float32, + device=self.device, + ).repeat(self.num_envs, 1) + + ## Initialize the intermediate state + self._update_intermediate_state() + + def _reset_idx(self, env_ids: Sequence[int]): + super()._reset_idx(env_ids) + + # Update the initial height of the objects + self._initial_obj_height_w[env_ids] = self._object.data.root_pos_w[env_ids, 2] + + def _get_dones(self) -> Tuple[torch.Tensor, torch.Tensor]: + # Note: This assumes that `_get_dones()` is called before `_get_rewards()` and `_get_observations()` in `step()` + self._update_intermediate_state() + + if not self.cfg.enable_truncation: + self._truncations = torch.zeros_like(self._truncations) + + return self._terminations, self._truncations + + def _get_rewards(self) -> torch.Tensor: + return self._rewards + + def _get_observations(self) -> Dict[str, torch.Tensor]: + return _construct_observations( + remaining_time=self._remaining_time, + robot_joint_pos_arm=self._robot_joint_pos_arm, + robot_joint_pos_hand=self._robot_joint_pos_hand, + robot_ee_pos_wrt_base=self._robot_ee_pos_wrt_base, + robot_ee_rotmat_wrt_base=self._robot_ee_rotmat_wrt_base, + robot_hand_wrench=self._robot_hand_wrench, + obj_com_pos_wrt_robot_ee=self._obj_com_pos_wrt_robot_ee, + obj_com_rotmat_wrt_robot_ee=self._obj_com_rotmat_wrt_robot_ee, + hole_entrance_pos_wrt_peg_ends=self._hole_entrance_pos_wrt_peg_ends, + hole_bottom_pos_wrt_peg_ends=self._hole_bottom_pos_wrt_peg_ends, + hole_rotmat_wrt_peg=self._hole_rotmat_wrt_peg, + ) + + ######################## + ### Helper Functions ### + ######################## + + def _update_intermediate_state(self): + ## Extract intermediate states + self._robot_ee_pos_wrt_base = self._tf_robot_ee.data.target_pos_source[:, 0, :] + self._robot_hand_wrench = ( + self._robot.root_physx_view.get_link_incoming_joint_force()[ + :, self._robot_hand_joint_indices + ] + ) + + ## Compute other intermediate states + target_pos_w, target_quat_w = self._target.get_world_poses() + ( + self._remaining_time, + self._robot_joint_pos_arm, + self._robot_joint_pos_hand, + self._robot_ee_rotmat_wrt_base, + self._obj_com_pos_wrt_robot_ee, + self._obj_com_rotmat_wrt_robot_ee, + self._hole_entrance_pos_wrt_peg_ends, + self._hole_bottom_pos_wrt_peg_ends, + self._hole_rotmat_wrt_peg, + self._rewards, + self._terminations, + self._truncations, + ) = _compute_intermediate_state( + current_action=self.action_manager.action, + previous_action=self.action_manager.prev_action, + episode_length_buf=self.episode_length_buf, + max_episode_length=self._max_episode_length, + robot_arm_joint_indices=self._robot_arm_joint_indices, + robot_hand_joint_indices=self._robot_hand_joint_indices, + joint_pos=self._robot.data.joint_pos, + soft_joint_pos_limits=self._robot.data.soft_joint_pos_limits, + robot_ee_pos_w=self._tf_robot_ee.data.target_pos_w[:, 0, :], + robot_ee_quat_w=self._tf_robot_ee.data.target_quat_w[:, 0, :], + robot_ee_quat_wrt_base=self._tf_robot_ee.data.target_quat_source[:, 0, :], + robot_arm_contact_net_forces=self._contacts_robot.data.net_forces_w, + robot_hand_obj_contact_force_matrix=self._contacts_robot_hand_obj.data.force_matrix_w, + obj_pos_w=self._object.data.root_pos_w, + obj_quat_w=self._object.data.root_quat_w, + obj_com_offset=self._obj_com_offset, + target_pos_w=target_pos_w, + target_quat_w=target_quat_w, + initial_obj_height_w=self._initial_obj_height_w, + peg_offset_pos_ends=self._peg_offset_pos_ends, + peg_rot_symmetry_n=self._peg_rot_symmetry_n, + hole_offset_pos_bottom=self._hole_offset_pos_bottom, + hole_offset_pos_entrance=self._hole_offset_pos_entrance, + ) + + +############################# +### TorchScript functions ### +############################# + + +@torch.jit.script +def _compute_intermediate_state( + *, + current_action: torch.Tensor, + previous_action: torch.Tensor, + episode_length_buf: torch.Tensor, + max_episode_length: int, + robot_arm_joint_indices: List[int], + robot_hand_joint_indices: List[int], + joint_pos: torch.Tensor, + soft_joint_pos_limits: torch.Tensor, + robot_ee_pos_w: torch.Tensor, + robot_ee_quat_w: torch.Tensor, + robot_ee_quat_wrt_base: torch.Tensor, + robot_arm_contact_net_forces: torch.Tensor, + robot_hand_obj_contact_force_matrix: torch.Tensor, + obj_pos_w: torch.Tensor, + obj_quat_w: torch.Tensor, + obj_com_offset: torch.Tensor, + target_pos_w: torch.Tensor, + target_quat_w: torch.Tensor, + initial_obj_height_w: torch.Tensor, + peg_offset_pos_ends: torch.Tensor, + peg_rot_symmetry_n: torch.Tensor, + hole_offset_pos_bottom: torch.Tensor, + hole_offset_pos_entrance: torch.Tensor, +) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + ## Intermediate states + # Time + remaining_time = 1 - (episode_length_buf / max_episode_length).unsqueeze(-1) + + # Robot joint positions + joint_pos_normalized = math_utils.scale_transform( + joint_pos, + soft_joint_pos_limits[:, :, 0], + soft_joint_pos_limits[:, :, 1], + ) + robot_joint_pos_arm, robot_joint_pos_hand = ( + joint_pos_normalized[:, robot_arm_joint_indices], + joint_pos_normalized[:, robot_hand_joint_indices], + ) + + # End-effector pose (position and '6D' rotation) + robot_ee_rotmat_wrt_base = math_utils.matrix_from_quat(robot_ee_quat_wrt_base) + + # Transformation | Object origin -> Object CoM + obj_com_pos_w, obj_com_quat_w = math_utils.combine_frame_transforms( + t01=obj_pos_w, + q01=obj_quat_w, + t12=obj_com_offset[:, :3], + q12=obj_com_offset[:, 3:], + ) + + # Transformation | Object origin -> Peg ends + _peg_end0_pos_w, _ = math_utils.combine_frame_transforms( + t01=obj_pos_w, + q01=obj_quat_w, + t12=peg_offset_pos_ends[:, 0], + ) + _peg_end1_pos_w, _ = math_utils.combine_frame_transforms( + t01=obj_pos_w, + q01=obj_quat_w, + t12=peg_offset_pos_ends[:, 1], + ) + peg_ends_pos_w = torch.stack([_peg_end0_pos_w, _peg_end1_pos_w], dim=1) + + # Transformation | Target origin -> Hole entrance + hole_entrance_pos_w, _ = math_utils.combine_frame_transforms( + t01=target_pos_w, + q01=target_quat_w, + t12=hole_offset_pos_entrance, + ) + + # Transformation | Target origin -> Hole bottom + hole_bottom_pos_w, _ = math_utils.combine_frame_transforms( + t01=target_pos_w, + q01=target_quat_w, + t12=hole_offset_pos_bottom, + ) + + # Transformation | End-effector -> Object CoM + obj_com_pos_wrt_robot_ee, obj_com_quat_wrt_robot_ee = ( + math_utils.subtract_frame_transforms( + t01=robot_ee_pos_w, + q01=robot_ee_quat_w, + t02=obj_com_pos_w, + q02=obj_com_quat_w, + ) + ) + obj_com_rotmat_wrt_robot_ee = math_utils.matrix_from_quat(obj_com_quat_wrt_robot_ee) + + # Transformation | Peg ends -> Hole entrance + _hole_entrance_pos_wrt_peg_end0, hole_quat_wrt_peg = ( + math_utils.subtract_frame_transforms( + t01=peg_ends_pos_w[:, 0], + q01=obj_quat_w, + t02=hole_entrance_pos_w, + q02=target_quat_w, + ) + ) + _hole_entrance_pos_wrt_peg_end1, _ = math_utils.subtract_frame_transforms( + t01=peg_ends_pos_w[:, 1], + q01=obj_quat_w, + t02=hole_entrance_pos_w, + ) + hole_entrance_pos_wrt_peg_ends = torch.stack( + [_hole_entrance_pos_wrt_peg_end0, _hole_entrance_pos_wrt_peg_end1], dim=1 + ) + hole_rotmat_wrt_peg = math_utils.matrix_from_quat(hole_quat_wrt_peg) + + # Transformation | Peg ends -> Hole bottom + _hole_bottom_pos_wrt_peg_end0, _ = math_utils.subtract_frame_transforms( + t01=peg_ends_pos_w[:, 0], + q01=obj_quat_w, + t02=hole_bottom_pos_w, + ) + _hole_bottom_pos_wrt_peg_end1, _ = math_utils.subtract_frame_transforms( + t01=peg_ends_pos_w[:, 1], + q01=obj_quat_w, + t02=hole_bottom_pos_w, + ) + hole_bottom_pos_wrt_peg_ends = torch.stack( + [_hole_bottom_pos_wrt_peg_end0, _hole_bottom_pos_wrt_peg_end1], dim=1 + ) + + ## Rewards + # Penalty: Action rate + WEIGHT_ACTION_RATE = -0.05 + penalty_action_rate = WEIGHT_ACTION_RATE * torch.sum( + torch.square(current_action - previous_action), dim=1 + ) + + # Penalty: Undesired robot arm contacts + WEIGHT_UNDERSIRED_ROBOT_ARM_CONTACTS = -0.1 + THRESHOLD_UNDERSIRED_ROBOT_ARM_CONTACTS = 10.0 + penalty_undersired_robot_arm_contacts = WEIGHT_UNDERSIRED_ROBOT_ARM_CONTACTS * ( + torch.max(torch.norm(robot_arm_contact_net_forces, dim=-1), dim=1)[0] + > THRESHOLD_UNDERSIRED_ROBOT_ARM_CONTACTS + ) + + # Reward: Distance | End-effector <--> Object + WEIGHT_DISTANCE_EE_TO_OBJ = 1.0 + TANH_STD_DISTANCE_EE_TO_OBJ = 0.25 + reward_distance_ee_to_obj = WEIGHT_DISTANCE_EE_TO_OBJ * ( + 1.0 + - torch.tanh( + torch.norm(obj_com_pos_wrt_robot_ee, dim=-1) / TANH_STD_DISTANCE_EE_TO_OBJ + ) + ) + + # Reward: Object grasped + WEIGHT_OBJ_GRASPED = 4.0 + THRESHOLD_OBJ_GRASPED = 5.0 + reward_obj_grasped = WEIGHT_OBJ_GRASPED * ( + torch.mean( + torch.max(torch.norm(robot_hand_obj_contact_force_matrix, dim=-1), dim=-1)[ + 0 + ], + dim=1, + ) + > THRESHOLD_OBJ_GRASPED + ) + + # Reward: Object lifted + WEIGHT_OBJ_LIFTED = 8.0 + HEIGHT_OFFSET_OBJ_LIFTED = 0.3 + HEIGHT_SPAN_OBJ_LIFTED = 0.25 + TAHN_STD_HEIGHT_OBJ_LIFTED = 0.05 + obj_target_height_offset = ( + torch.abs(obj_com_pos_w[:, 2] - initial_obj_height_w - HEIGHT_OFFSET_OBJ_LIFTED) + - HEIGHT_SPAN_OBJ_LIFTED + ).clamp(min=0.0) + reward_obj_lifted = WEIGHT_OBJ_LIFTED * ( + 1.0 - torch.tanh(obj_target_height_offset / TAHN_STD_HEIGHT_OBJ_LIFTED) + ) + + # Reward: Alignment | Peg -> Hole | Primary Z axis + WEIGHT_ALIGN_PEG_TO_HOLE_PRIMARY = 8.0 + TANH_STD_ALIGN_PEG_TO_HOLE_PRIMARY = 0.5 + _peg_to_hole_primary_axis_similarity = torch.abs(hole_rotmat_wrt_peg[:, 2, 2]) + reward_align_peg_to_hole_primary = WEIGHT_ALIGN_PEG_TO_HOLE_PRIMARY * ( + 1.0 + - torch.tanh( + (1.0 - _peg_to_hole_primary_axis_similarity) + / TANH_STD_ALIGN_PEG_TO_HOLE_PRIMARY + ) + ) + + # Reward: Alignment | Peg -> Hole | Secondary XY axes (affected by primary via power) + WEIGHT_ALIGN_PEG_TO_HOLE_SECONDARY = 4.0 + TANH_STD_ALIGN_PEG_TO_HOLE_SECONDARY = 0.2 + _peg_to_hole_yaw = torch.atan2( + hole_rotmat_wrt_peg[:, 0, 1], hole_rotmat_wrt_peg[:, 0, 0] + ) + _symmetry_step = 2 * torch.pi / peg_rot_symmetry_n + _peg_to_hole_yaw_symmetric_directional = _peg_to_hole_yaw % _symmetry_step + # Note: Lines above might result in NaN/inf when `peg_rot_symmetry_n=0` (infinite circular symmetry) + # However, the following `torch.where()` will handle this case + _peg_to_hole_yaw_symmetric_normalized = torch.where( + peg_rot_symmetry_n <= 0, + 0.0, + torch.min( + _peg_to_hole_yaw_symmetric_directional, + _symmetry_step - _peg_to_hole_yaw_symmetric_directional, + ) + / (_symmetry_step / 2.0), + ) + reward_align_peg_to_hole_secondary = WEIGHT_ALIGN_PEG_TO_HOLE_SECONDARY * ( + 1.0 + - torch.tanh( + _peg_to_hole_yaw_symmetric_normalized.pow( + _peg_to_hole_primary_axis_similarity + ) + / TANH_STD_ALIGN_PEG_TO_HOLE_SECONDARY + ) + ) + + # Reward: Distance | Peg -> Hole entrance + WEIGHT_DISTANCE_PEG_TO_HOLE_ENTRANCE = 16.0 + TANH_STD_DISTANCE_PEG_TO_HOLE_ENTRANCE = 0.05 + reward_distance_peg_to_hole_entrance = WEIGHT_DISTANCE_PEG_TO_HOLE_ENTRANCE * ( + 1.0 + - torch.tanh( + torch.min(torch.norm(hole_entrance_pos_wrt_peg_ends, dim=-1), dim=1)[0] + / TANH_STD_DISTANCE_PEG_TO_HOLE_ENTRANCE + ) + ) + + # Reward: Distance | Peg -> Hole bottom + WEIGHT_DISTANCE_PEG_TO_HOLE_BOTTOM = 128.0 + TANH_STD_DISTANCE_PEG_TO_HOLE_BOTTOM = 0.005 + reward_distance_peg_to_hole_bottom = WEIGHT_DISTANCE_PEG_TO_HOLE_BOTTOM * ( + 1.0 + - torch.tanh( + torch.min(torch.norm(hole_bottom_pos_wrt_peg_ends, dim=-1), dim=1)[0] + / TANH_STD_DISTANCE_PEG_TO_HOLE_BOTTOM + ) + ) + + # Total reward + rewards = torch.sum( + torch.stack( + [ + penalty_action_rate, + penalty_undersired_robot_arm_contacts, + reward_distance_ee_to_obj, + reward_obj_grasped, + reward_obj_lifted, + reward_align_peg_to_hole_primary, + reward_align_peg_to_hole_secondary, + reward_distance_peg_to_hole_entrance, + reward_distance_peg_to_hole_bottom, + ], + dim=-1, + ), + dim=-1, + ) + + ## Termination and truncation + truncations = episode_length_buf > (max_episode_length - 1) + terminations = torch.zeros_like(truncations) + + # print( + # f""" + # penalty | action_rate: {float(penalty_action_rate[0])} + # penalty | undersired_robot_arm_contacts: {float(penalty_undersired_robot_arm_contacts[0])} + # reward | distance_ee_to_obj: {float(reward_distance_ee_to_obj[0])} + # reward | obj_grasped: {float(reward_obj_grasped[0])} + # reward | obj_lifted: {float(reward_obj_lifted[0])} + # reward | align_peg_to_hole_primary: {float(reward_align_peg_to_hole_primary[0])} + # reward | align_peg_to_hole_secondary: {float(reward_align_peg_to_hole_secondary[0])} + # reward | distance_peg_to_hole_entrance: {float(reward_distance_peg_to_hole_entrance[0])} + # reward | distance_peg_to_hole_bottom: {float(reward_distance_peg_to_hole_bottom[0])} + # total: {float(rewards[0])} + # """ + # ) + + return ( + remaining_time, + robot_joint_pos_arm, + robot_joint_pos_hand, + robot_ee_rotmat_wrt_base, + obj_com_pos_wrt_robot_ee, + obj_com_rotmat_wrt_robot_ee, + hole_entrance_pos_wrt_peg_ends, + hole_bottom_pos_wrt_peg_ends, + hole_rotmat_wrt_peg, + rewards, + terminations, + truncations, + ) + + +@torch.jit.script +def _construct_observations( + *, + remaining_time: torch.Tensor, + robot_joint_pos_arm: torch.Tensor, + robot_joint_pos_hand: torch.Tensor, + robot_ee_pos_wrt_base: torch.Tensor, + robot_ee_rotmat_wrt_base: torch.Tensor, + robot_hand_wrench: torch.Tensor, + obj_com_pos_wrt_robot_ee: torch.Tensor, + obj_com_rotmat_wrt_robot_ee: torch.Tensor, + hole_entrance_pos_wrt_peg_ends: torch.Tensor, + hole_bottom_pos_wrt_peg_ends: torch.Tensor, + hole_rotmat_wrt_peg: torch.Tensor, +) -> Dict[str, torch.Tensor]: + """ + Note: The `robot_hand_wrench` is considered as state (robot without force-torque sensors) + """ + + num_envs = remaining_time.size(0) + + # Robot joint positions + robot_joint_pos_hand_mean = robot_joint_pos_hand.mean(dim=-1, keepdim=True) + + # End-effector pose (position and '6D' rotation) + robot_ee_rot6d = math_utils.rotmat_to_rot6d(robot_ee_rotmat_wrt_base) + + # Wrench + robot_hand_wrench_full = robot_hand_wrench.view(num_envs, -1) + robot_hand_wrench_mean = robot_hand_wrench.mean(dim=1) + + # Transformation | End-effector -> Object CoM + obj_com_rot6d_wrt_robot_ee = math_utils.rotmat_to_rot6d(obj_com_rotmat_wrt_robot_ee) + + # Transformation | Object -> Target + hole_6d_wrt_peg = math_utils.rotmat_to_rot6d(hole_rotmat_wrt_peg) + + return { + "state": torch.cat( + [ + obj_com_pos_wrt_robot_ee, + obj_com_rot6d_wrt_robot_ee, + hole_entrance_pos_wrt_peg_ends.view(num_envs, -1), + hole_bottom_pos_wrt_peg_ends.view(num_envs, -1), + hole_6d_wrt_peg, + robot_hand_wrench_mean, + ], + dim=-1, + ), + "state_dyn": torch.cat([robot_hand_wrench_full], dim=-1), + "proprio": torch.cat( + [ + remaining_time, + robot_ee_pos_wrt_base, + robot_ee_rot6d, + robot_joint_pos_hand_mean, + ], + dim=-1, + ), + "proprio_dyn": torch.cat([robot_joint_pos_arm, robot_joint_pos_hand], dim=-1), + } diff --git a/space_robotics_bench/tasks/peg_in_hole/task_visual.py b/space_robotics_bench/tasks/peg_in_hole/task_visual.py new file mode 100644 index 0000000..9b765bc --- /dev/null +++ b/space_robotics_bench/tasks/peg_in_hole/task_visual.py @@ -0,0 +1,32 @@ +from typing import Dict + +import torch +from omni.isaac.lab.utils import configclass + +from space_robotics_bench.envs import ( + VisualManipulationEnvExt, + VisualManipulationEnvExtCfg, +) + +from .task import Task, TaskCfg + + +@configclass +class VisualTaskCfg(TaskCfg, VisualManipulationEnvExtCfg): + def __post_init__(self): + TaskCfg.__post_init__(self) + VisualManipulationEnvExtCfg.__post_init__(self) + + +class VisualTask(Task, VisualManipulationEnvExt): + cfg: VisualTaskCfg + + def __init__(self, cfg: VisualTaskCfg, **kwargs): + Task.__init__(self, cfg, **kwargs) + VisualManipulationEnvExt.__init__(self, cfg, **kwargs) + + def _get_observations(self) -> Dict[str, torch.Tensor]: + return { + **Task._get_observations(self), + **VisualManipulationEnvExt._get_observations(self), + } diff --git a/space_robotics_bench/tasks/peg_in_hole_multi/__init__.py b/space_robotics_bench/tasks/peg_in_hole_multi/__init__.py new file mode 100644 index 0000000..b380c63 --- /dev/null +++ b/space_robotics_bench/tasks/peg_in_hole_multi/__init__.py @@ -0,0 +1,17 @@ +from space_robotics_bench.utils.registry import register_tasks + +from .task import Task, TaskCfg +from .task_visual import VisualTask, VisualTaskCfg + +BASE_TASK_NAME = __name__.split(".")[-1] +register_tasks( + { + BASE_TASK_NAME: {}, + f"{BASE_TASK_NAME}_visual": { + "entry_point": VisualTask, + "task_cfg": VisualTaskCfg, + }, + }, + default_entry_point=Task, + default_task_cfg=TaskCfg, +) diff --git a/space_robotics_bench/tasks/peg_in_hole_multi/task.py b/space_robotics_bench/tasks/peg_in_hole_multi/task.py new file mode 100644 index 0000000..9015fd3 --- /dev/null +++ b/space_robotics_bench/tasks/peg_in_hole_multi/task.py @@ -0,0 +1,708 @@ +from dataclasses import MISSING as DELAYED_CFG +from typing import Dict, List, Optional, Sequence, Tuple + +import torch +from omni.isaac.core.prims.xform_prim_view import XFormPrimView +from omni.isaac.lab.managers import EventTermCfg, SceneEntityCfg +from omni.isaac.lab.sensors import ContactSensor, ContactSensorCfg +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.utils.math as math_utils +import space_robotics_bench.utils.sampling as spacing_utils +from space_robotics_bench.core.assets import RigidObject, RigidObjectCfg +from space_robotics_bench.envs import ( + BaseManipulationEnv, + BaseManipulationEnvCfg, + BaseManipulationEnvEventCfg, + mdp, +) + +from ..peg_in_hole.task import peg_and_hole_cfg + +############## +### Config ### +############## + + +@configclass +class TaskCfg(BaseManipulationEnvCfg): + ## Environment + episode_length_s: float = DELAYED_CFG + num_problems_per_env: int = 6 + problem_spacing = 0.15 + + ## Task + is_finite_horizon: bool = False + + ## Events + @configclass + class EventCfg(BaseManipulationEnvEventCfg): + pass + + events = EventCfg() + + def __post_init__(self): + super().__post_init__() + + ## Environment + self.episode_length_s = self.num_problems_per_env * 10.0 + + ## Scene + (num_rows, num_cols), (grid_spacing_pos, grid_spacing_rot) = ( + spacing_utils.compute_grid_spacing( + num_instances=self.num_problems_per_env, + spacing=self.problem_spacing, + global_pos_offset=(0.5, 0.0, 0.1), + ) + ) + self.problem_cfgs = [ + peg_and_hole_cfg( + self.env_cfg, + prim_path_peg=f"{{ENV_REGEX_NS}}/peg{i}", + prim_path_hole=f"{{ENV_REGEX_NS}}/hole{i}", + asset_cfg_peg=SceneEntityCfg(f"object{i}"), + num_assets=self.scene.num_envs, + init_state=RigidObjectCfg.InitialStateCfg( + pos=grid_spacing_pos[i], + rot=grid_spacing_rot[i], + ), + procgen_seed_offset=i * self.scene.num_envs, + spawn_kwargs_peg={ + "activate_contact_sensors": True, + }, + ) + for i in range(self.num_problems_per_env) + ] + for i, problem_cfg in enumerate(self.problem_cfgs): + setattr( + self.scene, + f"object{i}", + problem_cfg[0].asset_cfg.replace( + init_state=RigidObjectCfg.InitialStateCfg( + pos=(0.5, 0.0, 0.13), + ) + ), + ) + setattr( + self.scene, + f"target{i}", + problem_cfg[1].asset_cfg, + ) + + ## Sensors + self.scene.contacts_robot_hand_obj = ContactSensorCfg( + prim_path=f"{self.scene.robot.prim_path}/{self.robot_cfg.regex_links_hand}", + update_period=0.0, + # Note: This causes error 'Filter pattern did not match the correct number of entries' + # However, it seems to function properly anyway... + filter_prim_paths_expr=[ + asset.prim_path + for asset in [ + getattr(self.scene, f"object{i}") + for i in range(self.num_problems_per_env) + ] + ], + ) + + ## Events + self.events.reset_rand_object_state_multi = EventTermCfg( + func=mdp.reset_root_state_uniform_poisson_disk_2d, + mode="reset", + params={ + "asset_cfgs": [ + SceneEntityCfg(f"object{i}") + for i in range(self.num_problems_per_env) + ], + "pose_range": { + "x": ( + -0.5 * (num_rows - 0.5) * self.problem_spacing, + 0.5 * (num_rows - 0.5) * self.problem_spacing, + ), + "y": ( + -0.5 * (num_cols - 0.5) * self.problem_spacing, + 0.5 * (num_cols - 0.5) * self.problem_spacing, + ), + "roll": (torch.pi / 2, torch.pi / 2), + "pitch": (-torch.pi, torch.pi), + "yaw": (-torch.pi, torch.pi), + }, + "velocity_range": {}, + "radius": 0.1, + }, + ) + + +############ +### Task ### +############ + + +class Task(BaseManipulationEnv): + cfg: TaskCfg + + def __init__(self, cfg: TaskCfg, **kwargs): + super().__init__(cfg, **kwargs) + + # Get handles to scene assets + self._contacts_robot_hand_obj: ContactSensor = self.scene[ + "contacts_robot_hand_obj" + ] + self._objects: List[RigidObject] = [ + self.scene[f"object{i}"] for i in range(self.cfg.num_problems_per_env) + ] + self._targets: List[XFormPrimView] = [ + self.scene[f"target{i}"] for i in range(self.cfg.num_problems_per_env) + ] + + ## Pre-compute metrics used in hot loops + self._robot_arm_joint_indices, _ = self._robot.find_joints( + self.cfg.robot_cfg.regex_joints_arm + ) + self._robot_hand_joint_indices, _ = self._robot.find_joints( + self.cfg.robot_cfg.regex_joints_hand + ) + self._max_episode_length = self.max_episode_length + self._obj_com_offset = torch.stack( + [ + self._objects[i].data._root_physx_view.get_coms().to(self.device) + for i in range(self.cfg.num_problems_per_env) + ], + dim=1, + ) + + ## Initialize buffers + self._initial_obj_height_w = torch.zeros( + (self.num_envs, self.cfg.num_problems_per_env), + dtype=torch.float32, + device=self.device, + ) + self._peg_offset_pos_ends = torch.tensor( + [ + self.cfg.problem_cfgs[i][0].offset_pos_ends + for i in range(self.cfg.num_problems_per_env) + ], + dtype=torch.float32, + device=self.device, + ).repeat(self.num_envs, 1, 1, 1) + + self._peg_rot_symmetry_n = torch.tensor( + [ + self.cfg.problem_cfgs[i][0].rot_symmetry_n + for i in range(self.cfg.num_problems_per_env) + ], + dtype=torch.int32, + device=self.device, + ).repeat(self.num_envs, 1) + self._hole_offset_pos_bottom = torch.tensor( + [ + self.cfg.problem_cfgs[i][1].offset_pos_bottom + for i in range(self.cfg.num_problems_per_env) + ], + dtype=torch.float32, + device=self.device, + ).repeat(self.num_envs, 1, 1) + self._hole_offset_pos_entrance = torch.tensor( + [ + self.cfg.problem_cfgs[i][1].offset_pos_entrance + for i in range(self.cfg.num_problems_per_env) + ], + dtype=torch.float32, + device=self.device, + ).repeat(self.num_envs, 1, 1) + + ## Initialize the intermediate state + self._update_intermediate_state() + + def _reset_idx(self, env_ids: Sequence[int]): + super()._reset_idx(env_ids) + + # Update the initial height of the objects + self._initial_obj_height_w[env_ids] = torch.stack( + [ + self._objects[i].data.root_pos_w[env_ids, 2] + for i in range(self.cfg.num_problems_per_env) + ], + dim=1, + ) + + def _get_dones(self) -> Tuple[torch.Tensor, torch.Tensor]: + # Note: This assumes that `_get_dones()` is called before `_get_rewards()` and `_get_observations()` in `step()` + self._update_intermediate_state() + + if not self.cfg.enable_truncation: + self._truncations = torch.zeros_like(self._truncations) + + return self._terminations, self._truncations + + def _get_rewards(self) -> torch.Tensor: + return self._rewards + + def _get_observations(self) -> Dict[str, torch.Tensor]: + return _construct_observations( + remaining_time=self._remaining_time, + robot_joint_pos_arm=self._robot_joint_pos_arm, + robot_joint_pos_hand=self._robot_joint_pos_hand, + robot_ee_pos_wrt_base=self._robot_ee_pos_wrt_base, + robot_ee_rotmat_wrt_base=self._robot_ee_rotmat_wrt_base, + robot_hand_wrench=self._robot_hand_wrench, + obj_com_pos_wrt_robot_ee=self._obj_com_pos_wrt_robot_ee, + obj_com_rotmat_wrt_robot_ee=self._obj_com_rotmat_wrt_robot_ee, + hole_entrance_pos_wrt_peg_ends=self._hole_entrance_pos_wrt_peg_ends, + hole_bottom_pos_wrt_peg_ends=self._hole_bottom_pos_wrt_peg_ends, + hole_rotmat_wrt_peg=self._hole_rotmat_wrt_peg, + ) + + ######################## + ### Helper Functions ### + ######################## + + def _update_intermediate_state(self): + ## Extract intermediate states + self._robot_ee_pos_wrt_base = self._tf_robot_ee.data.target_pos_source[:, 0, :] + self._robot_hand_wrench = ( + self._robot.root_physx_view.get_link_incoming_joint_force()[ + :, self._robot_hand_joint_indices + ] + ) + + ## Compute other intermediate states + target_pos_w = torch.stack( + [ + self._targets[i].get_world_poses()[0] + for i in range(self.cfg.num_problems_per_env) + ], + dim=1, + ) + target_quat_w = torch.stack( + [ + self._targets[i].get_world_poses()[1] + for i in range(self.cfg.num_problems_per_env) + ], + dim=1, + ) + ( + self._remaining_time, + self._robot_joint_pos_arm, + self._robot_joint_pos_hand, + self._robot_ee_rotmat_wrt_base, + self._obj_com_pos_wrt_robot_ee, + self._obj_com_rotmat_wrt_robot_ee, + self._hole_entrance_pos_wrt_peg_ends, + self._hole_bottom_pos_wrt_peg_ends, + self._hole_rotmat_wrt_peg, + self._rewards, + self._terminations, + self._truncations, + ) = _compute_intermediate_state( + current_action=self.action_manager.action, + previous_action=self.action_manager.prev_action, + episode_length_buf=self.episode_length_buf, + max_episode_length=self._max_episode_length, + robot_arm_joint_indices=self._robot_arm_joint_indices, + robot_hand_joint_indices=self._robot_hand_joint_indices, + joint_pos=self._robot.data.joint_pos, + soft_joint_pos_limits=self._robot.data.soft_joint_pos_limits, + robot_ee_pos_w=self._tf_robot_ee.data.target_pos_w[:, 0, :], + robot_ee_quat_w=self._tf_robot_ee.data.target_quat_w[:, 0, :], + robot_ee_quat_wrt_base=self._tf_robot_ee.data.target_quat_source[:, 0, :], + robot_arm_contact_net_forces=self._contacts_robot.data.net_forces_w, + robot_hand_obj_contact_force_matrix=self._contacts_robot_hand_obj.data.force_matrix_w, + obj_pos_w=torch.stack( + [ + self._objects[i].data.root_pos_w + for i in range(self.cfg.num_problems_per_env) + ], + dim=1, + ), + obj_quat_w=torch.stack( + [ + self._objects[i].data.root_quat_w + for i in range(self.cfg.num_problems_per_env) + ], + dim=1, + ), + obj_com_offset=self._obj_com_offset, + target_pos_w=target_pos_w, + target_quat_w=target_quat_w, + initial_obj_height_w=self._initial_obj_height_w, + peg_offset_pos_ends=self._peg_offset_pos_ends, + peg_rot_symmetry_n=self._peg_rot_symmetry_n, + hole_offset_pos_bottom=self._hole_offset_pos_bottom, + hole_offset_pos_entrance=self._hole_offset_pos_entrance, + ) + + +############################# +### TorchScript functions ### +############################# + + +@torch.jit.script +def _compute_intermediate_state( + *, + current_action: torch.Tensor, + previous_action: torch.Tensor, + episode_length_buf: torch.Tensor, + max_episode_length: int, + robot_arm_joint_indices: List[int], + robot_hand_joint_indices: List[int], + joint_pos: torch.Tensor, + soft_joint_pos_limits: torch.Tensor, + robot_ee_pos_w: torch.Tensor, + robot_ee_quat_w: torch.Tensor, + robot_ee_quat_wrt_base: torch.Tensor, + robot_arm_contact_net_forces: torch.Tensor, + robot_hand_obj_contact_force_matrix: torch.Tensor, + obj_pos_w: torch.Tensor, + obj_quat_w: torch.Tensor, + obj_com_offset: torch.Tensor, + target_pos_w: torch.Tensor, + target_quat_w: torch.Tensor, + initial_obj_height_w: torch.Tensor, + peg_offset_pos_ends: torch.Tensor, + peg_rot_symmetry_n: torch.Tensor, + hole_offset_pos_bottom: torch.Tensor, + hole_offset_pos_entrance: torch.Tensor, +) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + ## Intermediate states + # Time + remaining_time = 1 - (episode_length_buf / max_episode_length).unsqueeze(-1) + + # Robot joint positions + joint_pos_normalized = math_utils.scale_transform( + joint_pos, + soft_joint_pos_limits[:, :, 0], + soft_joint_pos_limits[:, :, 1], + ) + robot_joint_pos_arm, robot_joint_pos_hand = ( + joint_pos_normalized[:, robot_arm_joint_indices], + joint_pos_normalized[:, robot_hand_joint_indices], + ) + + # End-effector pose (position and '6D' rotation) + robot_ee_rotmat_wrt_base = math_utils.matrix_from_quat(robot_ee_quat_wrt_base) + + # Transformation | Object origin -> Object CoM + obj_com_pos_w, obj_com_quat_w = math_utils.combine_frame_transforms( + t01=obj_pos_w, + q01=obj_quat_w, + t12=obj_com_offset[:, :, :3], + q12=obj_com_offset[:, :, 3:], + ) + + # Transformation | Object origin -> Peg ends + _peg_end0_pos_w, _ = math_utils.combine_frame_transforms( + t01=obj_pos_w, + q01=obj_quat_w, + t12=peg_offset_pos_ends[:, :, 0], + ) + _peg_end1_pos_w, _ = math_utils.combine_frame_transforms( + t01=obj_pos_w, + q01=obj_quat_w, + t12=peg_offset_pos_ends[:, :, 1], + ) + peg_ends_pos_w = torch.stack([_peg_end0_pos_w, _peg_end1_pos_w], dim=1) + + # Transformation | Target origin -> Hole entrance + hole_entrance_pos_w, _ = math_utils.combine_frame_transforms( + t01=target_pos_w, + q01=target_quat_w, + t12=hole_offset_pos_entrance, + ) + + # Transformation | Target origin -> Hole bottom + hole_bottom_pos_w, _ = math_utils.combine_frame_transforms( + t01=target_pos_w, + q01=target_quat_w, + t12=hole_offset_pos_bottom, + ) + + # Transformation | End-effector -> Object CoM + obj_com_pos_wrt_robot_ee, obj_com_quat_wrt_robot_ee = ( + math_utils.subtract_frame_transforms( + t01=robot_ee_pos_w.unsqueeze(1).repeat(1, obj_pos_w.shape[1], 1), + q01=robot_ee_quat_w.unsqueeze(1).repeat(1, obj_pos_w.shape[1], 1), + t02=obj_com_pos_w, + q02=obj_com_quat_w, + ) + ) + obj_com_rotmat_wrt_robot_ee = math_utils.matrix_from_quat(obj_com_quat_wrt_robot_ee) + + # Transformation | Peg ends -> Hole entrance + _hole_entrance_pos_wrt_peg_end0, hole_quat_wrt_peg = ( + math_utils.subtract_frame_transforms( + t01=peg_ends_pos_w[:, 0], + q01=obj_quat_w, + t02=hole_entrance_pos_w, + q02=target_quat_w, + ) + ) + _hole_entrance_pos_wrt_peg_end1, _ = math_utils.subtract_frame_transforms( + t01=peg_ends_pos_w[:, 1], + q01=obj_quat_w, + t02=hole_entrance_pos_w, + ) + hole_entrance_pos_wrt_peg_ends = torch.stack( + [_hole_entrance_pos_wrt_peg_end0, _hole_entrance_pos_wrt_peg_end1], dim=1 + ) + hole_rotmat_wrt_peg = math_utils.matrix_from_quat(hole_quat_wrt_peg) + + # Transformation | Peg ends -> Hole bottom + _hole_bottom_pos_wrt_peg_end0, _ = math_utils.subtract_frame_transforms( + t01=peg_ends_pos_w[:, 0], + q01=obj_quat_w, + t02=hole_bottom_pos_w, + ) + _hole_bottom_pos_wrt_peg_end1, _ = math_utils.subtract_frame_transforms( + t01=peg_ends_pos_w[:, 1], + q01=obj_quat_w, + t02=hole_bottom_pos_w, + ) + hole_bottom_pos_wrt_peg_ends = torch.stack( + [_hole_bottom_pos_wrt_peg_end0, _hole_bottom_pos_wrt_peg_end1], dim=1 + ) + + ## Rewards + # Penalty: Action rate + WEIGHT_ACTION_RATE = -0.05 + penalty_action_rate = WEIGHT_ACTION_RATE * torch.sum( + torch.square(current_action - previous_action), dim=1 + ) + + # Penalty: Undesired robot arm contacts + WEIGHT_UNDERSIRED_ROBOT_ARM_CONTACTS = -0.1 + THRESHOLD_UNDERSIRED_ROBOT_ARM_CONTACTS = 10.0 + penalty_undersired_robot_arm_contacts = WEIGHT_UNDERSIRED_ROBOT_ARM_CONTACTS * ( + torch.max(torch.norm(robot_arm_contact_net_forces, dim=-1), dim=1)[0] + > THRESHOLD_UNDERSIRED_ROBOT_ARM_CONTACTS + ) + + # Reward: Distance | End-effector <--> Object + WEIGHT_DISTANCE_EE_TO_OBJ = 1.0 + TANH_STD_DISTANCE_EE_TO_OBJ = 0.25 + reward_distance_ee_to_obj = WEIGHT_DISTANCE_EE_TO_OBJ * ( + 1.0 + - torch.tanh( + torch.norm(obj_com_pos_wrt_robot_ee, dim=-1) / TANH_STD_DISTANCE_EE_TO_OBJ + ) + ).sum(dim=-1) + + # Reward: Object grasped + WEIGHT_OBJ_GRASPED = 4.0 + THRESHOLD_OBJ_GRASPED = 5.0 + reward_obj_grasped = WEIGHT_OBJ_GRASPED * ( + torch.mean( + torch.max(torch.norm(robot_hand_obj_contact_force_matrix, dim=-1), dim=-1)[ + 0 + ], + dim=1, + ) + > THRESHOLD_OBJ_GRASPED + ) + + # Reward: Object lifted + WEIGHT_OBJ_LIFTED = 8.0 + HEIGHT_OFFSET_OBJ_LIFTED = 0.3 + HEIGHT_SPAN_OBJ_LIFTED = 0.25 + TAHN_STD_HEIGHT_OBJ_LIFTED = 0.05 + obj_target_height_offset = ( + torch.abs( + obj_com_pos_w[:, :, 2] - initial_obj_height_w - HEIGHT_OFFSET_OBJ_LIFTED + ) + - HEIGHT_SPAN_OBJ_LIFTED + ).clamp(min=0.0) + reward_obj_lifted = WEIGHT_OBJ_LIFTED * ( + 1.0 - torch.tanh(obj_target_height_offset / TAHN_STD_HEIGHT_OBJ_LIFTED) + ).sum(dim=-1) + + # Reward: Alignment | Peg -> Hole | Primary Z axis + WEIGHT_ALIGN_PEG_TO_HOLE_PRIMARY = 8.0 + TANH_STD_ALIGN_PEG_TO_HOLE_PRIMARY = 0.5 + _peg_to_hole_primary_axis_similarity = torch.abs(hole_rotmat_wrt_peg[:, :, 2, 2]) + reward_align_peg_to_hole_primary = WEIGHT_ALIGN_PEG_TO_HOLE_PRIMARY * ( + 1.0 + - torch.tanh( + (1.0 - _peg_to_hole_primary_axis_similarity) + / TANH_STD_ALIGN_PEG_TO_HOLE_PRIMARY + ) + ).sum(dim=-1) + + # Reward: Alignment | Peg -> Hole | Secondary XY axes (affected by primary via power) + WEIGHT_ALIGN_PEG_TO_HOLE_SECONDARY = 4.0 + TANH_STD_ALIGN_PEG_TO_HOLE_SECONDARY = 0.2 + _peg_to_hole_yaw = torch.atan2( + hole_rotmat_wrt_peg[:, :, 0, 1], hole_rotmat_wrt_peg[:, :, 0, 0] + ) + _symmetry_step = 2 * torch.pi / peg_rot_symmetry_n + _peg_to_hole_yaw_symmetric_directional = _peg_to_hole_yaw % _symmetry_step + # Note: Lines above might result in NaN/inf when `peg_rot_symmetry_n=0` (infinite circular symmetry) + # However, the following `torch.where()` will handle this case + _peg_to_hole_yaw_symmetric_normalized = torch.where( + peg_rot_symmetry_n <= 0, + 0.0, + torch.min( + _peg_to_hole_yaw_symmetric_directional, + _symmetry_step - _peg_to_hole_yaw_symmetric_directional, + ) + / (_symmetry_step / 2.0), + ) + reward_align_peg_to_hole_secondary = WEIGHT_ALIGN_PEG_TO_HOLE_SECONDARY * ( + 1.0 + - torch.tanh( + _peg_to_hole_yaw_symmetric_normalized.pow( + _peg_to_hole_primary_axis_similarity + ) + / TANH_STD_ALIGN_PEG_TO_HOLE_SECONDARY + ) + ).sum(dim=-1) + + # Reward: Distance | Peg -> Hole entrance + WEIGHT_DISTANCE_PEG_TO_HOLE_ENTRANCE = 16.0 + TANH_STD_DISTANCE_PEG_TO_HOLE_ENTRANCE = 0.025 + reward_distance_peg_to_hole_entrance = WEIGHT_DISTANCE_PEG_TO_HOLE_ENTRANCE * ( + 1.0 + - torch.tanh( + torch.min(torch.norm(hole_entrance_pos_wrt_peg_ends, dim=-1), dim=1)[0] + / TANH_STD_DISTANCE_PEG_TO_HOLE_ENTRANCE + ) + ).sum(dim=-1) + + # Reward: Distance | Peg -> Hole bottom + WEIGHT_DISTANCE_PEG_TO_HOLE_BOTTOM = 128.0 + TANH_STD_DISTANCE_PEG_TO_HOLE_BOTTOM = 0.005 + reward_distance_peg_to_hole_bottom = WEIGHT_DISTANCE_PEG_TO_HOLE_BOTTOM * ( + 1.0 + - torch.tanh( + torch.min(torch.norm(hole_bottom_pos_wrt_peg_ends, dim=-1), dim=1)[0] + / TANH_STD_DISTANCE_PEG_TO_HOLE_BOTTOM + ) + ).sum(dim=-1) + + # Total reward + rewards = torch.sum( + torch.stack( + [ + penalty_action_rate, + penalty_undersired_robot_arm_contacts, + reward_distance_ee_to_obj, + reward_obj_grasped, + reward_obj_lifted, + reward_align_peg_to_hole_primary, + reward_align_peg_to_hole_secondary, + reward_distance_peg_to_hole_entrance, + reward_distance_peg_to_hole_bottom, + ], + dim=-1, + ), + dim=-1, + ) + + ## Termination and truncation + truncations = episode_length_buf > (max_episode_length - 1) + terminations = torch.zeros_like(truncations) + + # print( + # f""" + # penalty | action_rate: {float(penalty_action_rate[0])} + # penalty | undersired_robot_arm_contacts: {float(penalty_undersired_robot_arm_contacts[0])} + # reward | distance_ee_to_obj: {float(reward_distance_ee_to_obj[0])} + # reward | obj_grasped: {float(reward_obj_grasped[0])} + # reward | obj_lifted: {float(reward_obj_lifted[0])} + # reward | align_peg_to_hole_primary: {float(reward_align_peg_to_hole_primary[0])} + # reward | align_peg_to_hole_secondary: {float(reward_align_peg_to_hole_secondary[0])} + # reward | distance_peg_to_hole_entrance: {float(reward_distance_peg_to_hole_entrance[0])} + # reward | distance_peg_to_hole_bottom: {float(reward_distance_peg_to_hole_bottom[0])} + # total: {float(rewards[0])} + # """ + # ) + + return ( + remaining_time, + robot_joint_pos_arm, + robot_joint_pos_hand, + robot_ee_rotmat_wrt_base, + obj_com_pos_wrt_robot_ee, + obj_com_rotmat_wrt_robot_ee, + hole_entrance_pos_wrt_peg_ends, + hole_bottom_pos_wrt_peg_ends, + hole_rotmat_wrt_peg, + rewards, + terminations, + truncations, + ) + + +@torch.jit.script +def _construct_observations( + *, + remaining_time: torch.Tensor, + robot_joint_pos_arm: torch.Tensor, + robot_joint_pos_hand: torch.Tensor, + robot_ee_pos_wrt_base: torch.Tensor, + robot_ee_rotmat_wrt_base: torch.Tensor, + robot_hand_wrench: torch.Tensor, + obj_com_pos_wrt_robot_ee: torch.Tensor, + obj_com_rotmat_wrt_robot_ee: torch.Tensor, + hole_entrance_pos_wrt_peg_ends: torch.Tensor, + hole_bottom_pos_wrt_peg_ends: torch.Tensor, + hole_rotmat_wrt_peg: torch.Tensor, +) -> Dict[str, torch.Tensor]: + """ + Note: The `robot_hand_wrench` is considered as state (robot without force-torque sensors) + """ + + num_envs = remaining_time.size(0) + + # Robot joint positions + robot_joint_pos_hand_mean = robot_joint_pos_hand.mean(dim=-1, keepdim=True) + + # End-effector pose (position and '6D' rotation) + robot_ee_rot6d = math_utils.rotmat_to_rot6d(robot_ee_rotmat_wrt_base) + + # Wrench + robot_hand_wrench_full = robot_hand_wrench.view(num_envs, -1) + robot_hand_wrench_mean = robot_hand_wrench.mean(dim=1) + + # Transformation | End-effector -> Object CoM + obj_com_rot6d_wrt_robot_ee = math_utils.rotmat_to_rot6d(obj_com_rotmat_wrt_robot_ee) + + # Transformation | Object -> Target + hole_6d_wrt_peg = math_utils.rotmat_to_rot6d(hole_rotmat_wrt_peg) + + return { + "state": torch.cat( + [ + obj_com_pos_wrt_robot_ee.view(num_envs, -1), + obj_com_rot6d_wrt_robot_ee.view(num_envs, -1), + hole_entrance_pos_wrt_peg_ends.view(num_envs, -1), + hole_bottom_pos_wrt_peg_ends.view(num_envs, -1), + hole_6d_wrt_peg.view(num_envs, -1), + robot_hand_wrench_mean.view(num_envs, -1), + ], + dim=-1, + ), + "state_dyn": torch.cat([robot_hand_wrench_full], dim=-1), + "proprio": torch.cat( + [ + remaining_time, + robot_ee_pos_wrt_base, + robot_ee_rot6d, + robot_joint_pos_hand_mean, + ], + dim=-1, + ), + "proprio_dyn": torch.cat([robot_joint_pos_arm, robot_joint_pos_hand], dim=-1), + } diff --git a/space_robotics_bench/tasks/peg_in_hole_multi/task_visual.py b/space_robotics_bench/tasks/peg_in_hole_multi/task_visual.py new file mode 100644 index 0000000..9b765bc --- /dev/null +++ b/space_robotics_bench/tasks/peg_in_hole_multi/task_visual.py @@ -0,0 +1,32 @@ +from typing import Dict + +import torch +from omni.isaac.lab.utils import configclass + +from space_robotics_bench.envs import ( + VisualManipulationEnvExt, + VisualManipulationEnvExtCfg, +) + +from .task import Task, TaskCfg + + +@configclass +class VisualTaskCfg(TaskCfg, VisualManipulationEnvExtCfg): + def __post_init__(self): + TaskCfg.__post_init__(self) + VisualManipulationEnvExtCfg.__post_init__(self) + + +class VisualTask(Task, VisualManipulationEnvExt): + cfg: VisualTaskCfg + + def __init__(self, cfg: VisualTaskCfg, **kwargs): + Task.__init__(self, cfg, **kwargs) + VisualManipulationEnvExt.__init__(self, cfg, **kwargs) + + def _get_observations(self) -> Dict[str, torch.Tensor]: + return { + **Task._get_observations(self), + **VisualManipulationEnvExt._get_observations(self), + } diff --git a/space_robotics_bench/tasks/sample_collection/__init__.py b/space_robotics_bench/tasks/sample_collection/__init__.py new file mode 100644 index 0000000..b380c63 --- /dev/null +++ b/space_robotics_bench/tasks/sample_collection/__init__.py @@ -0,0 +1,17 @@ +from space_robotics_bench.utils.registry import register_tasks + +from .task import Task, TaskCfg +from .task_visual import VisualTask, VisualTaskCfg + +BASE_TASK_NAME = __name__.split(".")[-1] +register_tasks( + { + BASE_TASK_NAME: {}, + f"{BASE_TASK_NAME}_visual": { + "entry_point": VisualTask, + "task_cfg": VisualTaskCfg, + }, + }, + default_entry_point=Task, + default_task_cfg=TaskCfg, +) diff --git a/space_robotics_bench/tasks/sample_collection/task.py b/space_robotics_bench/tasks/sample_collection/task.py new file mode 100644 index 0000000..c231e38 --- /dev/null +++ b/space_robotics_bench/tasks/sample_collection/task.py @@ -0,0 +1,552 @@ +from typing import Any, Dict, List, Optional, Sequence, Tuple + +import torch +from omni.isaac.lab.managers import EventTermCfg, SceneEntityCfg +from omni.isaac.lab.sensors import ContactSensor, ContactSensorCfg +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.envs as env_utils +import space_robotics_bench.core.sim as sim_utils +import space_robotics_bench.utils.math as math_utils +from space_robotics_bench import assets +from space_robotics_bench.core.assets import AssetCfg, RigidObject, RigidObjectCfg +from space_robotics_bench.core.markers import ( + VisualizationMarkers, + VisualizationMarkersCfg, +) +from space_robotics_bench.core.sim.spawners import SphereCfg +from space_robotics_bench.envs import ( + BaseManipulationEnv, + BaseManipulationEnvCfg, + BaseManipulationEnvEventCfg, + mdp, +) + +############## +### Config ### +############## + + +class SampleCfg(AssetCfg): + class Config: + arbitrary_types_allowed = True # Due to EventTermCfg + + ## Model + asset_cfg: RigidObjectCfg + + ## Randomization + state_randomizer: EventTermCfg + + +def sample_cfg( + env_cfg: env_utils.EnvironmentConfig, + *, + prim_path: str = "{ENV_REGEX_NS}/sample", + asset_cfg: SceneEntityCfg = SceneEntityCfg("object"), + num_assets: int = 1, + size: Tuple[float, float] = (0.06, 0.06, 0.04), + spawn_kwargs: Dict[str, Any] = {}, + procgen_seed_offset: int = 0, + procgen_kwargs: Dict[str, Any] = {}, + **kwargs, +) -> SampleCfg: + pose_range = { + "x": (-0.2, 0.2), + "y": (-0.3, 0.3), + "roll": (-torch.pi, torch.pi), + "pitch": (-torch.pi, torch.pi), + "yaw": (-torch.pi, torch.pi), + } + + match env_cfg.assets.object.variant: + case env_utils.AssetVariant.PRIMITIVE: + pose_range["z"] = (0.1, 0.1) + case env_utils.AssetVariant.DATASET: + match env_cfg.scenario: + case env_utils.Scenario.MARS: + pose_range.update( + { + "z": (0.03, 0.03), + "roll": (torch.pi / 7, torch.pi / 7), + "pitch": ( + 87.5 * torch.pi / 180, + 87.5 * torch.pi / 180, + ), + "yaw": (-torch.pi, torch.pi), + } + ) + case _: + pose_range["z"] = (0.05, 0.05) + case _: + pose_range["z"] = (0.04, 0.04) + + return SampleCfg( + asset_cfg=assets.object_of_interest_from_env_cfg( + env_cfg, + prim_path=prim_path, + num_assets=num_assets, + size=size, + spawn_kwargs=spawn_kwargs, + procgen_seed_offset=procgen_seed_offset, + procgen_kwargs=procgen_kwargs, + **kwargs, + ), + state_randomizer=EventTermCfg( + func=mdp.reset_root_state_uniform, + mode="reset", + params={ + "asset_cfg": asset_cfg, + "pose_range": pose_range, + "velocity_range": {}, + }, + ), + ) + + +@configclass +class TaskCfg(BaseManipulationEnvCfg): + ## Environment + episode_length_s: float = 7.5 + + ## Task + is_finite_horizon: bool = False + + ## Goal + target_pos = (-0.6, 0.0, 0.0) + target_quat = (1.0, 0.0, 0.0, 0.0) + target_marker_cfg = VisualizationMarkersCfg( + prim_path="/Visuals/target", + markers={ + "target": SphereCfg( + visible=False, + radius=0.025, + visual_material=sim_utils.PreviewSurfaceCfg( + diffuse_color=(1.0, 0.0, 0.0) + ), + ) + }, + ) + + ## Events + @configclass + class EventCfg(BaseManipulationEnvEventCfg): + ## Object + reset_rand_object_state: Optional[EventTermCfg] = None + + events = EventCfg() + + def __post_init__(self): + super().__post_init__() + + ## Scene + self.object_cfg = sample_cfg( + self.env_cfg, + num_assets=self.scene.num_envs, + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.55, 0.0, 0.0)), + spawn_kwargs={ + "activate_contact_sensors": True, + }, + ) + self.scene.object = self.object_cfg.asset_cfg + if self.vehicle_cfg: + self.target_pos = self.vehicle_cfg.frame_cargo_bay.offset.translation + + ## Sensors + self.scene.contacts_robot_hand_obj = ContactSensorCfg( + prim_path=f"{self.scene.robot.prim_path}/{self.robot_cfg.regex_links_hand}", + update_period=0.0, + # Note: This causes error 'Filter pattern did not match the correct number of entries' + # However, it seems to function properly anyway... + filter_prim_paths_expr=[self.scene.object.prim_path], + ) + + ## Events + self.events.reset_rand_object_state = self.object_cfg.state_randomizer + + +############ +### Task ### +############ + + +class Task(BaseManipulationEnv): + cfg: TaskCfg + + def __init__(self, cfg: TaskCfg, **kwargs): + super().__init__(cfg, **kwargs) + + ## Get handles to scene assets + self._contacts_robot_hand_obj: ContactSensor = self.scene[ + "contacts_robot_hand_obj" + ] + self._object: RigidObject = self.scene["object"] + + ## Pre-compute metrics used in hot loops + self._robot_arm_joint_indices, _ = self._robot.find_joints( + self.cfg.robot_cfg.regex_joints_arm + ) + self._robot_hand_joint_indices, _ = self._robot.find_joints( + self.cfg.robot_cfg.regex_joints_hand + ) + self._max_episode_length = self.max_episode_length + self._obj_com_offset = self._object.data._root_physx_view.get_coms().to( + self.device + ) + + ## Initialize buffers + self._initial_obj_height_w = torch.zeros( + self.num_envs, dtype=torch.float32, device=self.device + ) + self._target_pos_w = ( + torch.tensor( + self.cfg.target_pos, dtype=torch.float32, device=self.device + ).repeat(self.num_envs, 1) + + self.scene.env_origins + ) + self._target_quat_w = torch.tensor( + self.cfg.target_quat, dtype=torch.float32, device=self.device + ).repeat(self.num_envs, 1) + + ## Create visualization markers + self._target_marker = VisualizationMarkers(self.cfg.target_marker_cfg) + self._target_marker.visualize(self._target_pos_w, self._target_quat_w) + + ## Initialize the intermediate state + self._update_intermediate_state() + + def _reset_idx(self, env_ids: Sequence[int]): + super()._reset_idx(env_ids) + + # Update the initial height of the objects + self._initial_obj_height_w[env_ids] = self._object.data.root_pos_w[env_ids, 2] + + def _get_dones(self) -> Tuple[torch.Tensor, torch.Tensor]: + # Note: This assumes that `_get_dones()` is called before `_get_rewards()` and `_get_observations()` in `step()` + self._update_intermediate_state() + + if not self.cfg.enable_truncation: + self._truncations = torch.zeros_like(self._truncations) + + return self._terminations, self._truncations + + def _get_rewards(self) -> torch.Tensor: + return self._rewards + + def _get_observations(self) -> Dict[str, torch.Tensor]: + return _construct_observations( + remaining_time=self._remaining_time, + robot_joint_pos_arm=self._robot_joint_pos_arm, + robot_joint_pos_hand=self._robot_joint_pos_hand, + robot_ee_pos_wrt_base=self._robot_ee_pos_wrt_base, + robot_ee_rotmat_wrt_base=self._robot_ee_rotmat_wrt_base, + robot_hand_wrench=self._robot_hand_wrench, + obj_com_pos_wrt_robot_ee=self._obj_com_pos_wrt_robot_ee, + obj_com_rotmat_wrt_robot_ee=self._obj_com_rotmat_wrt_robot_ee, + target_pos_wrt_obj_com=self._target_pos_wrt_obj_com, + target_rotmat_wrt_obj_com=self._target_rotmat_wrt_obj_com, + ) + + ######################## + ### Helper Functions ### + ######################## + + def _update_intermediate_state(self): + ## Extract intermediate states + self._robot_ee_pos_wrt_base = self._tf_robot_ee.data.target_pos_source[:, 0, :] + self._robot_hand_wrench = ( + self._robot.root_physx_view.get_link_incoming_joint_force()[ + :, self._robot_hand_joint_indices + ] + ) + + ## Compute other intermediate states + ( + self._remaining_time, + self._robot_joint_pos_arm, + self._robot_joint_pos_hand, + self._robot_ee_rotmat_wrt_base, + self._obj_com_pos_wrt_robot_ee, + self._obj_com_rotmat_wrt_robot_ee, + self._target_pos_wrt_obj_com, + self._target_rotmat_wrt_obj_com, + self._rewards, + self._terminations, + self._truncations, + ) = _compute_intermediate_state( + current_action=self.action_manager.action, + previous_action=self.action_manager.prev_action, + episode_length_buf=self.episode_length_buf, + max_episode_length=self._max_episode_length, + robot_arm_joint_indices=self._robot_arm_joint_indices, + robot_hand_joint_indices=self._robot_hand_joint_indices, + joint_pos=self._robot.data.joint_pos, + soft_joint_pos_limits=self._robot.data.soft_joint_pos_limits, + robot_ee_pos_w=self._tf_robot_ee.data.target_pos_w[:, 0, :], + robot_ee_quat_w=self._tf_robot_ee.data.target_quat_w[:, 0, :], + robot_ee_quat_wrt_base=self._tf_robot_ee.data.target_quat_source[:, 0, :], + robot_arm_contact_net_forces=self._contacts_robot.data.net_forces_w, + robot_hand_obj_contact_force_matrix=self._contacts_robot_hand_obj.data.force_matrix_w, + obj_pos_w=self._object.data.root_pos_w, + obj_quat_w=self._object.data.root_quat_w, + obj_com_offset=self._obj_com_offset, + target_pos_w=self._target_pos_w, + target_quat_w=self._target_quat_w, + initial_obj_height_w=self._initial_obj_height_w, + ) + + +############################# +### TorchScript functions ### +############################# + + +@torch.jit.script +def _compute_intermediate_state( + *, + current_action: torch.Tensor, + previous_action: torch.Tensor, + episode_length_buf: torch.Tensor, + max_episode_length: int, + robot_arm_joint_indices: List[int], + robot_hand_joint_indices: List[int], + joint_pos: torch.Tensor, + soft_joint_pos_limits: torch.Tensor, + robot_ee_pos_w: torch.Tensor, + robot_ee_quat_w: torch.Tensor, + robot_ee_quat_wrt_base: torch.Tensor, + robot_arm_contact_net_forces: torch.Tensor, + robot_hand_obj_contact_force_matrix: torch.Tensor, + obj_pos_w: torch.Tensor, + obj_quat_w: torch.Tensor, + obj_com_offset: torch.Tensor, + target_pos_w: torch.Tensor, + target_quat_w: torch.Tensor, + initial_obj_height_w: torch.Tensor, +) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + ## Intermediate states + # Time + remaining_time = 1 - (episode_length_buf / max_episode_length).unsqueeze(-1) + + # Robot joint positions + joint_pos_normalized = math_utils.scale_transform( + joint_pos, + soft_joint_pos_limits[:, :, 0], + soft_joint_pos_limits[:, :, 1], + ) + robot_joint_pos_arm, robot_joint_pos_hand = ( + joint_pos_normalized[:, robot_arm_joint_indices], + joint_pos_normalized[:, robot_hand_joint_indices], + ) + + # End-effector pose (position and '6D' rotation) + robot_ee_rotmat_wrt_base = math_utils.matrix_from_quat(robot_ee_quat_wrt_base) + + # Transformation | Object origin -> Object CoM + obj_com_pos_w, obj_com_quat_w = math_utils.combine_frame_transforms( + t01=obj_pos_w, + q01=obj_quat_w, + t12=obj_com_offset[:, :3], + q12=obj_com_offset[:, 3:], + ) + + # Transformation | End-effector -> Object CoM + obj_com_pos_wrt_robot_ee, obj_com_quat_wrt_robot_ee = ( + math_utils.subtract_frame_transforms( + t01=robot_ee_pos_w, + q01=robot_ee_quat_w, + t02=obj_com_pos_w, + q02=obj_com_quat_w, + ) + ) + obj_com_rotmat_wrt_robot_ee = math_utils.matrix_from_quat(obj_com_quat_wrt_robot_ee) + + # Transformation | Object CoM -> Target + target_pos_wrt_obj_com, target_quat_wrt_obj_com = ( + math_utils.subtract_frame_transforms( + t01=obj_com_pos_w, + q01=obj_com_quat_w, + t02=target_pos_w, + q02=target_quat_w, + ) + ) + target_rotmat_wrt_obj_com = math_utils.matrix_from_quat(target_quat_wrt_obj_com) + + ## Rewards + # Penalty: Action rate + WEIGHT_ACTION_RATE = -0.05 + penalty_action_rate = WEIGHT_ACTION_RATE * torch.sum( + torch.square(current_action - previous_action), dim=1 + ) + + # Penalty: Undesired robot arm contacts + WEIGHT_UNDERSIRED_ROBOT_ARM_CONTACTS = -0.1 + THRESHOLD_UNDERSIRED_ROBOT_ARM_CONTACTS = 10.0 + penalty_undersired_robot_arm_contacts = WEIGHT_UNDERSIRED_ROBOT_ARM_CONTACTS * ( + torch.max(torch.norm(robot_arm_contact_net_forces, dim=-1), dim=1)[0] + > THRESHOLD_UNDERSIRED_ROBOT_ARM_CONTACTS + ) + + # Reward: Distance | End-effector <--> Object + WEIGHT_DISTANCE_EE_TO_OBJ = 1.0 + TANH_STD_DISTANCE_EE_TO_OBJ = 0.25 + reward_distance_ee_to_obj = WEIGHT_DISTANCE_EE_TO_OBJ * ( + 1.0 + - torch.tanh( + torch.norm(obj_com_pos_wrt_robot_ee, dim=-1) / TANH_STD_DISTANCE_EE_TO_OBJ + ) + ) + + # Reward: Object grasped + WEIGHT_OBJ_GRASPED = 4.0 + THRESHOLD_OBJ_GRASPED = 5.0 + reward_obj_grasped = WEIGHT_OBJ_GRASPED * ( + torch.mean( + torch.max(torch.norm(robot_hand_obj_contact_force_matrix, dim=-1), dim=-1)[ + 0 + ], + dim=1, + ) + > THRESHOLD_OBJ_GRASPED + ) + + # Reward: Object lifted + WEIGHT_OBJ_LIFTED = 8.0 + HEIGHT_OFFSET_OBJ_LIFTED = 0.5 + HEIGHT_SPAN_OBJ_LIFTED = 0.25 + TAHN_STD_HEIGHT_OBJ_LIFTED = 0.1 + obj_target_height_offset = ( + torch.abs(obj_com_pos_w[:, 2] - initial_obj_height_w - HEIGHT_OFFSET_OBJ_LIFTED) + - HEIGHT_SPAN_OBJ_LIFTED + ).clamp(min=0.0) + reward_obj_lifted = WEIGHT_OBJ_LIFTED * ( + 1.0 - torch.tanh(obj_target_height_offset / TAHN_STD_HEIGHT_OBJ_LIFTED) + ) + + # Reward: Distance | Object <--> Target + WEIGHT_DISTANCE_OBJ_TO_TARGET = 64.0 + TANH_STD_DISTANCE_OBJ_TO_TARGET = 0.333 + reward_distance_obj_to_target = WEIGHT_DISTANCE_OBJ_TO_TARGET * ( + 1.0 + - torch.tanh( + torch.norm(target_pos_wrt_obj_com, dim=-1) / TANH_STD_DISTANCE_OBJ_TO_TARGET + ) + ) + + # Total reward + rewards = torch.sum( + torch.stack( + [ + penalty_action_rate, + penalty_undersired_robot_arm_contacts, + reward_distance_ee_to_obj, + reward_obj_grasped, + reward_obj_lifted, + reward_distance_obj_to_target, + ], + dim=-1, + ), + dim=-1, + ) + + ## Termination and truncation + truncations = episode_length_buf > (max_episode_length - 1) + terminations = torch.zeros_like(truncations) + + # print( + # f""" + # penalty | action_rate: {float(penalty_action_rate[0])} + # penalty | undersired_robot_arm_contacts: {float(penalty_undersired_robot_arm_contacts[0])} + # reward | distance_ee_to_obj: {float(reward_distance_ee_to_obj[0])} + # reward | obj_grasped: {float(reward_obj_grasped[0])} + # reward | obj_lifted: {float(reward_obj_lifted[0])} + # reward | distance_obj_to_target: {float(reward_distance_obj_to_target[0])} + # total: {float(rewards[0])} + # """ + # ) + + return ( + remaining_time, + robot_joint_pos_arm, + robot_joint_pos_hand, + robot_ee_rotmat_wrt_base, + obj_com_pos_wrt_robot_ee, + obj_com_rotmat_wrt_robot_ee, + target_pos_wrt_obj_com, + target_rotmat_wrt_obj_com, + rewards, + terminations, + truncations, + ) + + +@torch.jit.script +def _construct_observations( + *, + remaining_time: torch.Tensor, + robot_joint_pos_arm: torch.Tensor, + robot_joint_pos_hand: torch.Tensor, + robot_ee_pos_wrt_base: torch.Tensor, + robot_ee_rotmat_wrt_base: torch.Tensor, + robot_hand_wrench: torch.Tensor, + obj_com_pos_wrt_robot_ee: torch.Tensor, + obj_com_rotmat_wrt_robot_ee: torch.Tensor, + target_pos_wrt_obj_com: torch.Tensor, + target_rotmat_wrt_obj_com: torch.Tensor, +) -> Dict[str, torch.Tensor]: + """ + Note: The `robot_hand_wrench` is considered as state (robot without force-torque sensors) + """ + + num_envs = remaining_time.size(0) + + # Robot joint positions + robot_joint_pos_hand_mean = robot_joint_pos_hand.mean(dim=-1, keepdim=True) + + # End-effector pose (position and '6D' rotation) + robot_ee_rot6d = math_utils.rotmat_to_rot6d(robot_ee_rotmat_wrt_base) + + # Wrench + robot_hand_wrench_full = robot_hand_wrench.view(num_envs, -1) + robot_hand_wrench_mean = robot_hand_wrench.mean(dim=1) + + # Transformation | End-effector -> Object CoM + obj_com_rot6d_wrt_robot_ee = math_utils.rotmat_to_rot6d(obj_com_rotmat_wrt_robot_ee) + + # Transformation | Object CoM -> Target + target_rot6d_wrt_obj_com = math_utils.rotmat_to_rot6d(target_rotmat_wrt_obj_com) + + return { + "state": torch.cat( + [ + obj_com_pos_wrt_robot_ee, + obj_com_rot6d_wrt_robot_ee, + target_pos_wrt_obj_com, + target_rot6d_wrt_obj_com, + robot_hand_wrench_mean, + ], + dim=-1, + ), + "state_dyn": torch.cat([robot_hand_wrench_full], dim=-1), + "proprio": torch.cat( + [ + remaining_time, + robot_ee_pos_wrt_base, + robot_ee_rot6d, + robot_joint_pos_hand_mean, + ], + dim=-1, + ), + "proprio_dyn": torch.cat([robot_joint_pos_arm, robot_joint_pos_hand], dim=-1), + } diff --git a/space_robotics_bench/tasks/sample_collection/task_visual.py b/space_robotics_bench/tasks/sample_collection/task_visual.py new file mode 100644 index 0000000..9b765bc --- /dev/null +++ b/space_robotics_bench/tasks/sample_collection/task_visual.py @@ -0,0 +1,32 @@ +from typing import Dict + +import torch +from omni.isaac.lab.utils import configclass + +from space_robotics_bench.envs import ( + VisualManipulationEnvExt, + VisualManipulationEnvExtCfg, +) + +from .task import Task, TaskCfg + + +@configclass +class VisualTaskCfg(TaskCfg, VisualManipulationEnvExtCfg): + def __post_init__(self): + TaskCfg.__post_init__(self) + VisualManipulationEnvExtCfg.__post_init__(self) + + +class VisualTask(Task, VisualManipulationEnvExt): + cfg: VisualTaskCfg + + def __init__(self, cfg: VisualTaskCfg, **kwargs): + Task.__init__(self, cfg, **kwargs) + VisualManipulationEnvExt.__init__(self, cfg, **kwargs) + + def _get_observations(self) -> Dict[str, torch.Tensor]: + return { + **Task._get_observations(self), + **VisualManipulationEnvExt._get_observations(self), + } diff --git a/space_robotics_bench/tasks/sample_collection_multi/__init__.py b/space_robotics_bench/tasks/sample_collection_multi/__init__.py new file mode 100644 index 0000000..b380c63 --- /dev/null +++ b/space_robotics_bench/tasks/sample_collection_multi/__init__.py @@ -0,0 +1,17 @@ +from space_robotics_bench.utils.registry import register_tasks + +from .task import Task, TaskCfg +from .task_visual import VisualTask, VisualTaskCfg + +BASE_TASK_NAME = __name__.split(".")[-1] +register_tasks( + { + BASE_TASK_NAME: {}, + f"{BASE_TASK_NAME}_visual": { + "entry_point": VisualTask, + "task_cfg": VisualTaskCfg, + }, + }, + default_entry_point=Task, + default_task_cfg=TaskCfg, +) diff --git a/space_robotics_bench/tasks/sample_collection_multi/task.py b/space_robotics_bench/tasks/sample_collection_multi/task.py new file mode 100644 index 0000000..d1d79c9 --- /dev/null +++ b/space_robotics_bench/tasks/sample_collection_multi/task.py @@ -0,0 +1,544 @@ +from dataclasses import MISSING as DELAYED_CFG +from typing import Dict, List, Optional, Sequence, Tuple + +import torch +from omni.isaac.lab.managers import EventTermCfg, SceneEntityCfg +from omni.isaac.lab.sensors import ContactSensor, ContactSensorCfg +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.envs as env_utils +import space_robotics_bench.core.sim as sim_utils +import space_robotics_bench.utils.math as math_utils +from space_robotics_bench.core.assets import RigidObject, RigidObjectCfg +from space_robotics_bench.core.markers import ( + VisualizationMarkers, + VisualizationMarkersCfg, +) +from space_robotics_bench.core.sim.spawners import SphereCfg +from space_robotics_bench.envs import ( + BaseManipulationEnv, + BaseManipulationEnvCfg, + BaseManipulationEnvEventCfg, + mdp, +) + +from ..sample_collection.task import sample_cfg + +############## +### Config ### +############## + + +@configclass +class TaskCfg(BaseManipulationEnvCfg): + ## Environment + episode_length_s: float = DELAYED_CFG + num_problems_per_env: int = 8 + + ## Task + is_finite_horizon: bool = False + + ## Goal + target_pos = (-0.6, 0.0, 0.0) + target_quat = (1.0, 0.0, 0.0, 0.0) + target_marker_cfg = VisualizationMarkersCfg( + prim_path="/Visuals/target", + markers={ + "target": SphereCfg( + visible=False, + radius=0.025, + visual_material=sim_utils.PreviewSurfaceCfg( + diffuse_color=(1.0, 0.0, 0.0) + ), + ) + }, + ) + + ## Events + @configclass + class EventCfg(BaseManipulationEnvEventCfg): + pass + + events = EventCfg() + + def __post_init__(self): + super().__post_init__() + + ## Environment + self.episode_length_s = self.num_problems_per_env * 7.5 + + ## Scene + self.object_cfgs = [ + sample_cfg( + self.env_cfg, + prim_path=f"{{ENV_REGEX_NS}}/sample{i}", + asset_cfg=SceneEntityCfg(f"object{i}"), + num_assets=self.scene.num_envs, + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.55, 0.0, 0.0)), + procgen_seed_offset=i * self.scene.num_envs, + spawn_kwargs={ + "activate_contact_sensors": True, + }, + ) + for i in range(self.num_problems_per_env) + ] + for i, object_cfg in enumerate(self.object_cfgs): + setattr(self.scene, f"object{i}", object_cfg.asset_cfg) + if self.vehicle_cfg: + self.target_pos = self.vehicle_cfg.frame_cargo_bay.offset.translation + + ## Sensors + self.scene.contacts_robot_hand_obj = ContactSensorCfg( + prim_path=f"{self.scene.robot.prim_path}/{self.robot_cfg.regex_links_hand}", + update_period=0.0, + # Note: This causes error 'Filter pattern did not match the correct number of entries' + # However, it seems to function properly anyway... + filter_prim_paths_expr=[ + asset.prim_path + for asset in [ + getattr(self.scene, f"object{i}") + for i in range(self.num_problems_per_env) + ] + ], + ) + + ## Events + self.events.reset_rand_object_state_multi = EventTermCfg( + func=mdp.reset_root_state_uniform_poisson_disk_2d, + mode="reset", + params={ + "asset_cfgs": [ + SceneEntityCfg(f"object{i}") + for i in range(self.num_problems_per_env) + ], + "pose_range": self.object_cfgs[0].state_randomizer.params["pose_range"], + "velocity_range": self.object_cfgs[0].state_randomizer.params[ + "velocity_range" + ], + "radius": ( + 0.225 + if self.env_cfg.assets.object.variant + == env_utils.AssetVariant.DATASET + else 0.06 + ), + }, + ) + + +############ +### Task ### +############ + + +class Task(BaseManipulationEnv): + cfg: TaskCfg + + def __init__(self, cfg: TaskCfg, **kwargs): + super().__init__(cfg, **kwargs) + + ## Get handles to scene assets + self._contacts_robot_hand_obj: ContactSensor = self.scene[ + "contacts_robot_hand_obj" + ] + self._objects: List[RigidObject] = [ + self.scene[f"object{i}"] for i in range(self.cfg.num_problems_per_env) + ] + + ## Pre-compute metrics used in hot loops + self._robot_arm_joint_indices, _ = self._robot.find_joints( + self.cfg.robot_cfg.regex_joints_arm + ) + self._robot_hand_joint_indices, _ = self._robot.find_joints( + self.cfg.robot_cfg.regex_joints_hand + ) + self._max_episode_length = self.max_episode_length + self._obj_com_offset = torch.stack( + [ + self._objects[i].data._root_physx_view.get_coms().to(self.device) + for i in range(self.cfg.num_problems_per_env) + ], + dim=1, + ) + + ## Initialize buffers + self._initial_obj_height_w = torch.zeros( + (self.num_envs, self.cfg.num_problems_per_env), + dtype=torch.float32, + device=self.device, + ) + self._target_pos_w = torch.tensor( + self.cfg.target_pos, dtype=torch.float32, device=self.device + ).repeat( + self.num_envs, self.cfg.num_problems_per_env, 1 + ) + self.scene.env_origins.unsqueeze( + 1 + ) + self._target_quat_w = torch.tensor( + self.cfg.target_quat, dtype=torch.float32, device=self.device + ).repeat(self.num_envs, self.cfg.num_problems_per_env, 1) + + ## Create visualization markers + self._target_marker = VisualizationMarkers(self.cfg.target_marker_cfg) + self._target_marker.visualize( + self._target_pos_w[:, 0, :], self._target_quat_w[:, 0, :] + ) + + ## Initialize the intermediate state + self._update_intermediate_state() + + def _reset_idx(self, env_ids: Sequence[int]): + super()._reset_idx(env_ids) + + # Update the initial height of the objects + self._initial_obj_height_w[env_ids] = torch.stack( + [ + self._objects[i].data.root_pos_w[env_ids, 2] + for i in range(self.cfg.num_problems_per_env) + ], + dim=1, + ) + + def _get_dones(self) -> Tuple[torch.Tensor, torch.Tensor]: + # Note: This assumes that `_get_dones()` is called before `_get_rewards()` and `_get_observations()` in `step()` + self._update_intermediate_state() + + if not self.cfg.enable_truncation: + self._truncations = torch.zeros_like(self._truncations) + + return self._terminations, self._truncations + + def _get_rewards(self) -> torch.Tensor: + return self._rewards + + def _get_observations(self) -> Dict[str, torch.Tensor]: + return _construct_observations( + remaining_time=self._remaining_time, + robot_joint_pos_arm=self._robot_joint_pos_arm, + robot_joint_pos_hand=self._robot_joint_pos_hand, + robot_ee_pos_wrt_base=self._robot_ee_pos_wrt_base, + robot_ee_rotmat_wrt_base=self._robot_ee_rotmat_wrt_base, + robot_hand_wrench=self._robot_hand_wrench, + obj_com_pos_wrt_robot_ee=self._obj_com_pos_wrt_robot_ee, + obj_com_rotmat_wrt_robot_ee=self._obj_com_rotmat_wrt_robot_ee, + target_pos_wrt_obj_com=self._target_pos_wrt_obj_com, + target_rotmat_wrt_obj_com=self._target_rotmat_wrt_obj_com, + ) + + ######################## + ### Helper Functions ### + ######################## + + def _update_intermediate_state(self): + ## Extract intermediate states + self._robot_ee_pos_wrt_base = self._tf_robot_ee.data.target_pos_source[:, 0, :] + self._robot_hand_wrench = ( + self._robot.root_physx_view.get_link_incoming_joint_force()[ + :, self._robot_hand_joint_indices + ] + ) + + ## Compute other intermediate states + ( + self._remaining_time, + self._robot_joint_pos_arm, + self._robot_joint_pos_hand, + self._robot_ee_rotmat_wrt_base, + self._obj_com_pos_wrt_robot_ee, + self._obj_com_rotmat_wrt_robot_ee, + self._target_pos_wrt_obj_com, + self._target_rotmat_wrt_obj_com, + self._rewards, + self._terminations, + self._truncations, + ) = _compute_intermediate_state( + current_action=self.action_manager.action, + previous_action=self.action_manager.prev_action, + episode_length_buf=self.episode_length_buf, + max_episode_length=self._max_episode_length, + robot_arm_joint_indices=self._robot_arm_joint_indices, + robot_hand_joint_indices=self._robot_hand_joint_indices, + joint_pos=self._robot.data.joint_pos, + soft_joint_pos_limits=self._robot.data.soft_joint_pos_limits, + robot_ee_pos_w=self._tf_robot_ee.data.target_pos_w[:, 0, :], + robot_ee_quat_w=self._tf_robot_ee.data.target_quat_w[:, 0, :], + robot_ee_quat_wrt_base=self._tf_robot_ee.data.target_quat_source[:, 0, :], + robot_arm_contact_net_forces=self._contacts_robot.data.net_forces_w, + robot_hand_obj_contact_force_matrix=self._contacts_robot_hand_obj.data.force_matrix_w, + obj_pos_w=torch.stack( + [ + self._objects[i].data.root_pos_w + for i in range(self.cfg.num_problems_per_env) + ], + dim=1, + ), + obj_quat_w=torch.stack( + [ + self._objects[i].data.root_quat_w + for i in range(self.cfg.num_problems_per_env) + ], + dim=1, + ), + obj_com_offset=self._obj_com_offset, + target_pos_w=self._target_pos_w, + target_quat_w=self._target_quat_w, + initial_obj_height_w=self._initial_obj_height_w, + ) + + +############################# +### TorchScript functions ### +############################# + + +@torch.jit.script +def _compute_intermediate_state( + *, + current_action: torch.Tensor, + previous_action: torch.Tensor, + episode_length_buf: torch.Tensor, + max_episode_length: int, + robot_arm_joint_indices: List[int], + robot_hand_joint_indices: List[int], + joint_pos: torch.Tensor, + soft_joint_pos_limits: torch.Tensor, + robot_ee_pos_w: torch.Tensor, + robot_ee_quat_w: torch.Tensor, + robot_ee_quat_wrt_base: torch.Tensor, + robot_arm_contact_net_forces: torch.Tensor, + robot_hand_obj_contact_force_matrix: torch.Tensor, + obj_pos_w: torch.Tensor, + obj_quat_w: torch.Tensor, + obj_com_offset: torch.Tensor, + target_pos_w: torch.Tensor, + target_quat_w: torch.Tensor, + initial_obj_height_w: torch.Tensor, +) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + ## Intermediate states + # Time + remaining_time = 1 - (episode_length_buf / max_episode_length).unsqueeze(-1) + + # Robot joint positions + joint_pos_normalized = math_utils.scale_transform( + joint_pos, + soft_joint_pos_limits[:, :, 0], + soft_joint_pos_limits[:, :, 1], + ) + robot_joint_pos_arm, robot_joint_pos_hand = ( + joint_pos_normalized[:, robot_arm_joint_indices], + joint_pos_normalized[:, robot_hand_joint_indices], + ) + + # End-effector pose (position and '6D' rotation) + robot_ee_rotmat_wrt_base = math_utils.matrix_from_quat(robot_ee_quat_wrt_base) + + # Transformation | Object origin -> Object CoM + obj_com_pos_w, obj_com_quat_w = math_utils.combine_frame_transforms( + t01=obj_pos_w, + q01=obj_quat_w, + t12=obj_com_offset[:, :, :3], + q12=obj_com_offset[:, :, 3:], + ) + + # Transformation | End-effector -> Object CoM + obj_com_pos_wrt_robot_ee, obj_com_quat_wrt_robot_ee = ( + math_utils.subtract_frame_transforms( + t01=robot_ee_pos_w.unsqueeze(1).repeat(1, obj_pos_w.shape[1], 1), + q01=robot_ee_quat_w.unsqueeze(1).repeat(1, obj_pos_w.shape[1], 1), + t02=obj_com_pos_w, + q02=obj_com_quat_w, + ) + ) + obj_com_rotmat_wrt_robot_ee = math_utils.matrix_from_quat(obj_com_quat_wrt_robot_ee) + + # Transformation | Object CoM -> Target + target_pos_wrt_obj_com, target_quat_wrt_obj_com = ( + math_utils.subtract_frame_transforms( + t01=obj_com_pos_w, + q01=obj_com_quat_w, + t02=target_pos_w, + q02=target_quat_w, + ) + ) + target_rotmat_wrt_obj_com = math_utils.matrix_from_quat(target_quat_wrt_obj_com) + + ## Rewards + # Penalty: Action rate + WEIGHT_ACTION_RATE = -0.05 + penalty_action_rate = WEIGHT_ACTION_RATE * torch.sum( + torch.square(current_action - previous_action), dim=1 + ) + + # Penalty: Undesired robot arm contacts + WEIGHT_UNDERSIRED_ROBOT_ARM_CONTACTS = -0.1 + THRESHOLD_UNDERSIRED_ROBOT_ARM_CONTACTS = 10.0 + penalty_undersired_robot_arm_contacts = WEIGHT_UNDERSIRED_ROBOT_ARM_CONTACTS * ( + torch.max(torch.norm(robot_arm_contact_net_forces, dim=-1), dim=1)[0] + > THRESHOLD_UNDERSIRED_ROBOT_ARM_CONTACTS + ) + + # Reward: Distance | End-effector <--> Object + WEIGHT_DISTANCE_EE_TO_OBJ = 1.0 + TANH_STD_DISTANCE_EE_TO_OBJ = 0.25 + reward_distance_ee_to_obj = WEIGHT_DISTANCE_EE_TO_OBJ * ( + 1.0 + - torch.tanh( + torch.norm(obj_com_pos_wrt_robot_ee, dim=-1) / TANH_STD_DISTANCE_EE_TO_OBJ + ) + ).sum(dim=-1) + + # Reward: Object grasped + WEIGHT_OBJ_GRASPED = 4.0 + THRESHOLD_OBJ_GRASPED = 5.0 + reward_obj_grasped = WEIGHT_OBJ_GRASPED * ( + torch.mean( + torch.max(torch.norm(robot_hand_obj_contact_force_matrix, dim=-1), dim=-1)[ + 0 + ], + dim=1, + ) + > THRESHOLD_OBJ_GRASPED + ) + + # Reward: Object lifted + WEIGHT_OBJ_LIFTED = 8.0 + HEIGHT_OFFSET_OBJ_LIFTED = 0.5 + HEIGHT_SPAN_OBJ_LIFTED = 0.25 + TAHN_STD_HEIGHT_OBJ_LIFTED = 0.1 + obj_target_height_offset = ( + torch.abs( + obj_com_pos_w[:, :, 2] - initial_obj_height_w - HEIGHT_OFFSET_OBJ_LIFTED + ) + - HEIGHT_SPAN_OBJ_LIFTED + ).clamp(min=0.0) + reward_obj_lifted = WEIGHT_OBJ_LIFTED * ( + 1.0 - torch.tanh(obj_target_height_offset / TAHN_STD_HEIGHT_OBJ_LIFTED) + ).sum(dim=-1) + + # Reward: Distance | Object <--> Target + WEIGHT_DISTANCE_OBJ_TO_TARGET = 64.0 + TANH_STD_DISTANCE_OBJ_TO_TARGET = 0.333 + reward_distance_obj_to_target = WEIGHT_DISTANCE_OBJ_TO_TARGET * ( + 1.0 + - torch.tanh( + torch.norm(target_pos_wrt_obj_com, dim=-1) / TANH_STD_DISTANCE_OBJ_TO_TARGET + ) + ).sum(dim=-1) + + # Total reward + rewards = torch.sum( + torch.stack( + [ + penalty_action_rate, + penalty_undersired_robot_arm_contacts, + reward_distance_ee_to_obj, + reward_obj_grasped, + reward_obj_lifted, + reward_distance_obj_to_target, + ], + dim=-1, + ), + dim=-1, + ) + + ## Termination and truncation + truncations = episode_length_buf > (max_episode_length - 1) + terminations = torch.zeros_like(truncations) + + # print( + # f""" + # penalty | action_rate: {float(penalty_action_rate[0])} + # penalty | undersired_robot_arm_contacts: {float(penalty_undersired_robot_arm_contacts[0])} + # reward | distance_ee_to_obj: {float(reward_distance_ee_to_obj[0])} + # reward | obj_grasped: {float(reward_obj_grasped[0])} + # reward | obj_lifted: {float(reward_obj_lifted[0])} + # reward | distance_obj_to_target: {float(reward_distance_obj_to_target[0])} + # total: {float(rewards[0])} + # """ + # ) + + return ( + remaining_time, + robot_joint_pos_arm, + robot_joint_pos_hand, + robot_ee_rotmat_wrt_base, + obj_com_pos_wrt_robot_ee, + obj_com_rotmat_wrt_robot_ee, + target_pos_wrt_obj_com, + target_rotmat_wrt_obj_com, + rewards, + terminations, + truncations, + ) + + +@torch.jit.script +def _construct_observations( + *, + remaining_time: torch.Tensor, + robot_joint_pos_arm: torch.Tensor, + robot_joint_pos_hand: torch.Tensor, + robot_ee_pos_wrt_base: torch.Tensor, + robot_ee_rotmat_wrt_base: torch.Tensor, + robot_hand_wrench: torch.Tensor, + obj_com_pos_wrt_robot_ee: torch.Tensor, + obj_com_rotmat_wrt_robot_ee: torch.Tensor, + target_pos_wrt_obj_com: torch.Tensor, + target_rotmat_wrt_obj_com: torch.Tensor, +) -> Dict[str, torch.Tensor]: + """ + Note: The `robot_hand_wrench` is considered as state (robot without force-torque sensors) + """ + + num_envs = remaining_time.size(0) + + # Robot joint positions + robot_joint_pos_hand_mean = robot_joint_pos_hand.mean(dim=-1, keepdim=True) + + # End-effector pose (position and '6D' rotation) + robot_ee_rot6d = math_utils.rotmat_to_rot6d(robot_ee_rotmat_wrt_base) + + # Wrench + robot_hand_wrench_full = robot_hand_wrench.view(num_envs, -1) + robot_hand_wrench_mean = robot_hand_wrench.mean(dim=1) + + # Transformation | End-effector -> Object CoM + obj_com_rot6d_wrt_robot_ee = math_utils.rotmat_to_rot6d(obj_com_rotmat_wrt_robot_ee) + + # Transformation | Object CoM -> Target + target_rot6d_wrt_obj_com = math_utils.rotmat_to_rot6d(target_rotmat_wrt_obj_com) + + return { + "state": torch.cat( + [ + obj_com_pos_wrt_robot_ee.view(num_envs, -1), + obj_com_rot6d_wrt_robot_ee.view(num_envs, -1), + target_pos_wrt_obj_com.view(num_envs, -1), + target_rot6d_wrt_obj_com.view(num_envs, -1), + robot_hand_wrench_mean.view(num_envs, -1), + ], + dim=-1, + ), + "state_dyn": torch.cat([robot_hand_wrench_full], dim=-1), + "proprio": torch.cat( + [ + remaining_time, + robot_ee_pos_wrt_base, + robot_ee_rot6d, + robot_joint_pos_hand_mean, + ], + dim=-1, + ), + "proprio_dyn": torch.cat([robot_joint_pos_arm, robot_joint_pos_hand], dim=-1), + } diff --git a/space_robotics_bench/tasks/sample_collection_multi/task_visual.py b/space_robotics_bench/tasks/sample_collection_multi/task_visual.py new file mode 100644 index 0000000..9b765bc --- /dev/null +++ b/space_robotics_bench/tasks/sample_collection_multi/task_visual.py @@ -0,0 +1,32 @@ +from typing import Dict + +import torch +from omni.isaac.lab.utils import configclass + +from space_robotics_bench.envs import ( + VisualManipulationEnvExt, + VisualManipulationEnvExtCfg, +) + +from .task import Task, TaskCfg + + +@configclass +class VisualTaskCfg(TaskCfg, VisualManipulationEnvExtCfg): + def __post_init__(self): + TaskCfg.__post_init__(self) + VisualManipulationEnvExtCfg.__post_init__(self) + + +class VisualTask(Task, VisualManipulationEnvExt): + cfg: VisualTaskCfg + + def __init__(self, cfg: VisualTaskCfg, **kwargs): + Task.__init__(self, cfg, **kwargs) + VisualManipulationEnvExt.__init__(self, cfg, **kwargs) + + def _get_observations(self) -> Dict[str, torch.Tensor]: + return { + **Task._get_observations(self), + **VisualManipulationEnvExt._get_observations(self), + } diff --git a/space_robotics_bench/tasks/solar_panel_assembly/__init__.py b/space_robotics_bench/tasks/solar_panel_assembly/__init__.py new file mode 100644 index 0000000..b380c63 --- /dev/null +++ b/space_robotics_bench/tasks/solar_panel_assembly/__init__.py @@ -0,0 +1,17 @@ +from space_robotics_bench.utils.registry import register_tasks + +from .task import Task, TaskCfg +from .task_visual import VisualTask, VisualTaskCfg + +BASE_TASK_NAME = __name__.split(".")[-1] +register_tasks( + { + BASE_TASK_NAME: {}, + f"{BASE_TASK_NAME}_visual": { + "entry_point": VisualTask, + "task_cfg": VisualTaskCfg, + }, + }, + default_entry_point=Task, + default_task_cfg=TaskCfg, +) diff --git a/space_robotics_bench/tasks/solar_panel_assembly/task.py b/space_robotics_bench/tasks/solar_panel_assembly/task.py new file mode 100644 index 0000000..d4ae889 --- /dev/null +++ b/space_robotics_bench/tasks/solar_panel_assembly/task.py @@ -0,0 +1,928 @@ +import os +import sys +from typing import Any, Dict, List, Optional, Sequence, Tuple + +import torch +from omni.isaac.core.prims.xform_prim_view import XFormPrimView +from omni.isaac.lab.sensors import ContactSensor, ContactSensorCfg +from omni.isaac.lab.utils import configclass + +import space_robotics_bench.core.envs as env_utils +import space_robotics_bench.core.sim as sim_utils +import space_robotics_bench.utils.math as math_utils +from space_robotics_bench import assets +from space_robotics_bench.core.assets import ( + AssetBaseCfg, + AssetCfg, + RigidObject, + RigidObjectCfg, +) +from space_robotics_bench.core.managers import EventTermCfg, SceneEntityCfg +from space_robotics_bench.core.markers import ( + VisualizationMarkers, + VisualizationMarkersCfg, +) +from space_robotics_bench.core.sim.schemas import RigidBodyPropertiesCfg +from space_robotics_bench.envs import ( + BaseManipulationEnv, + BaseManipulationEnvCfg, + BaseManipulationEnvEventCfg, + mdp, +) + +from ..peg_in_hole.task import peg_and_hole_cfg + +############## +### Config ### +############## + + +class PanelCfg(AssetCfg): + class Config: + arbitrary_types_allowed = True # Due to EventTermCfg + + ## Geometry + offset_pos: Tuple[float, float, float] = (0.0, 0.0, 0.15) + + ## Randomization + state_randomizer: Optional[EventTermCfg] = None + + +@configclass +class TaskCfg(BaseManipulationEnvCfg): + ## Environment + episode_length_s: float = 50.0 + + ## Task + is_finite_horizon: bool = False + + ## Panel + panel_target_pos = (0.55, 0.0, 0.0) + panel_target_quat = (1.0, 0.0, 0.0, 0.0) + panel_target_marker_cfg = VisualizationMarkersCfg( + prim_path="/Visuals/panel_target", + markers={ + "target": assets.object.SolarPanelCfg( + visible=False, + visual_material=sim_utils.PreviewSurfaceCfg( + emissive_color=(0.2, 0.2, 0.2), + opacity=0.1, + diffuse_color=(0.0, 0.0, 0.0), + metallic=0.0, + roughness=1.0, + ), + ) + }, + ) + + ## Events + @configclass + class EventCfg(BaseManipulationEnvEventCfg): + pass + + events = EventCfg() + + def __post_init__(self): + if self.env_cfg.assets.object.variant != env_utils.AssetVariant.DATASET: + print( + f"[WARN] Environment requires DATASET object ({self.env_cfg.assets.object.variant} ignored)", + file=sys.stderr, + ) + self.env_cfg.assets.object.variant = env_utils.AssetVariant.DATASET + + super().__post_init__() + + ## Scene + self.problem_cfgs = [ + peg_and_hole_cfg( + self.env_cfg, + prim_path_peg=f"{{ENV_REGEX_NS}}/peg{i}", + prim_path_hole=f"{{ENV_REGEX_NS}}/hole{i}", + asset_cfg_peg=SceneEntityCfg(f"object{i}"), + init_state=RigidObjectCfg.InitialStateCfg(pos=init_pos), + spawn_kwargs_peg={ + "activate_contact_sensors": True, + }, + short_peg=short_peg, + ) + for i, (init_pos, short_peg) in enumerate( + [ + ((0.55 + 0.1, 0.2, 0.015), True), + ((0.55 + 0.1, -0.2, 0.015), True), + ((0.55 - 0.1, 0.2, 0.015), False), + ((0.55 - 0.1, -0.2, 0.015), False), + ] + ) + ] + for i, problem_cfg in enumerate(self.problem_cfgs): + setattr( + self.scene, + f"object{i}", + problem_cfg[0].asset_cfg.replace( + # init_state=RigidObjectCfg.InitialStateCfg( + # pos=(0.55, 0.0, 0.015), + # ) + ), + ) + setattr( + self.scene, + f"target{i}", + problem_cfg[1].asset_cfg, + ) + + self.panel_cfg = self._panel_cfg( + self.env_cfg, + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.55, 0.0, 0.015)), + ) + self.scene.panel = self.panel_cfg.asset_cfg + self.panel_target_pos = ( + self.panel_target_pos[0] + self.panel_cfg.offset_pos[0], + self.panel_target_pos[1] + self.panel_cfg.offset_pos[1], + self.panel_target_pos[2] + self.panel_cfg.offset_pos[2] + 0.015, + ) + + ## Sensors + self.scene.contacts_robot_hand_obj = ContactSensorCfg( + prim_path=f"{self.scene.robot.prim_path}/{self.robot_cfg.regex_links_hand}", + update_period=0.0, + # Note: This causes error 'Filter pattern did not match the correct number of entries' + # However, it seems to function properly anyway... + filter_prim_paths_expr=[ + asset.prim_path + for asset in [getattr(self.scene, f"object{i}") for i in range(4)] + ], + ) + + ## Events + self.events.reset_rand_panel_state = self.panel_cfg.state_randomizer + self.events.reset_rand_object_state_multi = EventTermCfg( + func=mdp.reset_root_state_uniform_poisson_disk_2d, + mode="reset", + params={ + "asset_cfgs": [SceneEntityCfg(f"object{i}") for i in range(4)], + "pose_range": { + "x": ( + -0.2, + 0.05, + ), + "y": ( + -0.05, + 0.05, + ), + "z": ( + 0.025, + 0.025, + ), + "roll": (torch.pi / 2, torch.pi / 2), + "pitch": (-torch.pi, torch.pi), + "yaw": (-torch.pi, torch.pi), + }, + "velocity_range": {}, + "radius": 0.2, + }, + ) + + ######################## + ### Helper Functions ### + ######################## + + @staticmethod + def _panel_cfg( + env_cfg: env_utils.EnvironmentConfig, + *, + asset_cfg: SceneEntityCfg = SceneEntityCfg("panel"), + prim_path: str = "{ENV_REGEX_NS}/panel", + spawn_kwargs: Dict[str, Any] = {}, + **kwargs, + ) -> PanelCfg: + offset_pos = (0.0, 0.0, 0.15) + pose_range = { + "x": (-0.1 - 0.025, -0.1 + 0.025), + "y": (-0.45 - 0.025, -0.45 + 0.025), + "z": (-0.14, -0.14), + "roll": (torch.pi, torch.pi), + "pitch": (-torch.pi + torch.pi / 7, -torch.pi + torch.pi / 7), + "yaw": ( + torch.pi / 2 - torch.pi / 32, + torch.pi / 2 + torch.pi / 32, + ), + } + + if kwargs.get("init_state") is not None: + kwargs["init_state"].pos = ( + *kwargs["init_state"].pos[:2], + kwargs["init_state"].pos[2] + offset_pos[2], + ) + + return PanelCfg( + asset_cfg=assets.solar_panel_from_env_cfg( + env_cfg=env_cfg, + prim_path=prim_path, + spawn_kwargs=spawn_kwargs, + **kwargs, + ), + offset_pos=offset_pos, + state_randomizer=EventTermCfg( + func=mdp.reset_root_state_uniform, + mode="reset", + params={ + "asset_cfg": asset_cfg, + "pose_range": pose_range, + "velocity_range": {}, + }, + ), + ) + + +############ +### Task ### +############ + + +class Task(BaseManipulationEnv): + cfg: TaskCfg + + def __init__(self, cfg: TaskCfg, **kwargs): + super().__init__(cfg, **kwargs) + + # Get handles to scene assets + self._contacts_robot_hand_obj: ContactSensor = self.scene[ + "contacts_robot_hand_obj" + ] + self._objects: List[RigidObject] = [self.scene[f"object{i}"] for i in range(4)] + self._targets: List[XFormPrimView] = [ + self.scene[f"target{i}"] for i in range(4) + ] + self._panel: RigidObject = self.scene["panel"] + + ## Pre-compute metrics used in hot loops + self._robot_arm_joint_indices, _ = self._robot.find_joints( + self.cfg.robot_cfg.regex_joints_arm + ) + self._robot_hand_joint_indices, _ = self._robot.find_joints( + self.cfg.robot_cfg.regex_joints_hand + ) + self._max_episode_length = self.max_episode_length + self._obj_com_offset = torch.stack( + [ + self._objects[i].data._root_physx_view.get_coms().to(self.device) + for i in range(4) + ], + dim=1, + ) + self._panel_com_offset = self._panel.data._root_physx_view.get_coms().to( + self.device + ) + + ## Initialize buffers + self._initial_obj_height_w = torch.zeros( + (self.num_envs, 4), + dtype=torch.float32, + device=self.device, + ) + self._peg_offset_pos_ends = torch.tensor( + [self.cfg.problem_cfgs[i][0].offset_pos_ends for i in range(4)], + dtype=torch.float32, + device=self.device, + ).repeat(self.num_envs, 1, 1, 1) + + self._peg_rot_symmetry_n = torch.tensor( + [self.cfg.problem_cfgs[i][0].rot_symmetry_n for i in range(4)], + dtype=torch.int32, + device=self.device, + ).repeat(self.num_envs, 1) + self._hole_offset_pos_bottom = torch.tensor( + [self.cfg.problem_cfgs[i][1].offset_pos_bottom for i in range(4)], + dtype=torch.float32, + device=self.device, + ).repeat(self.num_envs, 1, 1) + self._hole_offset_pos_entrance = torch.tensor( + [self.cfg.problem_cfgs[i][1].offset_pos_entrance for i in range(4)], + dtype=torch.float32, + device=self.device, + ).repeat(self.num_envs, 1, 1) + self._initial_panel_height_w = torch.zeros( + self.num_envs, dtype=torch.float32, device=self.device + ) + self._panel_target_pos_w = ( + torch.tensor( + self.cfg.panel_target_pos, dtype=torch.float32, device=self.device + ).repeat(self.num_envs, 1) + + self.scene.env_origins + ) + self._panel_target_quat_w = torch.tensor( + self.cfg.panel_target_quat, dtype=torch.float32, device=self.device + ).repeat(self.num_envs, 1) + + ## Create visualization markers + self._panel_target_marker = VisualizationMarkers( + self.cfg.panel_target_marker_cfg + ) + self._panel_target_marker.visualize( + self._panel_target_pos_w, self._panel_target_quat_w + ) + + ## Initialize the intermediate state + self._update_intermediate_state() + + def _reset_idx(self, env_ids: Sequence[int]): + super()._reset_idx(env_ids) + + # Update the initial height of the objects + self._initial_obj_height_w[env_ids] = torch.stack( + [self._objects[i].data.root_pos_w[env_ids, 2] for i in range(4)], + dim=1, + ) + self._initial_panel_height_w[env_ids] = self._panel.data.root_pos_w[env_ids, 2] + + def _get_dones(self) -> Tuple[torch.Tensor, torch.Tensor]: + # Note: This assumes that `_get_dones()` is called before `_get_rewards()` and `_get_observations()` in `step()` + self._update_intermediate_state() + + if not self.cfg.enable_truncation: + self._truncations = torch.zeros_like(self._truncations) + + return self._terminations, self._truncations + + def _get_rewards(self) -> torch.Tensor: + return self._rewards + + def _get_observations(self) -> Dict[str, torch.Tensor]: + return _construct_observations( + remaining_time=self._remaining_time, + robot_joint_pos_arm=self._robot_joint_pos_arm, + robot_joint_pos_hand=self._robot_joint_pos_hand, + robot_ee_pos_wrt_base=self._robot_ee_pos_wrt_base, + robot_ee_rotmat_wrt_base=self._robot_ee_rotmat_wrt_base, + robot_hand_wrench=self._robot_hand_wrench, + obj_com_pos_wrt_robot_ee=self._obj_com_pos_wrt_robot_ee, + obj_com_rotmat_wrt_robot_ee=self._obj_com_rotmat_wrt_robot_ee, + hole_entrance_pos_wrt_peg_ends=self._hole_entrance_pos_wrt_peg_ends, + hole_bottom_pos_wrt_peg_ends=self._hole_bottom_pos_wrt_peg_ends, + hole_rotmat_wrt_peg=self._hole_rotmat_wrt_peg, + panel_com_pos_wrt_robot_ee=self._panel_com_pos_wrt_robot_ee, + panel_com_rotmat_wrt_robot_ee=self._panel_com_rotmat_wrt_robot_ee, + panel_target_pos_wrt_panel_com=self._panel_target_pos_wrt_panel_com, + panel_target_rotmat_wrt_panel_com=self._panel_target_rotmat_wrt_panel_com, + ) + + ######################## + ### Helper Functions ### + ######################## + + def _update_intermediate_state(self): + ## Extract intermediate states + self._robot_ee_pos_wrt_base = self._tf_robot_ee.data.target_pos_source[:, 0, :] + self._robot_hand_wrench = ( + self._robot.root_physx_view.get_link_incoming_joint_force()[ + :, self._robot_hand_joint_indices + ] + ) + + ## Compute other intermediate states + target_pos_w = torch.stack( + [self._targets[i].get_world_poses()[0] for i in range(4)], + dim=1, + ) + target_quat_w = torch.stack( + [self._targets[i].get_world_poses()[1] for i in range(4)], + dim=1, + ) + ( + self._remaining_time, + self._robot_joint_pos_arm, + self._robot_joint_pos_hand, + self._robot_ee_rotmat_wrt_base, + self._obj_com_pos_wrt_robot_ee, + self._obj_com_rotmat_wrt_robot_ee, + self._hole_entrance_pos_wrt_peg_ends, + self._hole_bottom_pos_wrt_peg_ends, + self._hole_rotmat_wrt_peg, + self._panel_com_pos_wrt_robot_ee, + self._panel_com_rotmat_wrt_robot_ee, + self._panel_target_pos_wrt_panel_com, + self._panel_target_rotmat_wrt_panel_com, + self._rewards, + self._terminations, + self._truncations, + ) = _compute_intermediate_state( + current_action=self.action_manager.action, + previous_action=self.action_manager.prev_action, + episode_length_buf=self.episode_length_buf, + max_episode_length=self._max_episode_length, + robot_arm_joint_indices=self._robot_arm_joint_indices, + robot_hand_joint_indices=self._robot_hand_joint_indices, + joint_pos=self._robot.data.joint_pos, + soft_joint_pos_limits=self._robot.data.soft_joint_pos_limits, + robot_ee_pos_w=self._tf_robot_ee.data.target_pos_w[:, 0, :], + robot_ee_quat_w=self._tf_robot_ee.data.target_quat_w[:, 0, :], + robot_ee_quat_wrt_base=self._tf_robot_ee.data.target_quat_source[:, 0, :], + robot_arm_contact_net_forces=self._contacts_robot.data.net_forces_w, + robot_hand_obj_contact_force_matrix=self._contacts_robot_hand_obj.data.force_matrix_w, + obj_pos_w=torch.stack( + [self._objects[i].data.root_pos_w for i in range(4)], + dim=1, + ), + obj_quat_w=torch.stack( + [self._objects[i].data.root_quat_w for i in range(4)], + dim=1, + ), + obj_com_offset=self._obj_com_offset, + target_pos_w=target_pos_w, + target_quat_w=target_quat_w, + panel_pos_w=self._panel.data.root_pos_w, + panel_quat_w=self._panel.data.root_quat_w, + panel_com_offset=self._panel_com_offset, + panel_target_pos_w=self._panel_target_pos_w, + panel_target_quat_w=self._panel_target_quat_w, + initial_obj_height_w=self._initial_obj_height_w, + initial_panel_height_w=self._initial_panel_height_w, + peg_offset_pos_ends=self._peg_offset_pos_ends, + peg_rot_symmetry_n=self._peg_rot_symmetry_n, + hole_offset_pos_bottom=self._hole_offset_pos_bottom, + hole_offset_pos_entrance=self._hole_offset_pos_entrance, + ) + + +############################# +### TorchScript functions ### +############################# + + +@torch.jit.script +def _compute_intermediate_state( + *, + current_action: torch.Tensor, + previous_action: torch.Tensor, + episode_length_buf: torch.Tensor, + max_episode_length: int, + robot_arm_joint_indices: List[int], + robot_hand_joint_indices: List[int], + joint_pos: torch.Tensor, + soft_joint_pos_limits: torch.Tensor, + robot_ee_pos_w: torch.Tensor, + robot_ee_quat_w: torch.Tensor, + robot_ee_quat_wrt_base: torch.Tensor, + robot_arm_contact_net_forces: torch.Tensor, + robot_hand_obj_contact_force_matrix: torch.Tensor, + obj_pos_w: torch.Tensor, + obj_quat_w: torch.Tensor, + obj_com_offset: torch.Tensor, + target_pos_w: torch.Tensor, + target_quat_w: torch.Tensor, + panel_pos_w: torch.Tensor, + panel_quat_w: torch.Tensor, + panel_com_offset: torch.Tensor, + panel_target_pos_w: torch.Tensor, + panel_target_quat_w: torch.Tensor, + initial_obj_height_w: torch.Tensor, + initial_panel_height_w: torch.Tensor, + peg_offset_pos_ends: torch.Tensor, + peg_rot_symmetry_n: torch.Tensor, + hole_offset_pos_bottom: torch.Tensor, + hole_offset_pos_entrance: torch.Tensor, +) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + ## Intermediate states + # Time + remaining_time = 1 - (episode_length_buf / max_episode_length).unsqueeze(-1) + + # Robot joint positions + joint_pos_normalized = math_utils.scale_transform( + joint_pos, + soft_joint_pos_limits[:, :, 0], + soft_joint_pos_limits[:, :, 1], + ) + robot_joint_pos_arm, robot_joint_pos_hand = ( + joint_pos_normalized[:, robot_arm_joint_indices], + joint_pos_normalized[:, robot_hand_joint_indices], + ) + + # End-effector pose (position and '6D' rotation) + robot_ee_rotmat_wrt_base = math_utils.matrix_from_quat(robot_ee_quat_wrt_base) + + # Transformation | Object origin -> Object CoM + obj_com_pos_w, obj_com_quat_w = math_utils.combine_frame_transforms( + t01=obj_pos_w, + q01=obj_quat_w, + t12=obj_com_offset[:, :, :3], + q12=obj_com_offset[:, :, 3:], + ) + + # Transformation | Object origin -> Peg ends + _peg_end0_pos_w, _ = math_utils.combine_frame_transforms( + t01=obj_pos_w, + q01=obj_quat_w, + t12=peg_offset_pos_ends[:, :, 0], + ) + _peg_end1_pos_w, _ = math_utils.combine_frame_transforms( + t01=obj_pos_w, + q01=obj_quat_w, + t12=peg_offset_pos_ends[:, :, 1], + ) + peg_ends_pos_w = torch.stack([_peg_end0_pos_w, _peg_end1_pos_w], dim=1) + + # Transformation | Panel origin -> Panel CoM + panel_com_pos_w, panel_com_quat_w = math_utils.combine_frame_transforms( + t01=panel_pos_w, + q01=panel_quat_w, + t12=panel_com_offset[:, :3], + q12=panel_com_offset[:, 3:], + ) + + # Transformation | End-effector -> Panel CoM + panel_com_pos_wrt_robot_ee, panel_com_quat_wrt_robot_ee = ( + math_utils.subtract_frame_transforms( + t01=robot_ee_pos_w, + q01=robot_ee_quat_w, + t02=panel_com_pos_w, + q02=panel_com_quat_w, + ) + ) + panel_com_rotmat_wrt_robot_ee = math_utils.matrix_from_quat( + panel_com_quat_wrt_robot_ee + ) + + # Transformation | Panel CoM -> Panel Target + panel_target_pos_wrt_panel_com, panel_target_quat_wrt_panel_com = ( + math_utils.subtract_frame_transforms( + t01=panel_com_pos_w, + q01=panel_com_quat_w, + t02=panel_target_pos_w, + q02=panel_target_quat_w, + ) + ) + panel_target_rotmat_wrt_panel_com = math_utils.matrix_from_quat( + panel_target_quat_wrt_panel_com + ) + + # Transformation | Target origin -> Hole entrance + hole_entrance_pos_w, _ = math_utils.combine_frame_transforms( + t01=target_pos_w, + q01=target_quat_w, + t12=hole_offset_pos_entrance, + ) + + # Transformation | Target origin -> Hole bottom + hole_bottom_pos_w, _ = math_utils.combine_frame_transforms( + t01=target_pos_w, + q01=target_quat_w, + t12=hole_offset_pos_bottom, + ) + + # Transformation | End-effector -> Object CoM + obj_com_pos_wrt_robot_ee, obj_com_quat_wrt_robot_ee = ( + math_utils.subtract_frame_transforms( + t01=robot_ee_pos_w.unsqueeze(1).repeat(1, 4, 1), + q01=robot_ee_quat_w.unsqueeze(1).repeat(1, 4, 1), + t02=obj_com_pos_w, + q02=obj_com_quat_w, + ) + ) + obj_com_rotmat_wrt_robot_ee = math_utils.matrix_from_quat(obj_com_quat_wrt_robot_ee) + + # Transformation | Peg ends -> Hole entrance + _hole_entrance_pos_wrt_peg_end0, hole_quat_wrt_peg = ( + math_utils.subtract_frame_transforms( + t01=peg_ends_pos_w[:, 0], + q01=obj_quat_w, + t02=hole_entrance_pos_w, + q02=target_quat_w, + ) + ) + _hole_entrance_pos_wrt_peg_end1, _ = math_utils.subtract_frame_transforms( + t01=peg_ends_pos_w[:, 1], + q01=obj_quat_w, + t02=hole_entrance_pos_w, + ) + hole_entrance_pos_wrt_peg_ends = torch.stack( + [_hole_entrance_pos_wrt_peg_end0, _hole_entrance_pos_wrt_peg_end1], dim=1 + ) + hole_rotmat_wrt_peg = math_utils.matrix_from_quat(hole_quat_wrt_peg) + + # Transformation | Peg ends -> Hole bottom + _hole_bottom_pos_wrt_peg_end0, _ = math_utils.subtract_frame_transforms( + t01=peg_ends_pos_w[:, 0], + q01=obj_quat_w, + t02=hole_bottom_pos_w, + ) + _hole_bottom_pos_wrt_peg_end1, _ = math_utils.subtract_frame_transforms( + t01=peg_ends_pos_w[:, 1], + q01=obj_quat_w, + t02=hole_bottom_pos_w, + ) + hole_bottom_pos_wrt_peg_ends = torch.stack( + [_hole_bottom_pos_wrt_peg_end0, _hole_bottom_pos_wrt_peg_end1], dim=1 + ) + + ## Rewards + # Penalty: Action rate + WEIGHT_ACTION_RATE = -0.05 + penalty_action_rate = WEIGHT_ACTION_RATE * torch.sum( + torch.square(current_action - previous_action), dim=1 + ) + + # Penalty: Undesired robot arm contacts + WEIGHT_UNDERSIRED_ROBOT_ARM_CONTACTS = -0.1 + THRESHOLD_UNDERSIRED_ROBOT_ARM_CONTACTS = 10.0 + penalty_undersired_robot_arm_contacts = WEIGHT_UNDERSIRED_ROBOT_ARM_CONTACTS * ( + torch.max(torch.norm(robot_arm_contact_net_forces, dim=-1), dim=1)[0] + > THRESHOLD_UNDERSIRED_ROBOT_ARM_CONTACTS + ) + + # Reward: Distance | End-effector <--> Object + WEIGHT_DISTANCE_EE_TO_OBJ = 1.0 + TANH_STD_DISTANCE_EE_TO_OBJ = 0.25 + reward_distance_ee_to_obj = WEIGHT_DISTANCE_EE_TO_OBJ * ( + 1.0 + - torch.tanh( + torch.norm(obj_com_pos_wrt_robot_ee, dim=-1) / TANH_STD_DISTANCE_EE_TO_OBJ + ) + ).sum(dim=-1) + + # Reward: Distance | End-effector <--> Panel + WEIGHT_DISTANCE_EE_TO_PANEL = 4.0 * WEIGHT_DISTANCE_EE_TO_OBJ + TANH_STD_DISTANCE_EE_TO_PANEL = 0.25 + reward_distance_ee_to_panel = WEIGHT_DISTANCE_EE_TO_PANEL * ( + 1.0 + - torch.tanh( + torch.norm(panel_com_pos_wrt_robot_ee, dim=-1) + / TANH_STD_DISTANCE_EE_TO_PANEL + ) + ) + + # Reward: Object grasped + WEIGHT_OBJ_GRASPED = 4.0 + THRESHOLD_OBJ_GRASPED = 5.0 + reward_obj_grasped = WEIGHT_OBJ_GRASPED * ( + torch.mean( + torch.max(torch.norm(robot_hand_obj_contact_force_matrix, dim=-1), dim=-1)[ + 0 + ], + dim=1, + ) + > THRESHOLD_OBJ_GRASPED + ) + + # Reward: Object lifted + WEIGHT_OBJ_LIFTED = 8.0 + HEIGHT_OFFSET_OBJ_LIFTED = 0.3 + HEIGHT_SPAN_OBJ_LIFTED = 0.25 + TAHN_STD_HEIGHT_OBJ_LIFTED = 0.05 + obj_target_height_offset = ( + torch.abs( + obj_com_pos_w[:, :, 2] - initial_obj_height_w - HEIGHT_OFFSET_OBJ_LIFTED + ) + - HEIGHT_SPAN_OBJ_LIFTED + ).clamp(min=0.0) + reward_obj_lifted = WEIGHT_OBJ_LIFTED * ( + 1.0 - torch.tanh(obj_target_height_offset / TAHN_STD_HEIGHT_OBJ_LIFTED) + ).sum(dim=-1) + + # Reward: Panel lifted + WEIGHT_PANEL_LIFTED = 2 * WEIGHT_OBJ_LIFTED + HEIGHT_OFFSET_PANEL_LIFTED = 0.5 + HEIGHT_SPAN_PANEL_LIFTED = 0.25 + TAHN_STD_HEIGHT_PANEL_LIFTED = 0.1 + panel_target_height_offset = ( + torch.abs( + panel_com_pos_w[:, 2] - initial_panel_height_w - HEIGHT_OFFSET_PANEL_LIFTED + ) + - HEIGHT_SPAN_PANEL_LIFTED + ).clamp(min=0.0) + reward_panel_lifted = WEIGHT_PANEL_LIFTED * ( + 1.0 - torch.tanh(panel_target_height_offset / TAHN_STD_HEIGHT_PANEL_LIFTED) + ) + + # Reward: Alignment | Peg -> Hole | Primary Z axis + WEIGHT_ALIGN_PEG_TO_HOLE_PRIMARY = 8.0 + TANH_STD_ALIGN_PEG_TO_HOLE_PRIMARY = 0.5 + _peg_to_hole_primary_axis_similarity = torch.abs(hole_rotmat_wrt_peg[:, :, 2, 2]) + reward_align_peg_to_hole_primary = WEIGHT_ALIGN_PEG_TO_HOLE_PRIMARY * ( + 1.0 + - torch.tanh( + (1.0 - _peg_to_hole_primary_axis_similarity) + / TANH_STD_ALIGN_PEG_TO_HOLE_PRIMARY + ) + ).sum(dim=-1) + + # Reward: Alignment | Peg -> Hole | Secondary XY axes (affected by primary via power) + WEIGHT_ALIGN_PEG_TO_HOLE_SECONDARY = 4.0 + TANH_STD_ALIGN_PEG_TO_HOLE_SECONDARY = 0.2 + _peg_to_hole_yaw = torch.atan2( + hole_rotmat_wrt_peg[:, :, 0, 1], hole_rotmat_wrt_peg[:, :, 0, 0] + ) + _symmetry_step = 2 * torch.pi / peg_rot_symmetry_n + _peg_to_hole_yaw_symmetric_directional = _peg_to_hole_yaw % _symmetry_step + # Note: Lines above might result in NaN/inf when `peg_rot_symmetry_n=0` (infinite circular symmetry) + # However, the following `torch.where()` will handle this case + _peg_to_hole_yaw_symmetric_normalized = torch.where( + peg_rot_symmetry_n <= 0, + 0.0, + torch.min( + _peg_to_hole_yaw_symmetric_directional, + _symmetry_step - _peg_to_hole_yaw_symmetric_directional, + ) + / (_symmetry_step / 2.0), + ) + reward_align_peg_to_hole_secondary = WEIGHT_ALIGN_PEG_TO_HOLE_SECONDARY * ( + 1.0 + - torch.tanh( + _peg_to_hole_yaw_symmetric_normalized.pow( + _peg_to_hole_primary_axis_similarity + ) + / TANH_STD_ALIGN_PEG_TO_HOLE_SECONDARY + ) + ).sum(dim=-1) + + # Reward: Distance | Peg -> Hole entrance + WEIGHT_DISTANCE_PEG_TO_HOLE_ENTRANCE = 16.0 + TANH_STD_DISTANCE_PEG_TO_HOLE_ENTRANCE = 0.025 + reward_distance_peg_to_hole_entrance = WEIGHT_DISTANCE_PEG_TO_HOLE_ENTRANCE * ( + 1.0 + - torch.tanh( + torch.min(torch.norm(hole_entrance_pos_wrt_peg_ends, dim=-1), dim=1)[0] + / TANH_STD_DISTANCE_PEG_TO_HOLE_ENTRANCE + ) + ).sum(dim=-1) + + # Reward: Distance | Peg -> Hole bottom + WEIGHT_DISTANCE_PEG_TO_HOLE_BOTTOM = 128.0 + TANH_STD_DISTANCE_PEG_TO_HOLE_BOTTOM = 0.005 + reward_distance_peg_to_hole_bottom = WEIGHT_DISTANCE_PEG_TO_HOLE_BOTTOM * ( + 1.0 + - torch.tanh( + torch.min(torch.norm(hole_bottom_pos_wrt_peg_ends, dim=-1), dim=1)[0] + / TANH_STD_DISTANCE_PEG_TO_HOLE_BOTTOM + ) + ).sum(dim=-1) + + # Reward: Distance | Panel <--> Panel Target + WEIGHT_DISTANCE_PANEL_TO_TARGET = 4.0 * WEIGHT_DISTANCE_PEG_TO_HOLE_BOTTOM + TANH_STD_DISTANCE_PANEL_TO_TARGET = 0.05 + reward_distance_panel_to_target = WEIGHT_DISTANCE_PANEL_TO_TARGET * ( + 1.0 + - torch.tanh( + torch.norm(panel_target_pos_wrt_panel_com, dim=-1) + / TANH_STD_DISTANCE_PANEL_TO_TARGET + ) + ) + + # Total reward + rewards = torch.sum( + torch.stack( + [ + penalty_action_rate, + penalty_undersired_robot_arm_contacts, + reward_distance_ee_to_obj, + reward_distance_ee_to_panel, + reward_obj_grasped, + reward_obj_lifted, + reward_panel_lifted, + reward_align_peg_to_hole_primary, + reward_align_peg_to_hole_secondary, + reward_distance_peg_to_hole_entrance, + reward_distance_peg_to_hole_bottom, + reward_distance_panel_to_target, + ], + dim=-1, + ), + dim=-1, + ) + + ## Termination and truncation + truncations = episode_length_buf > (max_episode_length - 1) + terminations = torch.zeros_like(truncations) + + # print( + # f""" + # penalty | action_rate: {float(penalty_action_rate[0])} + # penalty | undersired_robot_arm_contacts: {float(penalty_undersired_robot_arm_contacts[0])} + # reward | distance_ee_to_obj: {float(reward_distance_ee_to_obj[0])} + # reward | distance_ee_to_panel: {float(reward_distance_ee_to_panel[0])} + # reward | obj_grasped: {float(reward_obj_grasped[0])} + # reward | obj_lifted: {float(reward_obj_lifted[0])} + # reward | panel_lifted: {float(reward_panel_lifted[0])} + # reward | align_peg_to_hole_primary: {float(reward_align_peg_to_hole_primary[0])} + # reward | align_peg_to_hole_secondary: {float(reward_align_peg_to_hole_secondary[0])} + # reward | distance_peg_to_hole_entrance: {float(reward_distance_peg_to_hole_entrance[0])} + # reward | distance_peg_to_hole_bottom: {float(reward_distance_peg_to_hole_bottom[0])} + # reward | distance_panel_to_target: {float(reward_distance_panel_to_target[0])} + # total: {float(rewards[0])} + # """ + # ) + + return ( + remaining_time, + robot_joint_pos_arm, + robot_joint_pos_hand, + robot_ee_rotmat_wrt_base, + obj_com_pos_wrt_robot_ee, + obj_com_rotmat_wrt_robot_ee, + hole_entrance_pos_wrt_peg_ends, + hole_bottom_pos_wrt_peg_ends, + hole_rotmat_wrt_peg, + panel_com_pos_wrt_robot_ee, + panel_com_rotmat_wrt_robot_ee, + panel_target_pos_wrt_panel_com, + panel_target_rotmat_wrt_panel_com, + rewards, + terminations, + truncations, + ) + + +@torch.jit.script +def _construct_observations( + *, + remaining_time: torch.Tensor, + robot_joint_pos_arm: torch.Tensor, + robot_joint_pos_hand: torch.Tensor, + robot_ee_pos_wrt_base: torch.Tensor, + robot_ee_rotmat_wrt_base: torch.Tensor, + robot_hand_wrench: torch.Tensor, + obj_com_pos_wrt_robot_ee: torch.Tensor, + obj_com_rotmat_wrt_robot_ee: torch.Tensor, + hole_entrance_pos_wrt_peg_ends: torch.Tensor, + hole_bottom_pos_wrt_peg_ends: torch.Tensor, + hole_rotmat_wrt_peg: torch.Tensor, + panel_com_pos_wrt_robot_ee: torch.Tensor, + panel_com_rotmat_wrt_robot_ee: torch.Tensor, + panel_target_pos_wrt_panel_com: torch.Tensor, + panel_target_rotmat_wrt_panel_com: torch.Tensor, +) -> Dict[str, torch.Tensor]: + """ + Note: The `robot_hand_wrench` is considered as state (robot without force-torque sensors) + """ + + num_envs = remaining_time.size(0) + + # Robot joint positions + robot_joint_pos_hand_mean = robot_joint_pos_hand.mean(dim=-1, keepdim=True) + + # End-effector pose (position and '6D' rotation) + robot_ee_rot6d = math_utils.rotmat_to_rot6d(robot_ee_rotmat_wrt_base) + + # Wrench + robot_hand_wrench_full = robot_hand_wrench.view(num_envs, -1) + robot_hand_wrench_mean = robot_hand_wrench.mean(dim=1) + + # Transformation | End-effector -> Object CoM + obj_com_rot6d_wrt_robot_ee = math_utils.rotmat_to_rot6d(obj_com_rotmat_wrt_robot_ee) + + # Transformation | Object -> Target + hole_rot6d_wrt_peg = math_utils.rotmat_to_rot6d(hole_rotmat_wrt_peg) + + # Transformation | End-effector -> Panel Target + panel_com_rot6d_wrt_robot_ee = math_utils.rotmat_to_rot6d( + panel_com_rotmat_wrt_robot_ee + ) + # Transformation | Panel CoM -> Panel Target + panel_target_rot6d_wrt_panel_com = math_utils.rotmat_to_rot6d( + panel_target_rotmat_wrt_panel_com + ) + + return { + "state": torch.cat( + [ + obj_com_pos_wrt_robot_ee.view(num_envs, -1), + obj_com_rot6d_wrt_robot_ee.view(num_envs, -1), + hole_entrance_pos_wrt_peg_ends.view(num_envs, -1), + hole_bottom_pos_wrt_peg_ends.view(num_envs, -1), + hole_rot6d_wrt_peg.view(num_envs, -1), + panel_com_pos_wrt_robot_ee, + panel_com_rot6d_wrt_robot_ee, + panel_target_pos_wrt_panel_com, + panel_target_rot6d_wrt_panel_com, + robot_hand_wrench_mean.view(num_envs, -1), + ], + dim=-1, + ), + "state_dyn": torch.cat([robot_hand_wrench_full], dim=-1), + "proprio": torch.cat( + [ + remaining_time, + robot_ee_pos_wrt_base, + robot_ee_rot6d, + robot_joint_pos_hand_mean, + ], + dim=-1, + ), + "proprio_dyn": torch.cat([robot_joint_pos_arm, robot_joint_pos_hand], dim=-1), + } diff --git a/space_robotics_bench/tasks/solar_panel_assembly/task_visual.py b/space_robotics_bench/tasks/solar_panel_assembly/task_visual.py new file mode 100644 index 0000000..9b765bc --- /dev/null +++ b/space_robotics_bench/tasks/solar_panel_assembly/task_visual.py @@ -0,0 +1,32 @@ +from typing import Dict + +import torch +from omni.isaac.lab.utils import configclass + +from space_robotics_bench.envs import ( + VisualManipulationEnvExt, + VisualManipulationEnvExtCfg, +) + +from .task import Task, TaskCfg + + +@configclass +class VisualTaskCfg(TaskCfg, VisualManipulationEnvExtCfg): + def __post_init__(self): + TaskCfg.__post_init__(self) + VisualManipulationEnvExtCfg.__post_init__(self) + + +class VisualTask(Task, VisualManipulationEnvExt): + cfg: VisualTaskCfg + + def __init__(self, cfg: VisualTaskCfg, **kwargs): + Task.__init__(self, cfg, **kwargs) + VisualManipulationEnvExt.__init__(self, cfg, **kwargs) + + def _get_observations(self) -> Dict[str, torch.Tensor]: + return { + **Task._get_observations(self), + **VisualManipulationEnvExt._get_observations(self), + } diff --git a/space_robotics_bench/utils/__init__.py b/space_robotics_bench/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/space_robotics_bench/utils/cfg.py b/space_robotics_bench/utils/cfg.py new file mode 100644 index 0000000..452f3de --- /dev/null +++ b/space_robotics_bench/utils/cfg.py @@ -0,0 +1,57 @@ +import os +from typing import Dict + +SUPPORTED_FRAMEWORKS = { + "dreamerv3": {"multi_algo": False}, + "sb3": {"multi_algo": True}, + "robomimic": {"multi_algo": True}, +} +SUPPORTED_CFG_FILE_EXTENSIONS = ( + "json", + "toml", + "yaml", + "yml", +) +FRAMEWORK_CFG_ENTRYPOINT_KEY = "{FRAMEWORK}_cfg" +FRAMEWORK_MULTI_ALGO_CFG_ENTRYPOINT_KEY = "{FRAMEWORK}_{ALGO}_cfg" + + +def parse_algo_configs(cfg_dir: str) -> Dict[str, str]: + algo_config = {} + + for root, _, files in os.walk(cfg_dir): + for file in files: + if not file.endswith(SUPPORTED_CFG_FILE_EXTENSIONS): + continue + file = os.path.join(root, file) + + key = _identify_config(root, file) + if key is not None: + algo_config[key] = file + + return algo_config + + +def _identify_config(root: str, file) -> str: + basename = os.path.basename(file).split(".")[0] + + for framework, properties in SUPPORTED_FRAMEWORKS.items(): + if root.endswith(framework): + assert properties["multi_algo"] + if "_" in basename: + algo = basename.split("_")[0] + else: + algo = basename + return FRAMEWORK_MULTI_ALGO_CFG_ENTRYPOINT_KEY.format( + FRAMEWORK=framework, ALGO=algo + ) + elif basename.startswith(f"{framework}"): + if properties["multi_algo"]: + algo = basename[len(framework) + 1 :].split("_")[0] + return FRAMEWORK_MULTI_ALGO_CFG_ENTRYPOINT_KEY.format( + FRAMEWORK=framework, ALGO=algo + ) + else: + return FRAMEWORK_CFG_ENTRYPOINT_KEY.format(FRAMEWORK=framework) + + return None diff --git a/space_robotics_bench/utils/color.py b/space_robotics_bench/utils/color.py new file mode 100644 index 0000000..78d76c9 --- /dev/null +++ b/space_robotics_bench/utils/color.py @@ -0,0 +1,28 @@ +from typing import Tuple + +import space_robotics_bench.core.envs as env_utils +import space_robotics_bench.core.sim as sim_utils + + +def contrastive_color_from_env_cfg( + env_cfg: env_utils.EnvironmentConfig, +) -> Tuple[float, float, float]: + match env_cfg.scenario: + case env_utils.Scenario.ASTEROID | env_utils.Scenario.MOON: + return (0.8, 0.8, 0.8) + case ( + env_utils.Scenario.EARTH + | env_utils.Scenario.MARS + | env_utils.Scenario.ORBIT + ): + return (0.1, 0.1, 0.1) + case _: + return (0.7071, 0.7071, 0.7071) + + +def preview_surface_from_env_cfg( + env_cfg: env_utils.EnvironmentConfig, +) -> sim_utils.PreviewSurfaceCfg: + return sim_utils.PreviewSurfaceCfg( + diffuse_color=contrastive_color_from_env_cfg(env_cfg), + ) diff --git a/space_robotics_bench/utils/image_proc.py b/space_robotics_bench/utils/image_proc.py new file mode 100644 index 0000000..95d0e75 --- /dev/null +++ b/space_robotics_bench/utils/image_proc.py @@ -0,0 +1,47 @@ +from typing import Dict, Tuple + +import torch +from omni.isaac.lab.sensors import Camera + + +def extract_images(camera: Camera) -> Dict[str, torch.Tensor]: + output = camera.data.output + return { + "rgb": output["rgb"], + "depth": output["distance_to_camera"], + } + + +@torch.jit.script +def process_rgb( + rgb: torch.Tensor, + dtype: torch.dtype, +) -> torch.Tensor: + return rgb[..., :3].to(dtype) / 255.0 + + +@torch.jit.script +def process_depth( + depth: torch.Tensor, + depth_range: Tuple[float, float], +) -> torch.Tensor: + return ( + depth.nan_to_num( + nan=depth_range[1], posinf=depth_range[1], neginf=depth_range[1] + ).clamp(depth_range[0], depth_range[1]) + - depth_range[0] + ) / (depth_range[1] - depth_range[0]) + + +@torch.jit.script +def construct_observation( + *, + rgb: torch.Tensor, + depth: torch.Tensor, + depth_range: Tuple[float, float], + image_name: str, +) -> Dict[str, torch.Tensor]: + return { + f"{image_name}_rgb": process_rgb(rgb, depth.dtype), + f"{image_name}_depth": process_depth(depth, depth_range), + } diff --git a/space_robotics_bench/utils/importer.py b/space_robotics_bench/utils/importer.py new file mode 100644 index 0000000..9a20280 --- /dev/null +++ b/space_robotics_bench/utils/importer.py @@ -0,0 +1,39 @@ +import importlib +import pkgutil +import sys +from typing import List, Optional + + +def import_modules_recursively(module_name: str, ignorelist: List[str] = []): + package = importlib.import_module(module_name) + for _ in _walk_modules( + path=package.__path__, prefix=f"{package.__name__}.", ignorelist=ignorelist + ): + pass + + +def _walk_modules( + path: Optional[str] = None, + prefix: str = "", + ignorelist: List[str] = [], +): + def seen(p, m={}): + if p in m: + return True + m[p] = True + + for info in pkgutil.iter_modules(path, prefix): + if any(module_name in info.name for module_name in ignorelist): + continue + + yield info + + if info.ispkg: + try: + __import__(info.name) + except Exception: + raise + else: + path = getattr(sys.modules[info.name], "__path__", None) or [] + path = [p for p in path if not seen(p)] + yield from _walk_modules(path, f"{info.name}.", ignorelist) diff --git a/space_robotics_bench/utils/math.py b/space_robotics_bench/utils/math.py new file mode 100644 index 0000000..c658e08 --- /dev/null +++ b/space_robotics_bench/utils/math.py @@ -0,0 +1,181 @@ +import math +from typing import Iterable, Tuple + +import torch +from omni.isaac.lab.utils.math import * # noqa: F403 +from omni.isaac.lab.utils.math import matrix_from_quat, quat_apply, quat_inv, quat_mul + + +def quat_from_rpy( + *args: Iterable[float], deg: bool = True +) -> Tuple[float, float, float, float]: + """ + Returns wxyz quaternion from roll-pitch-yaw angles. + """ + + rpy = args[0] if len(args) == 1 else args + + roll, pitch, yaw = (r * (math.pi / 180.0) for r in rpy) if deg else rpy + cy = math.cos(yaw / 2.0) + sy = math.sin(yaw / 2.0) + cr = math.cos(roll / 2.0) + sr = math.sin(roll / 2.0) + cp = math.cos(pitch / 2.0) + sp = math.sin(pitch / 2.0) + + qw = cy * cr * cp + sy * sr * sp + qx = cy * sr * cp - sy * cr * sp + qy = cy * cr * sp + sy * sr * cp + qz = sy * cr * cp - cy * sr * sp + + return (qw, qx, qy, qz) + + +@torch.jit.script +def quat_to_rot6d(quaternions: torch.Tensor) -> torch.Tensor: + return matrix_from_quat(quaternions)[..., :2].reshape(-1, 6) + + +@torch.jit.script +def rotmat_to_rot6d(rotmat: torch.Tensor) -> torch.Tensor: + return rotmat[..., :2].reshape(-1, 6) + + +@torch.jit.script +def combine_frame_transforms( + t01: torch.Tensor, + q01: torch.Tensor, + t12: torch.Tensor | None = None, + q12: torch.Tensor | None = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + r"""Combine transformations between two reference frames into a stationary frame. + + It performs the following transformation operation: :math:`T_{02} = T_{01} \times T_{12}`, + where :math:`T_{AB}` is the homogeneous transformation matrix from frame A to B. + + Args: + t01: Position of frame 1 w.r.t. frame 0. Shape is (N, 3). + q01: Quaternion orientation of frame 1 w.r.t. frame 0 in (w, x, y, z). Shape is (N, 4). + t12: Position of frame 2 w.r.t. frame 1. Shape is (N, 3). + Defaults to None, in which case the position is assumed to be zero. + q12: Quaternion orientation of frame 2 w.r.t. frame 1 in (w, x, y, z). Shape is (N, 4). + Defaults to None, in which case the orientation is assumed to be identity. + + Returns: + A tuple containing the position and orientation of frame 2 w.r.t. frame 0. + Shape of the tensors are (N, 3) and (N, 4) respectively. + """ + # compute orientation + q02 = quat_mul(q01, q12) if q12 is not None else q01 + # compute translation + t02 = t01 + quat_apply(q01, t12) if t12 is not None else t01 + return t02, q02 + + +@torch.jit.script +def subtract_frame_transforms( + t01: torch.Tensor, + q01: torch.Tensor, + t02: torch.Tensor | None = None, + q02: torch.Tensor | None = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + r"""Subtract transformations between two reference frames into a stationary frame. + + It performs the following transformation operation: :math:`T_{12} = T_{01}^{-1} \times T_{02}`, + where :math:`T_{AB}` is the homogeneous transformation matrix from frame A to B. + + Args: + t01: Position of frame 1 w.r.t. frame 0. Shape is (N, 3). + q01: Quaternion orientation of frame 1 w.r.t. frame 0 in (w, x, y, z). Shape is (N, 4). + t02: Position of frame 2 w.r.t. frame 0. Shape is (N, 3). + Defaults to None, in which case the position is assumed to be zero. + q02: Quaternion orientation of frame 2 w.r.t. frame 0 in (w, x, y, z). Shape is (N, 4). + Defaults to None, in which case the orientation is assumed to be identity. + + Returns: + A tuple containing the position and orientation of frame 2 w.r.t. frame 1. + Shape of the tensors are (N, 3) and (N, 4) respectively. + """ + # compute orientation + q10 = quat_inv(q01) + q12 = quat_mul(q10, q02) if q02 is not None else q10 + # compute translation + t12 = quat_apply(q10, t02 - t01) if t02 is not None else quat_apply(q10, -t01) + return t12, q12 + + +@torch.jit.script +def transform_points( + points: torch.Tensor, + pos: torch.Tensor | None = None, + quat: torch.Tensor | None = None, +) -> torch.Tensor: + r"""Transform input points in a given frame to a target frame. + + This function transform points from a source frame to a target frame. The transformation is defined by the + position :math:`t` and orientation :math:`R` of the target frame in the source frame. + + .. math:: + p_{target} = R_{target} \times p_{source} + t_{target} + + If the input `points` is a batch of points, the inputs `pos` and `quat` must be either a batch of + positions and quaternions or a single position and quaternion. If the inputs `pos` and `quat` are + a single position and quaternion, the same transformation is applied to all points in the batch. + + If either the inputs :attr:`pos` and :attr:`quat` are None, the corresponding transformation is not applied. + + Args: + points: Points to transform. Shape is (N, P, 3) or (P, 3). + pos: Position of the target frame. Shape is (N, 3) or (3,). + Defaults to None, in which case the position is assumed to be zero. + quat: Quaternion orientation of the target frame in (w, x, y, z). Shape is (N, 4) or (4,). + Defaults to None, in which case the orientation is assumed to be identity. + + Returns: + Transformed points in the target frame. Shape is (N, P, 3) or (P, 3). + + Raises: + ValueError: If the inputs `points` is not of shape (N, P, 3) or (P, 3). + ValueError: If the inputs `pos` is not of shape (N, 3) or (3,). + ValueError: If the inputs `quat` is not of shape (N, 4) or (4,). + """ + points_batch = points.clone() + # check if inputs are batched + is_batched = points_batch.dim() == 3 + # -- check inputs + if points_batch.dim() == 2: + points_batch = points_batch[None] # (P, 3) -> (1, P, 3) + if points_batch.dim() != 3: + raise ValueError( + f"Expected points to have dim = 2 or dim = 3: got shape {points.shape}" + ) + if pos is not None and pos.dim() != 1 and pos.dim() != 2: + raise ValueError( + f"Expected pos to have dim = 1 or dim = 2: got shape {pos.shape}" + ) + if quat is not None and quat.dim() != 1 and quat.dim() != 2: + raise ValueError( + f"Expected quat to have dim = 1 or dim = 2: got shape {quat.shape}" + ) + # -- rotation + if quat is not None: + # convert to batched rotation matrix + rot_mat = matrix_from_quat(quat) + if rot_mat.dim() == 2: + rot_mat = rot_mat[None] # (3, 3) -> (1, 3, 3) + # convert points to matching batch size (N, P, 3) -> (N, 3, P) + # and apply rotation + points_batch = torch.matmul(rot_mat, points_batch.transpose_(1, 2)) + # (N, 3, P) -> (N, P, 3) + points_batch = points_batch.transpose_(1, 2) + # -- translation + if pos is not None: + # convert to batched translation vector + pos = pos[None, None, :] if pos.dim() == 1 else pos[:, None, :] + # apply translation + points_batch += pos + # -- return points in same shape as input + if not is_batched: + points_batch = points_batch.squeeze(0) # (1, P, 3) -> (P, 3) + + return points_batch diff --git a/space_robotics_bench/utils/parsing.py b/space_robotics_bench/utils/parsing.py new file mode 100644 index 0000000..d4528c6 --- /dev/null +++ b/space_robotics_bench/utils/parsing.py @@ -0,0 +1,182 @@ +import datetime +import importlib +import inspect +import os +from typing import Any, Dict, Union + +import gymnasium as gym +import yaml +from omni.isaac.lab.utils import update_class_from_dict, update_dict + +from space_robotics_bench.core.envs import BaseEnvCfg + + +def load_cfg_from_registry( + task_name: str, entry_point_key: str, unpack_callable: bool = True +) -> Union[ + BaseEnvCfg, + Dict[str, Any], +]: + """Load default configuration given its entry point from the gym registry. + + This function loads the configuration object from the gym registry for the given task name. + It supports both YAML and Python configuration files. + + It expects the configuration to be registered in the gym registry as: + + .. code-block:: python + + gym.register( + id="My-Awesome-Task-v0", + ... + kwargs={"env_entry_point_cfg": "path.to.config:ConfigClass"}, + ) + + The parsed configuration object for above example can be obtained as: + + .. code-block:: python + + from space_robotics_bench.utils.parsing import load_cfg_from_registry, + + cfg = load_cfg_from_registry("My-Awesome-Task-v0", "env_entry_point_cfg") + + Args: + task_name: The name of the environment. + entry_point_key: The entry point key to resolve the configuration file. + + Returns: + The parsed configuration object. This is either a dictionary or a class object. + + Raises: + ValueError: If the entry point key is not available in the gym registry for the task. + """ + # Obtain the configuration entry point + cfg_entry_point = gym.spec(task_name).kwargs.get(entry_point_key) + # Check if entry point exists + if cfg_entry_point is None: + raise ValueError( + f"Could not find configuration for the environment: '{task_name}'." + f" Please check that the gym registry has the entry point: '{entry_point_key}'." + ) + # Parse the default config file + if isinstance(cfg_entry_point, str) and cfg_entry_point.endswith(".yaml"): + if os.path.exists(cfg_entry_point): + # Absolute path for the config file + config_file = cfg_entry_point + else: + # Resolve path to the module location + mod_name, file_name = cfg_entry_point.split(":") + mod_path = os.path.dirname(importlib.import_module(mod_name).__file__) + # Obtain the configuration file path + config_file = os.path.join(mod_path, file_name) + # Load the configuration + print(f"[INFO]: Parsing configuration from: {config_file}") + with open(config_file, encoding="utf-8") as f: + cfg = yaml.full_load(f) + else: + if unpack_callable and callable(cfg_entry_point): + # Resolve path to the module location + mod_path = inspect.getfile(cfg_entry_point) + # Load the configuration + cfg_cls = cfg_entry_point() + elif isinstance(cfg_entry_point, str): + # Resolve path to the module location + mod_name, attr_name = cfg_entry_point.split(":") + mod = importlib.import_module(mod_name) + cfg_cls = getattr(mod, attr_name) + else: + cfg_cls = cfg_entry_point + # Load the configuration + print(f"[INFO]: Parsing configuration from: {cfg_entry_point}") + cfg = cfg_cls() if unpack_callable and callable(cfg_cls) else cfg_cls + return cfg + + +def parse_task_cfg( + task_name: str, + device: str = "cuda:0", + num_envs: int | None = None, + use_fabric: bool | None = None, +) -> Union[ + BaseEnvCfg, + Dict[str, Any], +]: + """Parse configuration for an environment and override based on inputs. + + Args: + task_name: The name of the environment. + device: The device to run the simulation on. Defaults to "cuda:0". + num_envs: Number of environments to create. Defaults to None, in which case it is left unchanged. + use_fabric: Whether to enable/disable fabric interface. If false, all read/write operations go through USD. + This slows down the simulation but allows seeing the changes in the USD through the USD stage. + Defaults to None, in which case it is left unchanged. + + Returns: + The parsed configuration object. This is either a dictionary or a class object. + + Raises: + ValueError: If the task name is not provided, i.e. None. + """ + # Create a dictionary to update from + args_cfg = {"sim": {}, "scene": {}} + + # Simulation device + args_cfg["sim"]["device"] = device + + # Disable fabric to read/write through USD + if use_fabric is not None: + args_cfg["sim"]["use_fabric"] = use_fabric + + # Number of environments + if num_envs is not None: + args_cfg["scene"]["num_envs"] = num_envs + + # Load the default configuration + cfg = load_cfg_from_registry(task_name, "task_cfg", unpack_callable=False) + # Update the main configuration + if callable(cfg): + default_cfg = cfg() + cfg = cfg( + sim=default_cfg.sim.replace(**args_cfg["sim"]), + scene=default_cfg.scene.replace(**args_cfg["scene"]), + ) + elif isinstance(cfg, dict): + cfg = update_dict(cfg, args_cfg) + else: + update_class_from_dict(cfg, args_cfg) + + return cfg + + +def create_logdir_path( + algo_name: str, + task_name: str, + prefix: str = "logs/", + timestamp_format="%Y%m%d-%H%M%S", +) -> str: + timestamp = datetime.datetime.now().strftime(timestamp_format) + logdir = os.path.realpath(os.path.join(prefix, algo_name, task_name, timestamp)) + os.makedirs(logdir, exist_ok=True) + return logdir + + +def get_last_run_logdir_path( + algo_name: str, + task_name: str, + prefix: str = "logs/", +) -> str: + logdir_root = os.path.abspath(os.path.join(prefix, algo_name, task_name)) + logdirs = [ + os.path.join(logdir_root, d) + for d in os.listdir(logdir_root) + if os.path.isdir(os.path.join(logdir_root, d)) + ] + logdirs.sort(key=os.path.getmtime, reverse=True) + if len(logdirs) == 0: + raise FileNotFoundError(f"No logdirs found in: {logdir_root}") + last_logdir = None + for d in logdirs: + if not d.endswith("eval"): + last_logdir = d + break + return last_logdir diff --git a/space_robotics_bench/utils/path.py b/space_robotics_bench/utils/path.py new file mode 100644 index 0000000..81c0fe7 --- /dev/null +++ b/space_robotics_bench/utils/path.py @@ -0,0 +1,6 @@ +from os import listdir, path +from typing import List + + +def abs_listdir(dir: str) -> List[str]: + return [path.realpath(path.join(dir, file)) for file in listdir(dir)] diff --git a/space_robotics_bench/utils/registry.py b/space_robotics_bench/utils/registry.py new file mode 100644 index 0000000..b6e9742 --- /dev/null +++ b/space_robotics_bench/utils/registry.py @@ -0,0 +1,37 @@ +from typing import Any, Dict, List, Literal, Optional, Union + +import gymnasium + +from space_robotics_bench.paths import SRB_HYPERPARAMS_DIR +from space_robotics_bench.utils.cfg import parse_algo_configs + + +def register_tasks( + tasks: Dict[ + str, + Dict[ + Literal["entry_point", "task_cfg", "cfg_dir"], + Union[gymnasium.Env, Any, str], + ], + ], + *, + default_entry_point: Optional[gymnasium.Env] = None, + default_task_cfg: Optional[Any] = None, + default_cfg_dir: Optional[str] = SRB_HYPERPARAMS_DIR, + namespace: str = "srb", +): + for id, cfg in tasks.items(): + entry_point = cfg.get("entry_point", default_entry_point) + gymnasium.register( + id=f"{namespace}/{id}", + entry_point=f"{entry_point.__module__}:{entry_point.__name__}", + kwargs={ + "task_cfg": cfg.get("task_cfg", default_task_cfg), + **parse_algo_configs(cfg.get("cfg_dir", default_cfg_dir)), + }, + disable_env_checker=True, + ) + + +def get_srb_tasks() -> List[str]: + return [env_id for env_id in gymnasium.registry.keys() if "srb" in env_id] diff --git a/space_robotics_bench/utils/ros.py b/space_robotics_bench/utils/ros.py new file mode 100644 index 0000000..c60a937 --- /dev/null +++ b/space_robotics_bench/utils/ros.py @@ -0,0 +1,27 @@ +from importlib.util import find_spec +from os import environ + +import omni.ext + + +def enable_ros2_bridge(): + if find_spec("rclpy") is not None: + # No-op if rclpy can already be imported, which means that ROS 2 is already sourced + return + + ld_library_path = environ.get("LD_LIBRARY_PATH", None) + ld_library_path = f":{ld_library_path}" if ld_library_path else "" + environ["LD_LIBRARY_PATH"] = ( + f"{omni.kit.paths.get_omni_path()}/exts/omni.isaac.ros2_bridge/humble/lib{ld_library_path}" + ) + + # Get the extension manager and list of available extensions + extension_manager = omni.kit.app.get_app().get_extension_manager() + extensions = extension_manager.get_extensions() + + # Extract the ROS extension + ros_extension = [ext for ext in extensions if "ros2_bridge" in ext["id"]][0] + + # Load the ROS extension if it is not already loaded + if not extension_manager.is_extension_enabled(ros_extension["id"]): + extension_manager.set_extension_enabled_immediate(ros_extension["id"], True) diff --git a/space_robotics_bench/utils/sampling.py b/space_robotics_bench/utils/sampling.py new file mode 100644 index 0000000..f3fd009 --- /dev/null +++ b/space_robotics_bench/utils/sampling.py @@ -0,0 +1,71 @@ +from typing import Iterable, List, Optional, Tuple, Union + +import numpy as np +import torch +from pxr import Gf + +from space_robotics_bench._rs.utils.sampling import * # noqa: F403 + + +def compute_grid_spacing( + num_instances: int, + spacing: float, + global_pos_offset: Optional[Union[np.ndarray, torch.Tensor, Iterable]] = None, + global_rot_offset: Optional[Union[np.ndarray, torch.Tensor, Iterable]] = None, +) -> Tuple[Tuple[int, int], Tuple[List[float], List[float]]]: + if global_pos_offset is not None: + if isinstance(global_pos_offset, torch.Tensor): + global_pos_offset = global_pos_offset.detach().cpu().numpy() + elif not isinstance(global_pos_offset, np.ndarray): + global_pos_offset = np.asarray(global_pos_offset) + if global_rot_offset is not None: + if isinstance(global_rot_offset, torch.Tensor): + global_rot_offset = global_rot_offset.detach().cpu().numpy() + elif not isinstance(global_rot_offset, np.ndarray): + global_rot_offset = np.asarray(global_rot_offset) + + num_per_row = np.ceil(np.sqrt(num_instances)) + num_rows = np.ceil(num_instances / num_per_row) + num_cols = np.ceil(num_instances / num_rows) + + row_offset = 0.5 * spacing * (num_rows - 1) + col_offset = 0.5 * spacing * (num_cols - 1) + + positions = [] + orientations = [] + + for i in range(num_instances): + # Compute transform + row = i // num_cols + col = i % num_cols + x = row_offset - row * spacing + y = col * spacing - col_offset + + position = [x, y, 0] + orientation = Gf.Quatd.GetIdentity() + + if global_pos_offset is not None: + translation = global_pos_offset + position + else: + translation = position + + if global_rot_offset is not None: + orientation = ( + Gf.Quatd( + global_rot_offset[0].item(), + Gf.Vec3d(global_rot_offset[1:].tolist()), + ) + * orientation + ) + + orientation = [ + orientation.GetReal(), + orientation.GetImaginary()[0], + orientation.GetImaginary()[1], + orientation.GetImaginary()[2], + ] + + positions.append(translation) + orientations.append(orientation) + + return ((num_rows, num_cols), (positions, orientations)) diff --git a/space_robotics_bench/utils/sim_app.py b/space_robotics_bench/utils/sim_app.py new file mode 100644 index 0000000..f50148c --- /dev/null +++ b/space_robotics_bench/utils/sim_app.py @@ -0,0 +1,5 @@ +from importlib.util import find_spec + + +def is_sim_app_started() -> bool: + return find_spec("omni.isaac.version") is not None diff --git a/space_robotics_bench/utils/string.py b/space_robotics_bench/utils/string.py new file mode 100644 index 0000000..785b01b --- /dev/null +++ b/space_robotics_bench/utils/string.py @@ -0,0 +1,26 @@ +import re + +from omni.isaac.lab.utils.string import * # noqa: F403 + +REGEX_CANONICALIZE_STR_PATTERN: re.Pattern = re.compile("[\W_]+") + + +def canonicalize_str(input: str) -> str: + """ + Canonicalizes a string by converting it to lowercase and removing unwanted characters. + + This function processes the input string to ensure it is in a standardized format, making it suitable for consistent usage in applications. It utilizes a predefined regular expression pattern to eliminate any characters that do not meet the specified criteria. + + Args: + input (str): The string to be canonicalized. + + Returns: + str: The canonicalized version of the input string. + """ + return REGEX_CANONICALIZE_STR_PATTERN.sub("", input.lower()) + + +def sanitize_camera_name(name: str) -> str: + for s in ["cam_", "camera_", "sensor_"]: + name = name.replace(s, "") + return name diff --git a/space_robotics_bench/utils/traceback.py b/space_robotics_bench/utils/traceback.py new file mode 100644 index 0000000..70932f9 --- /dev/null +++ b/space_robotics_bench/utils/traceback.py @@ -0,0 +1,20 @@ +from importlib.util import find_spec +from os import environ + + +def try_enable_rich_traceback(): + if environ.get("SRB_WITH_TRACEBACK", "false").lower() not in ["true", "1"]: + return + + if find_spec("rich") is None: + return + + from rich import traceback # isort:skip + import numpy + + traceback.install( + width=120, + show_locals=environ.get("SRB_WITH_TRACEBACK_LOCALS", "false").lower() + in ["true", "1"], + suppress=(numpy,), + )