You are correct: in the New API Stack, the policy_mapping_fn should use the signature (agent_id, episode, **kwargs) and not include the worker argument. Including worker leads to a TypeError, as you observed. The documentation currently does not clearly distinguish this, which can cause confusion for users migrating to or starting with the New API Stack. This has been confirmed by user reports and code examples in the Ray community and source code comments, which show the new stack omits worker, while the old stack includes it for backward compatibility and RolloutWorker-based workflows. See the discussion and examples in Ray GitHub issues and Ray Discourse.
Would you like more detail on how to structure your config for multi-agent environments in the New API Stack, or on how to avoid compatibility issues?
Sources:
- https://github.com/ray-project/ray/issues/44475
- https://discuss.ray.io/t/ray-2-4-0-rllib-completely-lost-with-documentation/13200
Hint: Mention @RunLLM in the post for followups.