Fix axes flattening for single subplot case

[Sgt-Hashtag]

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes # (issue)

Type of change

• Bug fix (non-breaking change which fixes an issue)

Screenshots

Please attach before and after screenshots of the change if applicable.

| Before | After |

Checklist

• [x] I have read and understood the contribution guidelines
• [x] I have run the pre-commit checks with ./isaaclab.sh --format
• [ ] I have made corresponding changes to the documentation
• [x] My changes generate no new warnings
• [ ] I have added tests that prove my fix is effective or that my feature works
• [ ] I have updated the changelog and the corresponding version in the extension's config/extension.toml file
• [ ] I have added my name to the CONTRIBUTORS.md or my name already exists there