anupamagarwal001/amc_allocator_env

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the BSD-style license found in the5# LICENSE file in the root directory of this source tree.6 7"""Client for the AMC allocator environment."""8 9from __future__ import annotations10 11from typing import Dict12 13from openenv.core import EnvClient14from openenv.core.client_types import StepResult15 16from .models import AllocatorObservation, AllocatorState, PortfolioAction17 18 19class AmcAllocatorEnv(EnvClient[PortfolioAction, AllocatorObservation, AllocatorState]):20    """Persistent client for allocator episodes over WebSocket."""21 22    def _step_payload(self, action: PortfolioAction) -> Dict:23        return {24            "target_weights": action.target_weights,25            "reason": action.reason,26        }27 28    def _parse_result(self, payload: Dict) -> StepResult[AllocatorObservation]:29        obs_data = payload.get("observation", {})30        observation = AllocatorObservation(31            task_id=obs_data.get("task_id", ""),32            task_description=obs_data.get("task_description", ""),33            step_index=obs_data.get("step_index", 0),34            steps_remaining=obs_data.get("steps_remaining", 0),35            prices=obs_data.get("prices", {}),36            signals=obs_data.get("signals", {}),37            current_weights=obs_data.get("current_weights", {}),38            cash_weight=obs_data.get("cash_weight", 1.0),39            portfolio_value=obs_data.get("portfolio_value", 1.0),40            turnover=obs_data.get("turnover", 0.0),41            risk_metrics=obs_data.get("risk_metrics", {}),42            done=payload.get("done", False),43            reward=payload.get("reward"),44            metadata=obs_data.get("metadata", {}),45        )46        return StepResult(47            observation=observation,48            reward=payload.get("reward"),49            done=payload.get("done", False),50        )51 52    def _parse_state(self, payload: Dict) -> AllocatorState:53        return AllocatorState.model_validate(payload)54