feat(functional_jax): support reset options for JAX envs

[aryanyk]

Description

This PR adds support for options in reset() for functional JAX environments.Currently, FunctionalJaxEnv and FunctionalJaxVectorEnv accept an options argument in reset() but ignore it. This makes reset-time configuration
inconsistent with the Gymnasium API.

This change wires reset options into FuncEnv.get_default_params, enabling
environment-specific parameter overrides on reset.

Changes

• Use FuncEnv.get_default_params(**options) during reset
• Preserve existing behavior when options is not provided
• Add tests covering option handling and backward compatibility

Fixes #1138

Type of change

Please delete options that are not relevant.

• New feature (non-breaking change which adds functionality)

Checklist:

• I have run the pre-commit checks with pre-commit run --all-files (see CONTRIBUTING.md instructions to set it up)
• I have commented my code, particularly in hard-to-understand areas
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[aryanyk]

Hi @pseudo-rnd-thoughts

Hi, the CI workflows are awaiting approval and the PR is ready from my side.
Could a maintainer please approve the workflows and review when convenient?
Thanks!

[Kallinteris-Andreas]

1. Separation of concerns regarding Parameter Generation
   I have concerns regarding the modification of get_default_params(). Semantically, this method implies the retrieval of static, immutable default hyperparameters. Overloading it to handle dynamic state or initialization logic violates the Principle of Least Surprise. I strongly suggest decoupling the retrieval of defaults from the generation of instance-specific parameters. A factory pattern, such as a static or class method generate_params() or get_params(), would be more appropriate for this logic.

2. Violating Functional Purity in FuncEnv
   The proposed implementation for options introduces state mutation within FuncEnv. This stands in direct contradiction to the functional programming paradigm, where the environment instance is expected to remain stateless and methods should act as pure functions. Injecting state via options creates side effects.

Instead of a generic, state-mutating implementation in the base class, I recommend that environments requiring dynamic configuration implement a specialized generate_params(step) or pass configuration explicitly through the transition functions. This preserves the stateless nature of the FuncEnv abstraction.

3. Use Case Context
   Out of curiosity, which specific environment are you working with that necessitates options support in FuncEnv? The options argument isn't commonly used across Gymnasium environments in general, so I’m interested in understanding the specific use case driving this requirement.

[aryanyk]

Hi @Kallinteris-Andreas

Thanks for the detailed architectural feedback — I agree with the separation-of-concerns point.

I’ve refactored the implementation to introduce a dedicated generate_params(**kwargs) hook for reset-time parameter generation. By default this delegates to get_default_params(**kwargs) to preserve backward compatibility and semantic clarity.

reset() now calls generate_params() rather than get_default_params() directly.

Additionally, I identified a functional gap: although reset options were being accepted, the generated parameters were not consistently threaded through the JAX functional call chain. This has been corrected — params are now explicitly passed through initial, transition, observation, reward, terminal, info, and render.

No instance state is mutated; parameters remain explicit function inputs, preserving FuncEnv’s stateless design.

Added tests verify:

Reset options can change initialization behavior

generate_params is invoked during reset

Backward compatibility when no options are provided

Use case context: this enables reset-time configuration such as stochastic initial state ranges and domain randomization in functional JAX environments, aligning behavior with the broader Gymnasium API.

Happy to adjust further if needed.

P.S : Due to of push force it is comapring itself with previous PR . Please comapre this PR to latest main