Model design: IsaacLab application

[JulienHansen]

Hi,

I was going through the documentation for model design and was wondering what I was doing wrong with the following model:

seed: 42
# Models are instantiated using skrl's model instantiator utility
# https://skrl.readthedocs.io/en/latest/api/utils/model_instantiators.html
models:
  separate: False
  policy:  # see gaussian_model parameters
    class: GaussianMixin
    clip_actions: False
    clip_log_std: True
    min_log_std: -20.0
    max_log_std: 2.0
    initial_log_std: 0.5
    network:
    - name: features_extractor
      input: permute(STATES["camera"], (0, 3, 1, 2))  # PyTorch NHWC -> NCHW
      layers:
        - conv2d: {out_channels: 32, kernel_size: 8, stride: 4, padding: 0}
        - conv2d: {out_channels: 64, kernel_size: 4, stride: 2, padding: 0}
        - conv2d: {out_channels: 64, kernel_size: 3, stride: 1, padding: 0}
        - flatten
      activations: relu
    - name: net
      input: concatenate([features_extractor, STATES["robot-state"]])
      layers: [512, 256, 128]
      activations: elu
    output: ACTIONS
  value:  # see deterministic_model parameters
    class: DeterministicMixin
    clip_actions: False
    network:
      - name: net
        input: STATES
        layers: [64, 64]
        activations: elu
    output: ONE

  


// Rollout memory

- [ ]  https://skrl.readthedocs.io/en/latest/api/memories/random.html

memory:
  class: RandomMemory
  memory_size: -1  # automatically determined (same as agent:rollouts)


# PPO agent configuration (field names are from PPO_DEFAULT_CONFIG)
# https://skrl.readthedocs.io/en/latest/api/agents/ppo.html
agent:
  class: PPO
  rollouts: 24
  learning_epochs: 5
  mini_batches: 4
  discount_factor: 0.99
  lambda: 0.95
  learning_rate: 5.0e-04
  learning_rate_scheduler: KLAdaptiveLR
  learning_rate_scheduler_kwargs:
    kl_threshold: 0.016
  state_preprocessor: RunningStandardScaler
  state_preprocessor_kwargs: null
  value_preprocessor: RunningStandardScaler
  value_preprocessor_kwargs: null
  random_timesteps: 0
  learning_starts: 0
  grad_norm_clip: 1.0
  ratio_clip: 0.2
  value_clip: 0.2
  clip_predicted_values: True
  entropy_loss_scale: 0.3
  value_loss_scale: 1.0
  kl_threshold: 0.0
  rewards_shaper_scale: 0.01
  time_limit_bootstrap: False
  # logging and checkpoint
  experiment:
    directory: "quadcopter_direct"
    experiment_name: "100k"
    write_interval: auto
    checkpoint_interval: auto


//Sequential trainer
//https://skrl.readthedocs.io/en/latest/api/trainers/sequential.html
trainer:
  class: SequentialTrainer
  timesteps: 100000
  environment_info: log

My other application works fine, but for this particular model, where a CNN captures the image returned by a camera fixed on a drone and then concatenates the result with the robot's intrinsic values before passing it to an MLP, I am encountering an error when using the Jax framework:

  Traceback (most recent call last):
  File "/home/uliege/IsaacLab/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/utils/hydra.py", line 99, in hydra_main
    func(env_cfg, agent_cfg, *args, **kwargs)
  File "/home/uliege/Desktop/Isaac_drone/Isaac_drone/scripts/skrl/train.py", line 175, in main
    runner = Runner(env, agent_cfg)
  File "/home/uliege/miniconda3/envs/isaaclab_ju/lib/python3.10/site-packages/skrl/utils/runner/jax/runner.py", line 53, in __init__
    self._models = self._generate_models(self._env, copy.deepcopy(self._cfg))
  File "/home/uliege/miniconda3/envs/isaaclab_ju/lib/python3.10/site-packages/skrl/utils/runner/jax/runner.py", line 270, in _generate_models
    model.init_state_dict(role)
  File "/home/uliege/miniconda3/envs/isaaclab_ju/lib/python3.10/site-packages/flax/linen/module.py", line 699, in wrapped_module_method
    return self._call_wrapped_method(fun, args, kwargs)
  File "/home/uliege/miniconda3/envs/isaaclab_ju/lib/python3.10/site-packages/flax/linen/module.py", line 1216, in _call_wrapped_method
    y = run_fun(self, *args, **kwargs)
  File "/home/uliege/miniconda3/envs/isaaclab_ju/lib/python3.10/site-packages/skrl/models/jax/base.py", line 142, in init_state_dict
    self.state_dict = StateDict.create(apply_fn=self.apply, params=self.init(key, inputs, role))
  File "<string>", line 31, in __call__
IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices

And my observation space is designed as follow :

    def _get_observations(self) -> dict:
        cam_images = self.scene["camera"].data.output["rgb"]
        desired_pos_b, _ = subtract_frame_transforms(
            self._robot.data.root_state_w[:, :3], self._robot.data.root_state_w[:, 3:7], self._desired_pos_w
        )

        state = torch.cat([
            self._robot.data.root_lin_vel_b,        
            self._robot.data.root_ang_vel_b,       
            self._robot.data.projected_gravity_b,   
            desired_pos_b,
        ], dim=-1)

        observation = {
            "policy": {
                "robot-state": state,
                "camera": cam_images,
            }
        }
        return observation

I'm surely missing something. I started using IsaacLab a month ago and I'm not familiar with it, nor with SKRL. I have just defined the observation space and action space as follows:

    action_space = 4
    total_robot_state_dim = 12  # Adjust based on robot state features
    observation_space = total_robot_state_dim + (scene.camera.width * scene.camera.height * num_channels)

[Toni-SM]

Hi @JulienHansen

Because your observation is returned as a dictionary (in _get_observations), you need to define it as a gymnasium.spaces.Dict (or Isaac Lab's supported basic Python data types) in environment config too (see DirectRLEnvCfg.observation_space)...

E.g.:

• Using basic Python data types (for simplicity)

observation_space = {
    "robot-state": 12,
    "camera": [scene.camera.width, scene.camera.height, num_channels],
}

• Using gymnasium spaces

observation_space = gym.spaces.Dict({
    "robot-state": gym.spaces.Box(float("-inf"), float("inf"), shape=(12,)), 
    "camera": gym.spaces.Box(float("-inf"), float("inf"), shape=(scene.camera.width, scene.camera.height, num_channels)),
})

---

Btw, what skrl version are you using? (./isaaclab.sh -p -m pip show skrl)
Also, just out of curiosity, I see you are using JAX than PyTorch. In this case, there is no necessary to permute camera observation: see comment in https://github.com/isaac-sim/IsaacLab/blob/4d9914702d84577b59dc468485c6e31f61dcbfde/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/agents/skrl_camera_ppo_cfg.yaml#L17

[Toni-SM]

Hi @JulienHansen

Probably there is necessary to use jax.numpy.ndarray.at instead of raw slices when spaces utils methods are called in a jax.jit scope. I am creating a test scenario to test it.

I will let you know when ready/fixed.

[Toni-SM]

Hi @JulienHansen

Changed mad in develop...toni/model_instantiators#diff-5bcf41072a9caa3c578ae0ad9b9aae08a85b2044b76044b2a32081722b732dde in the branch toni/model_instantiators (file skrl/models/jax/base.py) should fix the issue.

Please, let me know if it works in your side.
I am going to do some other minor updates before merging it to develop

[JulienHansen]

I have tried with the new version of the develop branch and my initial error seems to be fixed, now using a composite observation space:

    observation_space = {
    "robot-state": 12,
    "camera": [scene.camera.width, scene.camera.height, num_channels],
    }
    
        def _get_observations(self) -> dict:
        cam_images = self.scene["camera"].data.output["rgb"]
        desired_pos_b, _ = subtract_frame_transforms(
            self._robot.data.root_state_w[:, :3], self._robot.data.root_state_w[:, 3:7], self._desired_pos_w
        )

        state = torch.cat([
            self._robot.data.root_lin_vel_b,        
            self._robot.data.root_ang_vel_b,        
            self._robot.data.projected_gravity_b,   
            desired_pos_b,
        ], dim=-1)

        observation = {
            "policy": {
                "robot-state": state,
                "camera": cam_images,
            }
        }
        return observation

And following model :

models:
  separate: False
  policy:  # see gaussian_model parameters
    class: GaussianMixin
    clip_actions: False
    clip_log_std: True
    min_log_std: -20.0
    max_log_std: 2.0
    initial_log_std: 0.5
    network:
    - name: features_extractor
      input: STATES["camera"] # TODO: NEED TO PERMUTE IF USING TORCH
      layers:
        - conv2d: {out_channels: 32, kernel_size: 8, stride: 4, padding: 0}
        - conv2d: {out_channels: 64, kernel_size: 4, stride: 2, padding: 0}
        - conv2d: {out_channels: 64, kernel_size: 3, stride: 1, padding: 0}
        - flatten
      activations: relu
    - name: net
      input: concatenate([features_extractor, STATES["robot-state"]])
      layers: [512, 256, 128]
      activations: elu
    output: ACTIONS
  value:  # see deterministic_model parameters
    class: DeterministicMixin
    clip_actions: False
    network:
      - name: net
        input: STATES
        layers: [64, 64]
        activations: elu
    output: ONE

I get the following error, i don't have time right now to investigate it so maybe it just come from something i miss on my side, this week-end i will test to see if it is still the observation space that it causing this:

    Traceback (most recent call last):
  File "/home/uliege/IsaacLab/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/utils/hydra.py", line 99, in hydra_main
    func(env_cfg, agent_cfg, *args, **kwargs)
  File "/home/uliege/Desktop/Isaac_drone/Isaac_drone/scripts/skrl/train.py", line 175, in main
    runner = Runner(env, agent_cfg)
  File "/home/uliege/miniconda3/envs/isaaclab_ju/lib/python3.10/site-packages/skrl/utils/runner/jax/runner.py", line 53, in __init__
    self._models = self._generate_models(self._env, copy.deepcopy(self._cfg))
  File "/home/uliege/miniconda3/envs/isaaclab_ju/lib/python3.10/site-packages/skrl/utils/runner/jax/runner.py", line 270, in _generate_models
    model.init_state_dict(role)
  File "/home/uliege/miniconda3/envs/isaaclab_ju/lib/python3.10/site-packages/flax/linen/module.py", line 699, in wrapped_module_method
    return self._call_wrapped_method(fun, args, kwargs)
  File "/home/uliege/miniconda3/envs/isaaclab_ju/lib/python3.10/site-packages/flax/linen/module.py", line 1216, in _call_wrapped_method
    y = run_fun(self, *args, **kwargs)
  File "/home/uliege/miniconda3/envs/isaaclab_ju/lib/python3.10/site-packages/skrl/models/jax/base.py", line 149, in init_state_dict
    self.state_dict = StateDict.create(apply_fn=self.apply, params=self.init(key, inputs, role))
  File "<string>", line 18, in __call__
  File "/home/uliege/miniconda3/envs/isaaclab_ju/lib/python3.10/site-packages/flax/linen/combinators.py", line 106, in __call__
    outputs = self.layers[0](*args, **kwargs)
  File "/home/uliege/miniconda3/envs/isaaclab_ju/lib/python3.10/site-packages/flax/linen/linear.py", line 254, in __call__
    (jnp.shape(inputs)[-1], self.features),
IndexError: tuple index out of range

[Toni-SM]

Hi @JulienHansen

The problem is with the value network definition.
Note that according to your environment, the model STATE is a dictionary and therefore should be treated as such, similar to what you do in policy

[JulienHansen]

Sorry for the time I took to continue this discussion, but I had to work on another project. Now, I have a few questions.

• I have changed my value network so that it is identical to the policy network, and I no longer get any errors. However, during testing, I observed that the agent never seems to learn anything. In this case, the drone never exhibits the behavior I achieved when training it without the composite observations, which is the only thing that has changed. Could this come from skrl ? Is there an available example of composite space environment right now ?

• Integrating the CNN dropped the number of environments from around 2000-3000 to 750–1000. Is there something I can do about these?

Thanks in advance for your time.