Update

Skylark0924 · Dec 21, 2023 · 172a276 · 172a276
1 parent f6f4aa7
commit 172a276
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/rofunc/learning/RofuncRL/agents/mixline/amp_agent.py b/rofunc/learning/RofuncRL/agents/mixline/amp_agent.py
@@ -12,16 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
+from typing import Callable, Union, Tuple, Optional, List
+
 import gym
 import gymnasium
-import math
+import rofunc as rf
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from omegaconf import DictConfig
-from typing import Callable, Union, Tuple, Optional, List
-
-import rofunc as rf
 from rofunc.learning.RofuncRL.agents.base_agent import BaseAgent
 from rofunc.learning.RofuncRL.models.actor_models import ActorAMP
 from rofunc.learning.RofuncRL.models.critic_models import Critic

diff --git a/rofunc/learning/RofuncRL/trainers/ppo_trainer.py b/rofunc/learning/RofuncRL/trainers/ppo_trainer.py
@@ -20,7 +20,7 @@
 class PPOTrainer(BaseTrainer):
     def __init__(self, cfg, env, device, env_name):
         super().__init__(cfg, env, device, env_name)
-        self.memory = RandomMemory(memory_size=cfg.Trainer.rollouts, num_envs=self.env.num_envs, device=device)
+        self.memory = RandomMemory(memory_size=self.rollouts, num_envs=self.env.num_envs, device=device)
         self.agent = PPOAgent(cfg, self.env.observation_space, self.env.action_space, self.memory,
                               device, self.exp_dir, self.rofunc_logger)
         self.setup_wandb()