Jump to content

Connect SuperML | Leeroopedia MCP: Equip your AI agents with best practices, code verification, and debugging knowledge. Powered by Leeroo — building Organizational Superintelligence. Contact us at founders@leeroo.com.

Implementation:Isaac sim IsaacGymEnvs ADR Get Set Env State

From Leeroopedia
Metadata
Knowledge Sources
Domains
Last Updated 2026-02-15 00:00 GMT

Overview

The get_env_state() and set_env_state() methods implement ADR training state serialization across the class hierarchy: VecTaskDextreme saves ADR parameter ranges, AllegroHandDextreme adds task-specific state (action moving average, save tensors). Together they enable seamless training resumption with identical ADR state.

Description

The state serialization is implemented as a chain of get_env_state() / set_env_state() calls across the class hierarchy:

  1. VecTaskDextreme.get_env_state() returns {'adr_params': self.adr_params} when ADR is enabled, or an empty dict otherwise.
  2. VecTaskDextreme.set_env_state() restores ADR parameters from the checkpoint, with an option to skip via adr_load_from_checkpoint.
  3. AllegroHandDextreme.get_env_state() extends the parent state with act_moving_average and calls super().get_env_state() to include ADR state.
  4. AllegroHandDextreme.get_save_tensors() captures per-environment runtime tensors (actions, cube state, goal state, joint positions/velocities, root state) for debugging or replay.

Usage

These methods are called by the rl_games training framework during checkpoint save/load:

# During checkpoint save (called by rl_games agent)
env_state = env.get_env_state()
# env_state is stored alongside model weights in the checkpoint file

# During checkpoint load (called by rl_games agent)
env.set_env_state(checkpoint['env_state'])

Code Reference

Source Location

  • File: isaacgymenvs/tasks/dextreme/adr_vec_task.py (lines 368--398) -- VecTaskDextreme.get_env_state() and VecTaskDextreme.set_env_state()
  • File: isaacgymenvs/tasks/dextreme/allegro_hand_dextreme.py (lines 134--163) -- AllegroHandDextreme.get_env_state() and AllegroHandDextreme.get_save_tensors()

Signatures

# VecTaskDextreme (adr_vec_task.py)
class VecTaskDextreme(EnvDextreme, VecTask):

    def get_env_state(self):
        """Return serializable environment state for checkpoint.

        Returns:
            dict with 'adr_params' key when ADR is enabled,
            empty dict otherwise.
        """

    def set_env_state(self, env_state):
        """Restore environment state from checkpoint.

        Respects adr_load_from_checkpoint flag -- if False,
        skips loading ADR params (uses fresh init_range values).

        Args:
            env_state: dict previously returned by get_env_state().
        """


# AllegroHandDextreme (allegro_hand_dextreme.py)
class AllegroHandDextreme(ADRVecTask):

    def get_env_state(self):
        """Return task-level environment state for checkpoint.

        Returns:
            dict with 'act_moving_average' plus all ADR state
            from super().get_env_state().
        """

    def get_save_tensors(self):
        """Capture per-environment runtime tensors for debugging/replay.

        Returns:
            dict with 'actions', 'cube_state', 'goal_state',
            'joint_positions', 'joint_velocities', 'root_state'.
        """

Import

from isaacgymenvs.tasks.dextreme.adr_vec_task import VecTaskDextreme, ADRVecTask
from isaacgymenvs.tasks.dextreme.allegro_hand_dextreme import AllegroHandDextreme

I/O Contract

Inputs

Input Contract -- set_env_state
Name Type Description
env_state dict or None State dictionary previously returned by get_env_state(). Contains adr_params and act_moving_average keys.

Outputs

Output Contract -- get_env_state
Name Type Description
adr_params dict Complete ADR parameter dictionary including current ranges, limits, deltas, and delta styles for all ADR parameters.
act_moving_average float Current action moving average coefficient used for action smoothing.
Output Contract -- get_save_tensors
Name Type Description
actions torch.Tensor (num_envs, 16) Current policy actions.
cube_state torch.Tensor (num_envs, 13) Object root state (pos, quat, lin_vel, ang_vel).
goal_state torch.Tensor (num_envs, 7+) Goal poses.
joint_positions torch.Tensor (num_envs, 16) DOF positions.
joint_velocities torch.Tensor (num_envs, 16) DOF velocities.
root_state torch.Tensor (num_envs, 13) Hand root state.

Implementation Details

VecTaskDextreme.get_env_state

def get_env_state(self):
    if self.use_adr:
        return dict(adr_params=self.adr_params)
    else:
        return {}

VecTaskDextreme.set_env_state

def set_env_state(self, env_state):
    if env_state is None:
        return

    for key in self.get_env_state().keys():
        if key == "adr_params" and self.use_adr and not self.adr_load_from_checkpoint:
            print("Skipping loading ADR params from checkpoint...")
            continue
        value = env_state.get(key, None)
        if value is None:
            continue
        self.__dict__[key] = value
        print(f'Loaded env state value {key}:{value}')

    if self.use_adr:
        print(f'ADR Params after loading from checkpoint: {self.adr_params}')

AllegroHandDextreme.get_env_state

def get_env_state(self):
    env_dict = dict(act_moving_average=self.act_moving_average)

    if self.use_adr:
        env_dict = dict(**env_dict, **super().get_env_state())

    return env_dict

AllegroHandDextreme.get_save_tensors

def get_save_tensors(self):
    if hasattr(self, 'actions'):
        actions = self.actions
    else:
        actions = torch.zeros(
            (self.num_envs, self.cfg["env"]["numActions"])
        ).to(self.device)

    return {
        'actions': actions,
        'cube_state': self.root_state_tensor[self.object_indices],
        'goal_state': self.goal_states,
        'joint_positions': self.dof_pos,
        'joint_velocities': self.dof_vel,
        'root_state': self.root_state_tensor[self.hand_indices],
    }

State Dictionary Structure

The complete state dictionary when ADR is enabled has the following structure:

{
    'act_moving_average': float,
    'adr_params': {
        'hand_damping': {
            'range': [float, float],     # Current ADR range
            'init_range': [float, float],
            'limits': [float, float],
            'delta': float,
            'delta_style': str,
            'range_path': str,
            'next_limits': [float, float],
        },
        'hand_stiffness': { ... },
        'object_mass': { ... },
        'affine_action_additive': { ... },
        # ... all ADR parameters
    }
}

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.
Principle
Implementation
Heuristic
Environment