KR-IQN Risk-Sensitive MORL

This repository contains JAX/Flax implementations of risk-sensitive multi-objective Soft Actor-Critic variants. The main model is KR-IQN, a multi-objective SAC critic that applies risk distortion over quantile samples and trains with sampled preference weights.

The code has three practical entry points:

• main.py: MO-Gymnasium continuous-control environments.
• test_constrained_RL.py: Safety-Gymnasium environments with reward/cost objectives.
• test_finance.py: portfolio-management environment built on FinRL's StockTradingEnv.

Installation

The Dockerfile installs the base dependencies used by the MO-Gymnasium experiments:

docker build -t kr-iqn .
docker run --gpus all -it -v "$(pwd)":/workspace kr-iqn

The finance environment imports FinRL and related packages. Install them inside the environment when running test_finance.py:

bash install_finrl.sh

Repository Layout

• risk_morl/: algorithms, replay buffer, policies, critic/actor networks, risk measures, and evaluation utilities.
• basics/: IQN layers, weight sampling, and scheduling utilities.
• portfolio/port_mo_env.py: multi-objective portfolio environment.
• mo_hopper_swap_env/: reward-order shuffle wrapper for MO-Hopper ablations.
• plottings/: plotting helpers.
• train_df.csv, valid_df.csv, test_df.csv: bundled finance data files.

MO-Gymnasium Experiments

main.py is a Fire CLI. Constructor arguments such as --env_name, --batch_size, --gamma, --num_env, --folder, and --seed can be passed after the method name.

Example KR-IQN run:

python3 main.py kriqn \
  --env_name='mo-hopper-v5' \
  --learning_steps=1000000 \
  --batch_size=256 \
  --num_env=1 \
  --risk='neutral' \
  --truncation_lower=0 \
  --truncation_upper=2 \
  --gamma=0.99 \
  --seed=0

The included run.sh repeats this command for seeds 0..4.

Available methods in main.py:

• kriqn: KR-IQN MOSAC.
• marginal: Marginal-IQN baseline.
• ewp: EWP baseline.
• zhang: Zhang-style MMD baseline.
• kriqn_n_taus, kriqn_n_critics: ablations for quantile count and critic count.
• kriqn_ablation, kriqn_pe, kriqn_no_tqc, kriqn_no_proj_tqc, kriqn_no_marginal_tqc: KR-IQN ablations.
• kriqn_load_test, kriqn_no_tqc_train_only: load/test or train-only helpers.

Results are written under:

results/<env_name>/

Each run saves a .pth model and a result CSV.

Reward-Shuffled MO-Hopper

run_shuffle.py trains on MO-Hopper with permuted reward dimensions. The required constructor argument is env_index, which selects one of six reward permutations.

python3 run_shuffle.py kriqn \
  --env_index=0 \
  --learning_steps=1000000 \
  --risk='cvar' \
  --risk_param=0.5 \
  --index='[1]' \
  --seed=0

run_shuffle.py remaps the requested risk index through RewardShuffleWrapper.permute_risk_index() before constructing the risk measure.

Safety-Gymnasium Experiments

test_constrained_RL.py wraps a Safety-Gymnasium environment as a two-objective Gymnasium environment:

reward = [-cost * scale_cost, reward * scale_reward]

Example:

python3 test_constrained_RL.py armdsac \
  --env_name='SafetyPointGoal1-v0' \
  --learning_steps=500000 \
  --risk='cvar' \
  --risk_param=0.5 \
  --index='[-1]' \
  --seed=0

The available methods are armdsac and armdsac_train_only.

Finance Experiments

test_finance.py trains and evaluates on StockTradingMOEnv.

The script currently loads:

• train_df.csv for randomized training environments and vectorized validation rollouts.
• valid_df.csv for the deterministic test environment.

Example KR-IQN run:

python3 test_finance.py kr_iqn \
  --learning_steps=75000 \
  --batch_size=256 \
  --num_env=1 \
  --risk='cvar' \
  --index='[-1]' \
  --risk_param=0.5 \
  --truncation_lower=1 \
  --truncation_upper=2 \
  --gamma=0.99 \
  --seed=0 \
  --test_interval=10000 \
  --actor_lr=3e-4 \
  --ent_coef=5e-2

The included run_finance.sh runs several marginal configurations over seeds 0..4, including simplex, wang, cvar, and neutral risk settings.

Available methods in test_finance.py:

• kr_iqn: KR-IQN MOSAC.
• marginal: Marginal-IQN baseline.
• ewp: EWP baseline.
• zhang: Zhang-style MMD baseline.
• load_test_kriqn, load_test_ewp: load a saved model and run evaluation helpers.

Finance outputs are written under:

results/finance/

Typical files are:

• <model>.pth
• <model>result_valid.csv
• <model>result_test.csv
• <model>test_info.csv

Portfolio Environment

StockTradingMOEnv groups portfolio value into four asset classes:

• Commodities: XLE, GLD, SLV
• Bonds: TLT, TIP, JNK
• Equities: SPY, QQQ, SOXX
• Cash

The environment exposes reward_dim = 3:

1. Downside portfolio-value change: clip(after_value - prev_value, -inf, 0) * reward_scaling
2. Asset-class diversification: (entropy(asset_class_weights) - 0.5) * 0.1
3. Log portfolio return: log(after_value / prev_value) * portfolio_value_coef

Observations include scaled cash, current open prices, current share holdings, average buy prices relative to open prices, and lagged technical indicators.

Training uses randomized start days, adversarial-uniform bidding, and low-price stop-loss calculation. Testing uses a fixed day order, uniform bidding, and close-price stop-loss calculation.

Risk Measures

Risk distortion is implemented in risk_morl/utils/risk_measures.py.

Commonly used values:

Name	Behavior
neutral	Identity transform.
cvar	Scales selected quantile dimensions by alpha.
wang	Applies a Wang transform with normal CDF/PPF shift alpha.
triangle	Simplex-style distortion over two selected reward indices.
simplex	Simplex-style distortion over three selected reward indices.

Some experimental transforms are also present in the generator, including power, polar, polar_power, and polar_wang.

The index argument selects which reward dimensions the risk measure acts on. Negative indices follow normal Python indexing, so [-1] targets the last reward dimension.

Notes

• CLI commands are powered by Python Fire, so list-like arguments should be passed as strings, for example --index='[-1]' or --index='[0,1,2]'.
• truncation_lower and truncation_upper are forwarded into the policy through policy_kwargs.
• CSV logging is enabled for KR-IQN in main.py and test_finance.py when use_csv=True; logs are written below results/log_dir/<env_name>.